diff --git a/.github/workflows/integration-tests.yml b/.github/workflows/integration-tests.yml index 70995b5..536e59b 100644 --- a/.github/workflows/integration-tests.yml +++ b/.github/workflows/integration-tests.yml @@ -45,6 +45,58 @@ jobs: --cache-to type=gha,mode=max \ . + - name: Prepare SSH target on runner + run: | + set -euxo pipefail + APX_SSH_USER="apxci" + KEY_PATH="${RUNNER_TEMP}/apx_ci_key" + KNOWN_HOSTS_PATH="${RUNNER_TEMP}/apx_ci_known_hosts" + + ssh-keygen -t ed25519 -N "" -f "${KEY_PATH}" -C "apx-ci@github-actions" + chmod 600 "${KEY_PATH}" + + if ! id -u "${APX_SSH_USER}" >/dev/null 2>&1; then + sudo useradd -m -s /bin/bash "${APX_SSH_USER}" + fi + sudo mkdir -p "/home/${APX_SSH_USER}/.ssh" + sudo chmod 700 "/home/${APX_SSH_USER}/.ssh" + sudo touch "/home/${APX_SSH_USER}/.ssh/authorized_keys" + sudo chmod 600 "/home/${APX_SSH_USER}/.ssh/authorized_keys" + sudo sh -lc "cat '${KEY_PATH}.pub' >> '/home/${APX_SSH_USER}/.ssh/authorized_keys'" + sudo chown -R "${APX_SSH_USER}:${APX_SSH_USER}" "/home/${APX_SSH_USER}/.ssh" + echo "${APX_SSH_USER} ALL=(ALL) NOPASSWD:ALL" | sudo tee "/etc/sudoers.d/90-${APX_SSH_USER}-nopasswd" + sudo chmod 440 "/etc/sudoers.d/90-${APX_SSH_USER}-nopasswd" + sudo visudo -cf "/etc/sudoers.d/90-${APX_SSH_USER}-nopasswd" + + sudo apt-get update + sudo apt-get install -y openssh-server + sudo mkdir -p /run/sshd + sudo sed -i 's/^#\?PubkeyAuthentication .*/PubkeyAuthentication yes/' /etc/ssh/sshd_config + sudo sed -i 's/^#\?PasswordAuthentication .*/PasswordAuthentication no/' /etc/ssh/sshd_config + sudo sed -i 's/^#\?KbdInteractiveAuthentication .*/KbdInteractiveAuthentication no/' /etc/ssh/sshd_config + if grep -q '^AllowUsers ' /etc/ssh/sshd_config; then + sudo sed -i "s/^AllowUsers .*/AllowUsers ${APX_SSH_USER}/" /etc/ssh/sshd_config + else + echo "AllowUsers ${APX_SSH_USER}" | sudo tee -a /etc/ssh/sshd_config + fi + sudo systemctl restart ssh || sudo service ssh restart || sudo /usr/sbin/sshd + + ssh-keyscan -H 127.0.0.1 > "${KNOWN_HOSTS_PATH}" || true + ssh-keyscan -H 172.17.0.1 >> "${KNOWN_HOSTS_PATH}" || true + if [ ! -s "${KNOWN_HOSTS_PATH}" ]; then + touch "${KNOWN_HOSTS_PATH}" + fi + chmod 644 "${KNOWN_HOSTS_PATH}" + + ssh -o BatchMode=yes -o StrictHostKeyChecking=yes -o UserKnownHostsFile="${KNOWN_HOSTS_PATH}" -i "${KEY_PATH}" "${APX_SSH_USER}@127.0.0.1" "sudo -n true && echo ssh_ok" + + { + echo "APX_TEST_SSH_KEY_PATH=${KEY_PATH}" + echo "APX_TEST_KNOWN_HOSTS_PATH=${KNOWN_HOSTS_PATH}" + echo "APX_TEST_REMOTE_USER=${APX_SSH_USER}" + echo "APX_TEST_REMOTE_IP=172.17.0.1" + } >> "${GITHUB_ENV}" + - name: Run integration tests env: MCP_IMAGE: arm-mcp:latest diff --git a/mcp-local/Dockerfile b/mcp-local/Dockerfile index 37ef07f..2628a5c 100644 --- a/mcp-local/Dockerfile +++ b/mcp-local/Dockerfile @@ -38,9 +38,12 @@ RUN apt-get update && apt-get install -y --no-install-recommends \ build-essential libelf-dev \ ca-certificates file tar xz-utils jq libmagic1 && \ rm -rf /var/lib/apt/lists/* -RUN curl -sSL https://raw.githubusercontent.com/JoeStech/arm-linux-migration-tools/main/scripts/install.sh | bash +RUN bash -o pipefail -c 'set -eux; \ + curl -fsSL --retry 5 --retry-delay 2 --retry-all-errors \ + https://raw.githubusercontent.com/JoeStech/arm-linux-migration-tools/main/scripts/install.sh | bash' # Temp until migrate-ease is updated -RUN curl -sSL https://raw.githubusercontent.com/JoeStech/migrate-ease/main/js/advisor/main.py \ +RUN curl -fsSL --retry 5 --retry-delay 2 --retry-all-errors \ + https://raw.githubusercontent.com/JoeStech/migrate-ease/main/js/advisor/main.py \ -o /opt/arm-migration-tools/migrate-ease/js/advisor/main.py RUN rm -f /usr/local/bin/aperf \ @@ -114,7 +117,9 @@ ENV DEBIAN_FRONTEND=noninteractive \ APX_HOME=/opt/ArmPerformix-cli-current \ APX_BIN=/opt/ArmPerformix-cli-current/apx \ VIRTUAL_ENV=/app/.venv \ - PATH=/app/.venv/bin:$PATH + PATH=/app/.venv/bin:$PATH \ + KNOWN_HOSTS_PATH=/run/keys/known_hosts \ + SSH_KEY_PATH=/run/keys/ssh-key.pem RUN apt-get update && apt-get install -y --no-install-recommends \ python3 \ diff --git a/mcp-local/server.py b/mcp-local/server.py index b7fb2e0..78e1c99 100644 --- a/mcp-local/server.py +++ b/mcp-local/server.py @@ -292,8 +292,9 @@ def apx_recipe_run(cmd:str, remote_ip_addr:str, remote_usr:str, recipe:str="code }, ) apx_dir = os.environ.get("APX_HOME", "/opt/apx") - key_path = os.getenv("SSH_KEY_PATH") - known_hosts_path = os.getenv("KNOWN_HOSTS_PATH") + key_path = os.getenv("SSH_KEY_PATH", "/run/keys/ssh-key.pem") + known_hosts_path = os.getenv("KNOWN_HOSTS_PATH", "/run/keys/known_hosts") + include_debug_trace = os.getenv("APX_DEBUG_TRACE", "").strip().lower() in {"1", "true", "yes", "on"} if not key_path or not known_hosts_path: return { @@ -310,7 +311,7 @@ def apx_recipe_run(cmd:str, remote_ip_addr:str, remote_usr:str, recipe:str="code target_add_res = prepare_target(remote_ip_addr, remote_usr, key_path, apx_dir) if "error" in target_add_res: - return { + error_response = { "status": "error", "recipe": recipe, "stage": "target_prepare", @@ -322,10 +323,14 @@ def apx_recipe_run(cmd:str, remote_ip_addr:str, remote_usr:str, recipe:str="code "details": target_add_res.get("details", ""), "raw_output": target_add_res.get("raw_output", ""), } + if include_debug_trace: + error_response["debug_trace"] = target_add_res.get("debug_trace", []) + return error_response + prepare_debug_trace = target_add_res.get("debug_trace", []) run_res = run_workload(cmd, target_add_res["target_id"], recipe, apx_dir) if "error" in run_res: - return { + error_response = { "status": "error", "recipe": recipe, "stage": "workload_run", @@ -336,8 +341,19 @@ def apx_recipe_run(cmd:str, remote_ip_addr:str, remote_usr:str, recipe:str="code ), "details": run_res.get("details", ""), } + if include_debug_trace: + error_response["debug_trace"] = { + "prepare_target": prepare_debug_trace, + "run_workload": run_res.get("debug_trace", []), + } + return error_response results = get_results(run_res["run_id"], recipe, apx_dir) + if include_debug_trace: + results["debug_trace"] = { + "prepare_target": prepare_debug_trace, + "run_workload": run_res.get("debug_trace", []), + } return results diff --git a/mcp-local/tests/constants.py b/mcp-local/tests/constants.py index 677a687..3c6e707 100644 --- a/mcp-local/tests/constants.py +++ b/mcp-local/tests/constants.py @@ -243,3 +243,19 @@ }''' EXPECTED_CHECK_MCA_TOOL_RESPONSE_STATUS = "ok" + +CHECK_APX_RECIPE_RUN_REQUEST = { + "jsonrpc": "2.0", + "id": 8, + "method": "tools/call", + "params": { + "name": "apx_recipe_run", + "arguments": { + "cmd": "python3 /app/mcp-local/tests/test.py", + "remote_ip_addr": "localhost", + "remote_usr": "base", + "recipe": "code_hotspots", + "invocation_reason": "Run APX code hotspots recipe against the local test workload requested by the user.", + }, + }, + } diff --git a/mcp-local/tests/test_mcp.py b/mcp-local/tests/test_mcp.py index 096af6b..cf129f6 100644 --- a/mcp-local/tests/test_mcp.py +++ b/mcp-local/tests/test_mcp.py @@ -15,6 +15,9 @@ import json import constants import os +import shutil +import subprocess +import tempfile import time from pathlib import Path @@ -95,82 +98,152 @@ def test_mcp_stdio_transport_responds(platform): repo_root = Path(__file__).resolve().parents[1] print("\n***Repo Root: ", repo_root) - with ( - DockerContainer(image) - .with_volume_mapping(str(repo_root), "/workspace") - .with_kwargs(stdin_open=True, tty=False) - ) as container: - wait_for_logs(container, "Starting MCP server", timeout=60) - socket_wrapper = container.get_wrapped_container().attach_socket( - params={"stdin": 1, "stdout": 1, "stderr": 1, "stream": 1} - ) - raw_socket = socket_wrapper._sock - raw_socket.settimeout(10) - - raw_socket.sendall(_encode_mcp_message(constants.INIT_REQUEST)) - response = _read_mcp_message(raw_socket, timeout=20) - - #Check Container Init Test - assert response.get("id") == 1, "Test Failed: MCP initialize response id mismatch." - assert "result" in response, "Test Failed: MCP initialize response missing result field." - assert "serverInfo" in response["result"], "Test Failed: MCP initialize response missing serverInfo field." - raw_socket.sendall( - _encode_mcp_message({"jsonrpc": "2.0", "method": "initialized", "params": {}}) - ) - - def _read_response(expected_id: int, timeout: float = 10.0) -> dict: - deadline = time.time() + timeout - while time.time() < deadline: - message = _read_mcp_message(raw_socket, timeout=timeout) - if message.get("id") == expected_id: - return message - raise TimeoutError(f"Timed out waiting for MCP response id={expected_id}.") - - print("\n***Test Passed: arm-mcp container initilized and ran successfully") - - #Check Image Tool Test - raw_socket.sendall(_encode_mcp_message(constants.CHECK_IMAGE_REQUEST)) - check_image_response = _read_response(2, timeout=60) - assert check_image_response.get("result")["structuredContent"] == constants.EXPECTED_CHECK_IMAGE_RESPONSE, "Test Failed: MCP check_image tool failed: content mismatch. Expected: {}, Received: {}".format(json.dumps(constants.EXPECTED_CHECK_IMAGE_RESPONSE,indent=2), json.dumps(check_image_response.get("result")["structuredContent"],indent=2)) - print("\n***Test Passed: MCP check_image tool succeeded") - - #Check Skopeo Tool Test - raw_socket.sendall(_encode_mcp_message(constants.CHECK_SKOPEO_REQUEST)) - check_skopeo_response = _read_response(3, timeout=60) - actual_os = json.loads(check_skopeo_response.get("result")["structuredContent"]["stdout"]).get("Os") - actual_status = check_skopeo_response.get("result")["structuredContent"].get("status") - assert actual_os == json.loads(constants.EXPECTED_CHECK_SKOPEO_RESPONSE["stdout"]).get("Os"), "Test Failed: MCP check_skopeo tool failed: Os mismatch. Expected: {}, Received: {}".format(constants.EXPECTED_CHECK_SKOPEO_RESPONSE["Os"], actual_os) - assert actual_status == constants.EXPECTED_CHECK_SKOPEO_RESPONSE["status"], "Test Failed: MCP check_skopeo tool failed: Status mismatch. Expected: {}, Received: {}".format(constants.EXPECTED_CHECK_SKOPEO_RESPONSE["status"], actual_status) - print("\n***Test Passed: MCP check_skopeo tool succeeded") - - #Check NGINX Query Test - raw_socket.sendall(_encode_mcp_message(constants.CHECK_NGINX_REQUEST)) - check_nginx_response = _read_response(4, timeout=60) - urls = json.dumps(check_nginx_response["result"]["structuredContent"]) - assert any(expected in urls for expected in constants.EXPECTED_CHECK_NGINX_RESPONSE), "Test Failed: MCP check_nginx tool failed: content mismatch., Expected one of: {}, Received: {}".format(json.dumps(constants.EXPECTED_CHECK_NGINX_RESPONSE,indent=2), json.dumps(check_nginx_response.get("result")["structuredContent"],indent=2)) - print("\n***Test Passed: MCP check_nginx tool succeeded") - - #Check Migrate Ease Tool Test - raw_socket.sendall(_encode_mcp_message(constants.CHECK_MIGRATE_EASE_TOOL_REQUEST)) - check_migrate_ease_tool_response = _read_response(5, timeout=60) - #assert only the status field to avoid mismatches due to dynamic fields - assert check_migrate_ease_tool_response.get("result")["structuredContent"]["status"] == constants.EXPECTED_CHECK_MIGRATE_EASE_TOOL_RESPONSE_STATUS, "Test Failed: MCP check_migrate_ease_tool tool failed: status mismatch. Expected: {}, Received: {}".format(constants.EXPECTED_CHECK_MIGRATE_EASE_TOOL_RESPONSE_STATUS, check_migrate_ease_tool_response.get("result")["structuredContent"]["status"]) - print("\n***Test Passed: MCP check_migrate_ease_tool tool succeeded") - - #Check Sysreport Tool Test - raw_socket.sendall(_encode_mcp_message(constants.CHECK_SYSREPORT_TOOL_REQUEST)) - check_sysreport_response = _read_response(6, timeout=60) - assert check_sysreport_response.get("result")["structuredContent"] == constants.EXPECTED_CHECK_SYSREPORT_TOOL_RESPONSE, "Test Failed: MCP sysreport_instructions tool failed: content mismatch. Expected: {}, Received: {}".format(json.dumps(constants.EXPECTED_CHECK_SYSREPORT_TOOL_RESPONSE,indent=2), json.dumps(check_sysreport_response.get("result")["structuredContent"],indent=2)) - print("\n***Test Passed: MCP sysreport_instructions tool succeeded") - - #Check MCA Tool Test - works only on platform=linux/arm64 - if platform == constants.DEFAULT_PLATFORM: - raw_socket.sendall(_encode_mcp_message(constants.CHECK_MCA_TOOL_REQUEST)) - check_mca_response = _read_response(7, timeout=60) - assert check_mca_response.get("result")["structuredContent"]["status"] == constants.EXPECTED_CHECK_MCA_TOOL_RESPONSE_STATUS, "Test Failed: MCP mca tool failed: status mismatch.Expected: {}, Received: {}".format(json.dumps(constants.EXPECTED_CHECK_MCA_TOOL_RESPONSE_STATUS,indent=2), json.dumps(check_mca_response.get("result")["structuredContent"]["status"],indent=2)) - print("\n***Test Passed: MCP mca tool succeeded") + + with tempfile.TemporaryDirectory(prefix="apx-test-keys-") as temp_keys_dir: + temp_keys_path = Path(temp_keys_dir) + pem_path = temp_keys_path / "ssh-key.pem" + known_hosts_path = temp_keys_path / "known_hosts" + pub_key_path = temp_keys_path / "ssh-key.pem.pub" + + external_ssh_key_path = os.getenv("APX_TEST_SSH_KEY_PATH") + external_known_hosts_path = os.getenv("APX_TEST_KNOWN_HOSTS_PATH") + + if external_ssh_key_path: + shutil.copyfile(external_ssh_key_path, pem_path) + else: + subprocess.run( + ["ssh-keygen", "-t", "ed25519", "-N", "", "-f", str(pem_path)], + check=True, + capture_output=True, + text=True, + ) + if pub_key_path.exists(): + pub_key_path.unlink() + + if external_known_hosts_path: + shutil.copyfile(external_known_hosts_path, known_hosts_path) else: - print("\n***Test NA: MCP mca tool is not supported on this platform: {}".format(platform)) + known_hosts_path.write_text( + "172.17.0.1 ssh-ed25519 AAAAC3NzaC1lZDI1NTE5AAAAIFakeKnownHostKeyForIntegrationTestsOnly\n", + encoding="utf-8", + ) + + os.chmod(pem_path, 0o600) + os.chmod(known_hosts_path, 0o644) + + with ( + DockerContainer(image) + .with_volume_mapping(str(repo_root), "/workspace") + .with_volume_mapping(str(temp_keys_path), "/run/keys", mode="ro") + .with_env("SSH_KEY_PATH", "/tmp/ssh-key.pem") + .with_env("KNOWN_HOSTS_PATH", "/tmp/known_hosts") + .with_kwargs(stdin_open=True, tty=False) + ) as container: + prep_paths_cmd = ( + "cp /run/keys/ssh-key.pem /tmp/ssh-key.pem && " + "cp /run/keys/known_hosts /tmp/known_hosts && " + "chown 0:0 /tmp/ssh-key.pem /tmp/known_hosts && " + "chmod 600 /tmp/ssh-key.pem && " + "chmod 644 /tmp/known_hosts" + ) + prep_result = container.get_wrapped_container().exec_run( + ["/bin/sh", "-lc", prep_paths_cmd] + ) + prep_output = ( + prep_result.output.decode("utf-8", errors="replace") + if isinstance(prep_result.output, (bytes, bytearray)) + else str(prep_result.output) + ) + assert prep_result.exit_code == 0, ( + f"Failed to prepare SSH files in container. Exit code: {prep_result.exit_code}. " + f"Output: {prep_output}" + ) + + wait_for_logs(container, "Starting MCP server", timeout=60) + socket_wrapper = container.get_wrapped_container().attach_socket( + params={"stdin": 1, "stdout": 1, "stderr": 1, "stream": 1} + ) + raw_socket = socket_wrapper._sock + raw_socket.settimeout(10) + + raw_socket.sendall(_encode_mcp_message(constants.INIT_REQUEST)) + response = _read_mcp_message(raw_socket, timeout=20) + + #Check Container Init Test + assert response.get("id") == 1, "Test Failed: MCP initialize response id mismatch." + assert "result" in response, "Test Failed: MCP initialize response missing result field." + assert "serverInfo" in response["result"], "Test Failed: MCP initialize response missing serverInfo field." + raw_socket.sendall( + _encode_mcp_message({"jsonrpc": "2.0", "method": "initialized", "params": {}}) + ) + + def _read_response(expected_id: int, timeout: float = 10.0) -> dict: + deadline = time.time() + timeout + while time.time() < deadline: + message = _read_mcp_message(raw_socket, timeout=timeout) + if message.get("id") == expected_id: + return message + raise TimeoutError(f"Timed out waiting for MCP response id={expected_id}.") + + print("\n***Test Passed: arm-mcp container initilized and ran successfully") + + #Check Image Tool Test + raw_socket.sendall(_encode_mcp_message(constants.CHECK_IMAGE_REQUEST)) + check_image_response = _read_response(2, timeout=60) + assert check_image_response.get("result")["structuredContent"] == constants.EXPECTED_CHECK_IMAGE_RESPONSE, "Test Failed: MCP check_image tool failed: content mismatch. Expected: {}, Received: {}".format(json.dumps(constants.EXPECTED_CHECK_IMAGE_RESPONSE,indent=2), json.dumps(check_image_response.get("result")["structuredContent"],indent=2)) + print("\n***Test Passed: MCP check_image tool succeeded") + + #Check Skopeo Tool Test + raw_socket.sendall(_encode_mcp_message(constants.CHECK_SKOPEO_REQUEST)) + check_skopeo_response = _read_response(3, timeout=60) + actual_os = json.loads(check_skopeo_response.get("result")["structuredContent"]["stdout"]).get("Os") + actual_status = check_skopeo_response.get("result")["structuredContent"].get("status") + assert actual_os == json.loads(constants.EXPECTED_CHECK_SKOPEO_RESPONSE["stdout"]).get("Os"), "Test Failed: MCP check_skopeo tool failed: Os mismatch. Expected: {}, Received: {}".format(constants.EXPECTED_CHECK_SKOPEO_RESPONSE["Os"], actual_os) + assert actual_status == constants.EXPECTED_CHECK_SKOPEO_RESPONSE["status"], "Test Failed: MCP check_skopeo tool failed: Status mismatch. Expected: {}, Received: {}".format(constants.EXPECTED_CHECK_SKOPEO_RESPONSE["status"], actual_status) + print("\n***Test Passed: MCP check_skopeo tool succeeded") + + #Check NGINX Query Test + raw_socket.sendall(_encode_mcp_message(constants.CHECK_NGINX_REQUEST)) + check_nginx_response = _read_response(4, timeout=60) + urls = json.dumps(check_nginx_response["result"]["structuredContent"]) + assert any(expected in urls for expected in constants.EXPECTED_CHECK_NGINX_RESPONSE), "Test Failed: MCP check_nginx tool failed: content mismatch., Expected one of: {}, Received: {}".format(json.dumps(constants.EXPECTED_CHECK_NGINX_RESPONSE,indent=2), json.dumps(check_nginx_response.get("result")["structuredContent"],indent=2)) + print("\n***Test Passed: MCP check_nginx tool succeeded") + + #Check Migrate Ease Tool Test + raw_socket.sendall(_encode_mcp_message(constants.CHECK_MIGRATE_EASE_TOOL_REQUEST)) + check_migrate_ease_tool_response = _read_response(5, timeout=60) + #assert only the status field to avoid mismatches due to dynamic fields + assert check_migrate_ease_tool_response.get("result")["structuredContent"]["status"] == constants.EXPECTED_CHECK_MIGRATE_EASE_TOOL_RESPONSE_STATUS, "Test Failed: MCP check_migrate_ease_tool tool failed: status mismatch. Expected: {}, Received: {}".format(constants.EXPECTED_CHECK_MIGRATE_EASE_TOOL_RESPONSE_STATUS, check_migrate_ease_tool_response.get("result")["structuredContent"]["status"]) + print("\n***Test Passed: MCP check_migrate_ease_tool tool succeeded") + + #Check Sysreport Tool Test + raw_socket.sendall(_encode_mcp_message(constants.CHECK_SYSREPORT_TOOL_REQUEST)) + check_sysreport_response = _read_response(6, timeout=60) + assert check_sysreport_response.get("result")["structuredContent"] == constants.EXPECTED_CHECK_SYSREPORT_TOOL_RESPONSE, "Test Failed: MCP sysreport_instructions tool failed: content mismatch. Expected: {}, Received: {}".format(json.dumps(constants.EXPECTED_CHECK_SYSREPORT_TOOL_RESPONSE,indent=2), json.dumps(check_sysreport_response.get("result")["structuredContent"],indent=2)) + print("\n***Test Passed: MCP sysreport_instructions tool succeeded") + + #Check MCA Tool Test - works only on platform=linux/arm64 + if platform == constants.DEFAULT_PLATFORM: + raw_socket.sendall(_encode_mcp_message(constants.CHECK_MCA_TOOL_REQUEST)) + check_mca_response = _read_response(7, timeout=60) + assert check_mca_response.get("result")["structuredContent"]["status"] == constants.EXPECTED_CHECK_MCA_TOOL_RESPONSE_STATUS, "Test Failed: MCP mca tool failed: status mismatch.Expected: {}, Received: {}".format(json.dumps(constants.EXPECTED_CHECK_MCA_TOOL_RESPONSE_STATUS,indent=2), json.dumps(check_mca_response.get("result")["structuredContent"]["status"],indent=2)) + print("\n***Test Passed: MCP mca tool succeeded") + else: + print("\n***Test NA: MCP mca tool is not supported on this platform: {}".format(platform)) + + #Check APX Recipe Run Tool Test + apx_request = json.loads(json.dumps(constants.CHECK_APX_RECIPE_RUN_REQUEST)) + apx_args = apx_request["params"]["arguments"] + apx_args["remote_ip_addr"] = os.getenv("APX_TEST_REMOTE_IP", apx_args["remote_ip_addr"]) + apx_args["remote_usr"] = os.getenv("APX_TEST_REMOTE_USER", apx_args["remote_usr"]) + apx_args["cmd"] = os.getenv("APX_TEST_CMD", apx_args["cmd"]) + + raw_socket.sendall(_encode_mcp_message(apx_request)) + check_apx_recipe_run_response = _read_response(8, timeout=60) + apx_structured = check_apx_recipe_run_response.get("result", {}).get("structuredContent", {}) + assert apx_structured.get("recipe") == "code_hotspots", "Test Failed: MCP apx_recipe_run tool failed: recipe mismatch. Expected: code_hotspots, Received: {}".format(apx_structured.get("recipe")) + assert apx_structured.get("status") in {"success"}, "Test Failed: MCP apx_recipe_run tool failed: unexpected status. Received: {}".format(apx_structured.get("status")) + print("\n***Test Passed: MCP apx_recipe_run tool call completed") if __name__ == "__main__": pytest.main([__file__]) diff --git a/mcp-local/utils/apx.py b/mcp-local/utils/apx.py index b25e788..e62be96 100644 --- a/mcp-local/utils/apx.py +++ b/mcp-local/utils/apx.py @@ -9,6 +9,11 @@ QUERY_REGISTRY_PATH = Path(__file__).resolve().parent.parent / "sql" / "queries.sql" ANSI_ESCAPE_RE = re.compile(r"\x1B(?:[@-Z\\-_]|\[[0-?]*[ -/]*[@-~])") +SSH_PRIVATE_KEY_BLOCK_RE = re.compile( + r"-----BEGIN [A-Z ]*PRIVATE KEY-----.*?-----END [A-Z ]*PRIVATE KEY-----", + re.DOTALL, +) +SSH_KEY_PATH_RE = re.compile(r"(/[^\s:]+\.pem)\b") def load_recipe_query_map(sql_file_path: Path) -> Dict[str, Dict[str, str]]: @@ -92,6 +97,25 @@ def _sanitize_apx_output(output: str) -> str: return ANSI_ESCAPE_RE.sub("", output or "") +def _redact_sensitive_text(text: str) -> str: + if not text: + return "" + redacted = SSH_PRIVATE_KEY_BLOCK_RE.sub("[REDACTED_PRIVATE_KEY]", text) + redacted = SSH_KEY_PATH_RE.sub("[REDACTED_KEY_PATH]", redacted) + return redacted + + +def _redact_command(command: List[str]) -> List[str]: + redacted: List[str] = [] + for part in command: + if isinstance(part, str) and re.search(r"@[^\s:]+:\d+:[^\s]+\.pem$", part): + user_host, _, _ = part.rpartition(":") + redacted.append(f"{user_host}:[REDACTED_KEY_PATH]") + continue + redacted.append(_redact_sensitive_text(part) if isinstance(part, str) else part) + return redacted + + def _extract_session_id(render_output: str) -> Tuple[Optional[str], Optional[str]]: """Extract session_id from apx render output, trying full JSON then line-by-line JSON.""" clean_output = (render_output or "").strip() @@ -239,13 +263,13 @@ def _build_atp_error_response( "stage": stage, "message": message, "suggestion": suggestion, - "details": _trim_output(details), + "details": _trim_output(_redact_sensitive_text(details)), "warnings": [], } if query: response["query"] = query if raw_output: - response["raw_output"] = _trim_output(raw_output) + response["raw_output"] = _trim_output(_redact_sensitive_text(raw_output)) return response def extract_run_id(output: str) -> str: @@ -267,8 +291,7 @@ def run_command(command: list, cwd: str, parse_output=None) -> tuple: #print(command) result = subprocess.run(command, cwd=cwd, timeout=60*60*3, capture_output=True, text=True) except subprocess.TimeoutExpired as e: - print(f"Command timed out: {e}") - return -1, None + return -1, _redact_sensitive_text(str(e)) output = result.stdout if parse_output: @@ -283,36 +306,63 @@ def read_file_contents(file_path: str) -> str: def prepare_target(remote_ip_addr: str, remote_usr: str, ssh_key_path: str, apx_dir:str) -> dict: """Prepare the target machine for running workloads. Returns the target ID.""" - - #Check if target already exists + + debug_trace: List[Dict[str, Any]] = [] + + def _record_debug(command: List[str], status: int, output: Optional[str]) -> None: + debug_trace.append( + { + "command": _redact_command(command), + "status": status, + "output": _trim_output(_redact_sensitive_text(output or "")), + } + ) + + def _extract_targets(list_output: str) -> Dict[str, Any]: + if not list_output: + return {} + + candidates: List[str] = [list_output.strip()] + candidates.extend( + [line.strip() for line in list_output.splitlines() if line.strip().startswith("{")] + ) + + for candidate in candidates: + try: + data = json.loads(candidate) + except Exception: + continue + parsed_targets = data.get("data", {}) + if isinstance(parsed_targets, dict): + return parsed_targets + return {} + + canonical_host = "172.17.0.1" if remote_ip_addr in {"localhost", "127.0.0.1"} else remote_ip_addr + generated_name = f"{remote_usr}_{remote_ip_addr.replace('.', '_')}" + + # Check if target already exists list_command = ["./apx", "target", "list", "--json"] status, list_output = run_command(list_command, cwd=apx_dir) + _record_debug(list_command, status, list_output) if status == 0 and list_output: - try: - lines = list_output.strip().split("\n") - json_line = lines[1] if len(lines) > 1 else lines[0] - data = json.loads(json_line) - targets = data.get("data", {}) - for target_id, target_info in targets.items(): - value = target_info.get("value", {}) - jumps = value.get("jumps", []) - if not jumps: - continue - jump = jumps[0] - t_host = jump.get("host") - t_user = jump.get("username") - t_key = jump.get("private_key_filename") - if t_host == remote_ip_addr and t_user == remote_usr and t_key == ssh_key_path: - #print(f"Target already exists: {target_id}") - return { - "target_id": target_id - } - except Exception as e: - print(f"Failed to parse target list output: {e}") - - generated_name = f"{remote_usr}_{remote_ip_addr.replace('.', '_')}" + targets = _extract_targets(list_output) + for target_id, target_info in targets.items(): + value = target_info.get("value", {}) + jumps = value.get("jumps", []) + if not jumps: + continue + jump = jumps[0] + t_host = jump.get("host") + t_user = jump.get("username") + t_key = jump.get("private_key_filename") + if t_host == canonical_host and t_user == remote_usr and t_key == ssh_key_path: + return { + "target_id": target_id, + "debug_trace": debug_trace, + } + # Add the target if it doesn't exist - if remote_ip_addr in {"172.17.0.1", "localhost"}: + if remote_ip_addr in {"172.17.0.1", "localhost", "127.0.0.1"}: add_command = [ "./apx", "target", "add", f"{remote_usr}@172.17.0.1:22:{ssh_key_path}", @@ -325,28 +375,55 @@ def prepare_target(remote_ip_addr: str, remote_usr: str, ssh_key_path: str, apx_ "--name", generated_name ] add_status, add_output = run_command(add_command, cwd=apx_dir) + _record_debug(add_command, add_status, add_output) # Check for SSH key permission errors if add_output and ("engine.ssh.KEY_FILE_NOT_READABLE" in add_output): return { "error": "Check that the file permissions allow read access to the SSH key file. If ATP still cannot read the file, contact Arm support.", "details": f"Please run: chmod 0600 on your SSH key and then restart the mcp server.", - "raw_output": add_output + "raw_output": _redact_sensitive_text(add_output), + "debug_trace": debug_trace, } + if add_status != 0: + return { + "error": "Failed to add target before preparation.", + "details": _redact_sensitive_text(add_output or "apx target add returned a non-zero status."), + "raw_output": _redact_sensitive_text(add_output or ""), + "debug_trace": debug_trace, + } + + # Validate that target now exists before prepare to catch add/list state issues. + post_add_list_command = ["./apx", "target", "list", "--json"] + post_add_status, post_add_list_output = run_command(post_add_list_command, cwd=apx_dir) + _record_debug(post_add_list_command, post_add_status, post_add_list_output) + if post_add_status == 0: + post_add_targets = _extract_targets(post_add_list_output or "") + if generated_name not in post_add_targets: + return { + "error": "Target add reported success, but the target was not found in target list.", + "details": f"Expected target name '{generated_name}' was missing after add.", + "raw_output": _redact_sensitive_text(post_add_list_output or ""), + "debug_trace": debug_trace, + } + command = [ "./apx", "target", "prepare", "--target", f"{generated_name}" ] status, target_id = run_command(command, cwd=apx_dir) + _record_debug(command, status, target_id) if status != 0 or not target_id: return { "error": "Failed to prepare target. Check the connection details and make sure you have the correct username and ip address. Sometimes when you mean to connect to localhost, you are running from a docker container so the ip address needs to be 172.17.0.1", - "details": target_id + "details": _redact_sensitive_text(target_id or ""), + "debug_trace": debug_trace, } return { - "target_id": generated_name + "target_id": generated_name, + "debug_trace": debug_trace, } def run_workload(cmd:str, target: str, recipe:str, apx_dir:str) -> dict: @@ -354,24 +431,43 @@ def run_workload(cmd:str, target: str, recipe:str, apx_dir:str) -> dict: - 'Help my analyze my code's performance'. - 'Find the CPU hotspots in my application'. Returns the run ID of the workload execution.""" - + + debug_trace: List[Dict[str, Any]] = [] + + def _record_debug(command: List[str], status: int, output: Optional[str]) -> None: + debug_trace.append( + { + "command": _redact_command(command), + "status": status, + "output": _trim_output(_redact_sensitive_text(output or "")), + } + ) + # Check if the recipe is ready to run on the target ready_command = ["./apx", "recipe", "ready", recipe, "--target", target] ready_status, ready_output = run_command(ready_command, cwd=apx_dir) + _record_debug(ready_command, ready_status, ready_output) ready_output_text = (ready_output or "").lower() has_deploy_tools_hint = ( "--deploy-tools" in ready_output_text and "to deploy this tool on the target" in ready_output_text ) + has_missing_agent_hint = ( + "recipe is not ready to be run on your target machine" in ready_output_text + and "agent server" in ready_output_text + and "run `target prepare`" in ready_output_text + ) + is_expected_predeploy_state = has_deploy_tools_hint or has_missing_agent_hint # If readiness failed for reasons other than missing deployed tools, return early. # Missing tool deployment is expected because recipe run uses --deploy-tools. - if (ready_status != 0 or (ready_output and ready_output.strip())) and not has_deploy_tools_hint: + if (ready_status != 0 or (ready_output and ready_output.strip())) and not is_expected_predeploy_state: return { "error": "The recipe is not ready to run on the target machine.", - "details": ready_output if ready_output else "Recipe readiness check failed.", - "suggestion": "You may need to run 'target prepare' or use '--deploy-tools' flag." + "details": _redact_sensitive_text(ready_output) if ready_output else "Recipe readiness check failed.", + "suggestion": "You may need to run 'target prepare' or use '--deploy-tools' flag.", + "debug_trace": debug_trace, } command = [ @@ -383,14 +479,19 @@ def run_workload(cmd:str, target: str, recipe:str, apx_dir:str) -> dict: "--deploy-tools", "--param", "collect_java_stacks=true" ] status, output = run_command(command, cwd=apx_dir) + _record_debug(command, status, output) output_text = output or "" run_id = extract_run_id(output_text) if status == 0 else "" if not run_id or "Error" in output_text: return { - "error": output_text if output_text else "Failed to run workload.", - "details": output_text + "error": _redact_sensitive_text(output_text) if output_text else "Failed to run workload.", + "details": _redact_sensitive_text(output_text), + "debug_trace": debug_trace, } - return {"run_id": run_id} + return { + "run_id": run_id, + "debug_trace": debug_trace, + } def get_results(run_id: dict, recipe: str, apx_dir: str, default_table: str = "drilldown") -> Dict[str, Any]: """Get results from the target machine after running a workload. diff --git a/mcp-local/utils/tests/apx.pem b/mcp-local/utils/tests/apx.pem new file mode 100644 index 0000000..e69de29 diff --git a/mcp-local/utils/tests/known_hosts b/mcp-local/utils/tests/known_hosts new file mode 100644 index 0000000..e69de29 diff --git a/mcp-local/utils/tests/test.py b/mcp-local/utils/tests/test.py new file mode 100644 index 0000000..6d95fe9 --- /dev/null +++ b/mcp-local/utils/tests/test.py @@ -0,0 +1 @@ +print("Hello world") \ No newline at end of file