Python: Fix azurefunctions MCP tool invocation to use correct agent (#3339)

* MCP tool fix for azurefunctions

* Moving logic to check for thread id

python/packages/azurefunctions/tests/test_app.py

@@ -1056,6 +1056,70 @@ class TestMCPToolEndpoint:
             with pytest.raises(RuntimeError, match="Agent execution failed"):
                 await app._handle_mcp_tool_invocation("TestAgent", context, client)
 
+    async def test_handle_mcp_tool_invocation_ignores_agent_name_in_thread_id(self) -> None:
+        """Test that MCP tool invocation uses the agent_name parameter, not the name from thread_id."""
+        mock_agent = Mock()
+        mock_agent.name = "PlantAdvisor"
+
+        app = AgentFunctionApp(agents=[mock_agent])
+        client = AsyncMock()
+
+        # Mock the entity response
+        mock_state = Mock()
+        mock_state.entity_state = {
+            "schemaVersion": "1.0.0",
+            "data": {"conversationHistory": []},
+        }
+        client.read_entity_state.return_value = mock_state
+
+        # Thread ID contains a different agent name (@StockAdvisor@poc123)
+        # but we're invoking PlantAdvisor - it should use PlantAdvisor's entity
+        context = json.dumps({"arguments": {"query": "test query", "threadId": "@StockAdvisor@test123"}})
+
+        with patch.object(app, "_get_response_from_entity") as get_response_mock:
+            get_response_mock.return_value = {"status": "success", "response": "Test response"}
+
+            await app._handle_mcp_tool_invocation("PlantAdvisor", context, client)
+
+            # Verify signal_entity was called with PlantAdvisor's entity, not StockAdvisor's
+            client.signal_entity.assert_called_once()
+            call_args = client.signal_entity.call_args
+            entity_id = call_args[0][0]
+
+            # Entity name should be dafx-PlantAdvisor, not dafx-StockAdvisor
+            assert entity_id.name == "dafx-PlantAdvisor"
+            assert entity_id.key == "test123"
+
+    async def test_handle_mcp_tool_invocation_uses_plain_thread_id_as_key(self) -> None:
+        """Test that a plain thread_id (not in @name@key format) is used as-is for the key."""
+        mock_agent = Mock()
+        mock_agent.name = "TestAgent"
+
+        app = AgentFunctionApp(agents=[mock_agent])
+        client = AsyncMock()
+
+        mock_state = Mock()
+        mock_state.entity_state = {
+            "schemaVersion": "1.0.0",
+            "data": {"conversationHistory": []},
+        }
+        client.read_entity_state.return_value = mock_state
+
+        # Plain thread_id without @name@key format
+        context = json.dumps({"arguments": {"query": "test query", "threadId": "simple-thread-123"}})
+
+        with patch.object(app, "_get_response_from_entity") as get_response_mock:
+            get_response_mock.return_value = {"status": "success", "response": "Test response"}
+
+            await app._handle_mcp_tool_invocation("TestAgent", context, client)
+
+            client.signal_entity.assert_called_once()
+            call_args = client.signal_entity.call_args
+            entity_id = call_args[0][0]
+
+            assert entity_id.name == "dafx-TestAgent"
+            assert entity_id.key == "simple-thread-123"
+
     def test_health_check_includes_mcp_tool_enabled(self) -> None:
         """Test that health check endpoint includes mcp_tool_enabled field."""
         mock_agent = Mock()

python/packages/azurefunctions/tests/test_models.py

@@ -120,6 +120,34 @@ class TestAgentSessionId:
         assert parsed.name == original.name
         assert parsed.key == original.key
 
+    def test_parse_with_agent_name_override(self) -> None:
+        """Test parsing @name@key format with agent_name parameter overrides the name."""
+        session_id = AgentSessionId.parse("@OriginalAgent@test-key-123", agent_name="OverriddenAgent")
+
+        assert session_id.name == "OverriddenAgent"
+        assert session_id.key == "test-key-123"
+
+    def test_parse_without_agent_name_uses_parsed_name(self) -> None:
+        """Test parsing @name@key format without agent_name uses name from string."""
+        session_id = AgentSessionId.parse("@ParsedAgent@test-key-123")
+
+        assert session_id.name == "ParsedAgent"
+        assert session_id.key == "test-key-123"
+
+    def test_parse_plain_string_with_agent_name(self) -> None:
+        """Test parsing plain string with agent_name uses entire string as key."""
+        session_id = AgentSessionId.parse("simple-thread-123", agent_name="TestAgent")
+
+        assert session_id.name == "TestAgent"
+        assert session_id.key == "simple-thread-123"
+
+    def test_parse_plain_string_without_agent_name_raises(self) -> None:
+        """Test parsing plain string without agent_name raises ValueError."""
+        with pytest.raises(ValueError) as exc_info:
+            AgentSessionId.parse("simple-thread-123")
+
+        assert "Invalid agent session ID format" in str(exc_info.value)
+
     def test_to_entity_name_adds_prefix(self) -> None:
         """Test that to_entity_name adds the dafx- prefix."""
         entity_name = AgentSessionId.to_entity_name("TestAgent")