Merge pull request #737 from microsoft/psl-bug-rai1

fix: Resolve RAI reuse bug (#28329)

Author: Roopan-Microsoft

src/backend/v4/magentic_agents/common/lifecycle.py

@@ -43,6 +43,7 @@ def __init__(
         agent_description: str | None = None,
         agent_instructions: str | None = None,
         model_deployment_name: str | None = None,
+        project_client=None,
     ) -> None:
         self._stack: AsyncExitStack | None = None
         self.mcp_cfg: MCPConfig | None = mcp
@@ -58,6 +59,7 @@ def __init__(
         self.agent_description: str | None = agent_description
         self.agent_instructions: str | None = agent_instructions
         self.model_deployment_name: str | None = model_deployment_name
+        self.project_client = project_client
         self.logger = logging.getLogger(__name__)
 
     async def open(self) -> "MCPEnabledBase":
@@ -149,6 +151,55 @@ def get_chat_client(self, chat_client) -> AzureAIAgentClient:
         )
         return chat_client
 
+    async def resolve_agent_id(self, agent_id: str) -> Optional[str]:
+        """Resolve agent ID via Projects SDK first (for RAI agents), fallback to AgentsClient.
+
+        Args:
+            agent_id: The agent ID to resolve
+
+        Returns:
+            The resolved agent ID if found, None otherwise
+        """
+        # Try Projects SDK first (RAI agents were created via project_client)
+        try:
+            if self.project_client:
+                agent = await self.project_client.agents.get_agent(agent_id)
+                if agent and agent.id:
+                    self.logger.info(
+                        "RAI.AgentReuseSuccess: Resolved agent via Projects SDK (id=%s)",
+                        agent.id,
+                    )
+                    return agent.id
+        except Exception as ex:
+            self.logger.warning(
+                "RAI.AgentReuseMiss: Projects SDK get_agent failed (reason=ProjectsGetFailed, id=%s): %s",
+                agent_id,
+                ex,
+            )
+
+        # Fallback via AgentsClient (endpoint)
+        try:
+            if self.client:
+                agent = await self.client.get_agent(agent_id=agent_id)
+                if agent and agent.id:
+                    self.logger.info(
+                        "RAI.AgentReuseSuccess: Resolved agent via AgentsClient (id=%s)",
+                        agent.id,
+                    )
+                    return agent.id
+        except Exception as ex:
+            self.logger.warning(
+                "RAI.AgentReuseMiss: AgentsClient get_agent failed (reason=EndpointGetFailed, id=%s): %s",
+                agent_id,
+                ex,
+            )
+
+        self.logger.error(
+            "RAI.AgentReuseMiss: Agent ID not resolvable via any client (reason=ClientMismatch, id=%s)",
+            agent_id,
+        )
+        return None
+
     def get_agent_id(self, chat_client) -> str:
         """Return the underlying agent ID."""
         if chat_client and chat_client.agent_id is not None:
@@ -171,29 +222,69 @@ async def get_database_team_agent(self) -> Optional[AzureAIAgentClient]:
                 self.memory_store, self.team_config, self.agent_name
             )
 
-            if agent_id:
-                agent = await self.client.get_agent(agent_id=agent_id)
-                if agent and agent.id is not None:
-                    chat_client = AzureAIAgentClient(
-                        project_endpoint=self.project_endpoint,
-                        agent_id=agent.id,
-                        model_deployment_name=self.model_deployment_name,
-                        async_credential=self.creds,
-                    )
+            if not agent_id:
+                self.logger.info(
+                    "RAI reuse: no stored agent id (agent_name=%s)", self.agent_name
+                )
+                return None
+
+            # Use resolve_agent_id to try Projects SDK first, then AgentsClient
+            resolved = await self.resolve_agent_id(agent_id)
+            if not resolved:
+                self.logger.error(
+                    "RAI.AgentReuseMiss: stored id %s not resolvable (agent_name=%s)",
+                    agent_id,
+                    self.agent_name,
+                )
+                return None
+
+            # Create client with resolved ID, preferring project_client for RAI agents
+            if self.agent_name == "RAIAgent" and self.project_client:
+                chat_client = AzureAIAgentClient(
+                    project_client=self.project_client,
+                    agent_id=resolved,
+                    async_credential=self.creds,
+                )
+                self.logger.info(
+                    "RAI.AgentReuseSuccess: Created AzureAIAgentClient via Projects SDK (id=%s)",
+                    resolved,
+                )
+            else:
+                chat_client = AzureAIAgentClient(
+                    project_endpoint=self.project_endpoint,
+                    agent_id=resolved,
+                    model_deployment_name=self.model_deployment_name,
+                    async_credential=self.creds,
+                )
+                self.logger.info(
+                    "Created AzureAIAgentClient via endpoint (id=%s)", resolved
+                )
 
