refactor: Update SSH key retrieval process and improve documentation for KeyVault integration

Author: devanshjainms

src/agents/plugins/workspace.py

@@ -238,7 +238,6 @@ def get_system_configuration(self, workspace_id: Annotated[str, "Workspace name/
             "sources": [],
         }
 
-        # Read sap-parameters.yaml
         sap_params = self.store.read_file(workspace_id, "sap-parameters.yaml")
         if sap_params:
             try:
@@ -261,12 +260,10 @@ def get_system_configuration(self, workspace_id: Annotated[str, "Workspace name/
             except Exception:
                 result["sources"].append({"file": "sap-parameters.yaml", "error": "parse_error"})
 
-        # Read hosts.yaml
         hosts_raw = self.store.read_file(workspace_id, "hosts.yaml")
         if hosts_raw:
             try:
                 hosts_parsed = yaml.safe_load(hosts_raw)
-                # Expecting host groups or list
                 result["hosts"] = hosts_parsed if hosts_parsed else []
                 result["sources"].append(
                     {
@@ -373,7 +370,7 @@ def resolve_ssh_key(self, workspace_id: Annotated[str, "Workspace name/ID"]) ->
         for filename in files:
             filename_lower = filename.lower()
             if any(pattern in filename_lower for pattern in key_patterns):
-                if not filename.endswith(".pub"):  # Skip public keys
+                if not filename.endswith(".pub"):
                     key_path = workspace_path / filename
                     if key_path.exists():
                         logger.info(f"Resolved SSH key: {key_path}")
@@ -482,27 +479,19 @@ def get_execution_context(self, workspace_id: Annotated[str, "Workspace name/ID"
         if not result["ssh_key_path"] and self.keyvault_plugin:
             secret_id = result["sap_parameters"].get("secret_id", "")
             if secret_id:
-                logger.info(
-                    f"No local SSH key found, attempting to fetch from Key Vault "
-                    f"using secret_id: {secret_id}"
-                )
                 try:
                     parse_result = json.loads(
                         self.keyvault_plugin.parse_key_vault_id_and_secret_id(secret_id)
                     )
                     if "error" not in parse_result:
                         vault_name = parse_result.get("vault_name")
                         secret_name = parse_result.get("secret_name")
-                        managed_identity_client_id = result["sap_parameters"].get(
-                            "user_assigned_identity_client_id", ""
-                        )
                         if vault_name and secret_name:
                             key_result = json.loads(
                                 self.keyvault_plugin.get_ssh_private_key(
                                     secret_name=secret_name,
                                     vault_name=vault_name,
                                     key_filename=f"{workspace_id}_id_rsa",
-                                    managed_identity_client_id=managed_identity_client_id,
                                 )
                             )
                             if "key_path" in key_result:

src/agents/prompts.py

@@ -328,7 +328,8 @@
 
 EXECUTION TOOLS:
 - get_execution_context: Get ALL workspace context in ONE call (inventory, parameters, SSH key)
-- get_ssh_key_for_workspace: Fetch SSH key from Azure KeyVault when missing locally
+- parse_key_vault_id_and_secret_id: Parse KeyVault URL to extract vault_name and secret_name
+- get_ssh_private_key: Fetch SSH key from Azure KeyVault (requires vault_name, secret_name)
 - run_test_by_id: Run tests (auto-resolves SSH key and parameters)
 - run_readonly_command: Run diagnostic commands on SAP VMs (auto-resolves SSH key)
 - tail_log: Tail logs
@@ -489,7 +490,7 @@
 
 ERROR HANDLING:
 - Host unreachable: "Can't reach the host. Check if it's running and network is accessible."
-- SSH key missing: Check sap-parameters.yaml for kv_name. If present, call get_ssh_key_for_workspace(). If no KeyVault, ask user for key path.
+- SSH key missing: Read sap-parameters.yaml, find secret_id URL, parse with parse_key_vault_id_and_secret_id(), then call get_ssh_private_key(secret_name, vault_name).
 - Test not found: "That test doesn't exist. Available tests: [list]"
 - Keep errors user-friendly