feat: Add workspace_plugin parameter to agents for improved workspace management

Author: devanshjainms

src/agents/agents/action_planner_agent.py

@@ -59,10 +59,11 @@ def __init__(
         action_planner_plugin: Optional[ActionPlannerPlugin] = None,
         test_planner_plugin: Optional[TestPlannerPlugin] = None,
         troubleshooting_plugin: Optional[TroubleshootingPlugin] = None,
+        workspace_plugin: Optional[WorkspacePlugin] = None,
     ) -> None:
         action_planner = action_planner_plugin or ActionPlannerPlugin()
         test_planner = test_planner_plugin or TestPlannerPlugin()
-        workspace_plugin = WorkspacePlugin(workspace_store)
+        ws_plugin = workspace_plugin or WorkspacePlugin(workspace_store)
         troubleshooting = troubleshooting_plugin or TroubleshootingPlugin(workspace_store)
 
         super().__init__(
@@ -73,7 +74,7 @@ def __init__(
             ),
             kernel=kernel,
             instructions=ACTION_PLANNER_AGENT_SYSTEM_PROMPT,
-            plugins=[action_planner, test_planner, workspace_plugin, troubleshooting],
+            plugins=[action_planner, test_planner, ws_plugin, troubleshooting],
         )
 
         self.workspace_store: WorkspaceStore = workspace_store

src/agents/agents/base.py

@@ -246,18 +246,31 @@ def create_default_agent_registry(kernel: Optional[Kernel] = None) -> "AgentRegi
 
     registry = AgentRegistry()
     registry.register(EchoAgentSK(kernel=kernel))
-    registry.register(TestAdvisorAgentSK(kernel=kernel, workspace_store=workspace_store))
+    registry.register(
+        TestAdvisorAgentSK(
+            kernel=kernel,
+            workspace_store=workspace_store,
+            workspace_plugin=workspace_plugin,
+        )
+    )
     registry.register(
         ActionPlannerAgentSK(
             kernel=kernel,
             workspace_store=workspace_store,
+            workspace_plugin=workspace_plugin,
             troubleshooting_plugin=TroubleshootingPlugin(
                 workspace_store=workspace_store,
                 execution_plugin=execution_plugin,
             ),
         )
     )
-    registry.register(SystemContextAgentSK(kernel=kernel, workspace_store=workspace_store))
+    registry.register(
+        SystemContextAgentSK(
+            kernel=kernel,
+            workspace_store=workspace_store,
+            workspace_plugin=workspace_plugin,
+        )
+    )
     registry.register(
         ActionExecutorAgent(
             kernel=kernel,

src/agents/agents/system_context_agent.py

@@ -5,13 +5,10 @@
 
 This agent uses Semantic Kernel for agentic workspace management
 with function calling, replacing the custom tool loop.
-
-NOTE: This agent has NO hardcoded methods - all workspace operations
-are performed by the LLM autonomously calling WorkspacePlugin tools.
 """
 
+from typing import Optional
 from semantic_kernel import Kernel
-
 from src.agents.agents.base import SAPAutomationAgent
 from src.agents.workspace.workspace_store import WorkspaceStore
 from src.agents.plugins.workspace import WorkspacePlugin
@@ -35,15 +32,22 @@ class SystemContextAgentSK(SAPAutomationAgent):
     - create_workspace(): Create new workspace
     """
 
-    def __init__(self, kernel: Kernel, workspace_store: WorkspaceStore) -> None:
+    def __init__(
+        self,
+        kernel: Kernel,
+        workspace_store: WorkspaceStore,
+        workspace_plugin: Optional[WorkspacePlugin] = None,
+    ) -> None:
         """Initialize SystemContextAgent with Semantic Kernel.
 
         :param kernel: Configured Semantic Kernel instance
         :type kernel: Kernel
         :param workspace_store: WorkspaceStore instance for managing workspaces
         :type workspace_store: WorkspaceStore
+        :param workspace_plugin: Optional shared WorkspacePlugin (with KeyVault)
+        :type workspace_plugin: Optional[WorkspacePlugin]
         """
-        workspace_plugin = WorkspacePlugin(workspace_store)
+        ws_plugin = workspace_plugin or WorkspacePlugin(workspace_store)
         super().__init__(
             name="system_context",
             description=(
@@ -53,7 +57,7 @@ def __init__(self, kernel: Kernel, workspace_store: WorkspaceStore) -> None:
             ),
             kernel=kernel,
             instructions=SYSTEM_CONTEXT_AGENT_SYSTEM_PROMPT,
-            plugins=[workspace_plugin],
+            plugins=[ws_plugin],
         )
 
         self.workspace_store: WorkspaceStore = workspace_store

src/agents/agents/test_advisor_agent.py

@@ -7,9 +7,8 @@
 """
 
 from __future__ import annotations
-
+from typing import Optional
 from semantic_kernel import Kernel
-
 from src.agents.agents.base import SAPAutomationAgent
 from src.agents.observability import get_logger
 from src.agents.plugins.test import TestPlannerPlugin
@@ -28,9 +27,14 @@ class TestAdvisorAgentSK(SAPAutomationAgent):
     - WorkspacePlugin: Resolve SIDs, read workspace configuration (hosts.yaml, sap-parameters.yaml)
     """
 
-    def __init__(self, kernel: Kernel, workspace_store: WorkspaceStore):
+    def __init__(
+        self,
+        kernel: Kernel,
+        workspace_store: WorkspaceStore,
+        workspace_plugin: Optional[WorkspacePlugin] = None,
+    ):
         test_plugin = TestPlannerPlugin()
-        workspace_plugin = WorkspacePlugin(workspace_store)
+        ws_plugin = workspace_plugin or WorkspacePlugin(workspace_store)
 
         super().__init__(
             name="test_advisor",
@@ -40,7 +44,7 @@ def __init__(self, kernel: Kernel, workspace_store: WorkspaceStore):
             ),
             kernel=kernel,
             instructions=TEST_ADVISOR_AGENT_SYSTEM_PROMPT,
-            plugins=[test_plugin, workspace_plugin],
+            plugins=[test_plugin, ws_plugin],
         )
 
         self.workspace_store: WorkspaceStore = workspace_store