-        except (
-            Exception
-        ) as ex:  # Consider narrowing this to specific exceptions if possible
-            self.logger.error("Failed to initialize Get database team agent: %s", ex)
+        except Exception as ex:
+            self.logger.error(
+                "Failed to initialize Get database team agent (agent_name=%s): %s",
+                self.agent_name,
+                ex,
+            )
         return chat_client
 
     async def save_database_team_agent(self) -> None:
-        """Save current team agent to database."""
+        """Save current team agent to database (only if truly new or changed)."""
         try:
             if self._agent.id is None:
                 self.logger.error("Cannot save database team agent: agent_id is None")
                 return
 
+            # Check if stored ID matches current ID
+            stored_id = await get_database_team_agent_id(
+                self.memory_store, self.team_config, self.agent_name
+            )
+            if stored_id == self._agent.id:
+                self.logger.info(
+                    "RAI reuse: id unchanged (id=%s); skip save.", self._agent.id
+                )
+                return
+
             currentAgent = CurrentTeamAgent(
                 team_id=self.team_config.team_id,
                 team_name=self.team_config.name,
@@ -203,9 +294,14 @@ async def save_database_team_agent(self) -> None:
                 agent_instructions=self.agent_instructions,
             )
             await self.memory_store.add_team_agent(currentAgent)
+            self.logger.info(
+                "Saved team agent to database (agent_name=%s, id=%s)",
+                self.agent_name,
+                self._agent.id,
+            )
 
         except Exception as ex:
-            self.logger.error("Failed to save save database: %s", ex)
+            self.logger.error("Failed to save database: %s", ex)
 
     async def _prepare_mcp_tool(self) -> None:
         """Translate MCPConfig to a HostedMCPTool (agent_framework construct)."""
@@ -243,6 +339,7 @@ def __init__(
         agent_name: str | None = None,
         agent_description: str | None = None,
         agent_instructions: str | None = None,
+        project_client=None,
     ) -> None:
         super().__init__(
             mcp=mcp,
@@ -254,6 +351,7 @@ def __init__(
             agent_description=agent_description,
             agent_instructions=agent_instructions,
             model_deployment_name=model_deployment_name,
+            project_client=project_client,
         )
 
         self._created_ephemeral: bool = (

src/backend/v4/magentic_agents/foundry_agent.py

@@ -41,6 +41,9 @@ def __init__(
         team_config: TeamConfiguration | None = None,
         memory_store: DatabaseBase | None = None,
     ) -> None:
+        # Get project_client before calling super().__init__
+        project_client = config.get_ai_project_client()
+
         super().__init__(
             mcp=mcp_config,
             model_deployment_name=model_deployment_name,
@@ -51,12 +54,12 @@ def __init__(
             agent_name=agent_name,
             agent_description=agent_description,
             agent_instructions=agent_instructions,
+            project_client=project_client,
         )
 
         self.enable_code_interpreter = enable_code_interpreter
         self.search = search_config
         self.logger = logging.getLogger(__name__)
-        self.project_client = config.get_ai_project_client()
 
         # Decide early whether Azure Search mode should be activated
         self._use_azure_search = self._is_azure_search_requested()
@@ -269,10 +272,7 @@ async def _after_open(self) -> None:
                     temperature=temp,
                     model_id=self.model_deployment_name,
                 )
-
             self.logger.info("Initialized ChatAgent '%s'", self.agent_name)
-            if not chatClient:  # Only save if we didn't load from DB
-                await self.save_database_team_agent()
 
         except Exception as ex:
             self.logger.error("Failed to initialize ChatAgent: %s", ex)
@@ -299,7 +299,12 @@ async def invoke(self, prompt: str):
 
         messages = [ChatMessage(role=Role.USER, text=prompt)]
 
-        async for update in self._agent.run_stream(messages=messages):
+        agent_saved = False
+        async for update in self._agent.run_stream(messages):
+            # Save agent ID only once on first update (agent ID won't change during streaming)
+            if not agent_saved and self._agent.chat_client.agent_id:
+                await self.save_database_team_agent()
+                agent_saved = True
             yield update
 
     # -------------------------