diff --git a/plan/src/main/java/org/apache/flink/agents/plan/AgentPlan.java b/plan/src/main/java/org/apache/flink/agents/plan/AgentPlan.java index 054b2ed26..562bc5238 100644 --- a/plan/src/main/java/org/apache/flink/agents/plan/AgentPlan.java +++ b/plan/src/main/java/org/apache/flink/agents/plan/AgentPlan.java @@ -295,41 +295,7 @@ private void extractJavaMCPServer(Method method) throws Exception { descriptor.getModule(), JAVA_MCP_SERVER_CLASS_NAME, new HashMap<>(descriptor.getInitialArguments())); - JavaResourceProvider provider = new JavaResourceProvider(name, MCP_SERVER, descriptor); - - addResourceProvider(provider); - Object mcpServer = provider.provide(null); - - // Call listTools() via reflection - Method listToolsMethod = mcpServer.getClass().getMethod("listTools"); - @SuppressWarnings("unchecked") - Iterable tools = - (Iterable) listToolsMethod.invoke(mcpServer); - - for (SerializableResource tool : tools) { - Method getNameMethod = tool.getClass().getMethod("getName"); - String toolName = (String) getNameMethod.invoke(tool); - addResourceProvider( - JavaSerializableResourceProvider.createResourceProvider(toolName, TOOL, tool)); - } - - // Call listPrompts() via reflection - Method listPromptsMethod = mcpServer.getClass().getMethod("listPrompts"); - @SuppressWarnings("unchecked") - Iterable prompts = - (Iterable) listPromptsMethod.invoke(mcpServer); - - for (SerializableResource prompt : prompts) { - Method getNameMethod = prompt.getClass().getMethod("getName"); - String promptName = (String) getNameMethod.invoke(prompt); - addResourceProvider( - JavaSerializableResourceProvider.createResourceProvider( - promptName, PROMPT, prompt)); - } - - // Call close() via reflection - Method closeMethod = mcpServer.getClass().getMethod("close"); - closeMethod.invoke(mcpServer); + addResourceProvider(new JavaResourceProvider(name, MCP_SERVER, descriptor)); } private void extractResourceProvidersFromAgent(Agent agent) throws Exception { diff --git a/plan/src/test/java/org/apache/flink/agents/plan/AgentPlanDeclareMCPServerTest.java b/plan/src/test/java/org/apache/flink/agents/plan/AgentPlanDeclareMCPServerTest.java index e3c4fe2e4..9158a5d66 100644 --- a/plan/src/test/java/org/apache/flink/agents/plan/AgentPlanDeclareMCPServerTest.java +++ b/plan/src/test/java/org/apache/flink/agents/plan/AgentPlanDeclareMCPServerTest.java @@ -24,14 +24,12 @@ import org.apache.flink.agents.api.agents.Agent; import org.apache.flink.agents.api.annotation.Action; import org.apache.flink.agents.api.context.RunnerContext; -import org.apache.flink.agents.api.prompt.Prompt; -import org.apache.flink.agents.api.resource.Resource; import org.apache.flink.agents.api.resource.ResourceDescriptor; import org.apache.flink.agents.api.resource.ResourceName; import org.apache.flink.agents.api.resource.ResourceType; -import org.apache.flink.agents.api.tools.Tool; import org.apache.flink.agents.api.tools.ToolMetadata; import org.apache.flink.agents.integrations.mcp.MCPPrompt; +import org.apache.flink.agents.integrations.mcp.MCPServer; import org.apache.flink.agents.integrations.mcp.MCPTool; import org.apache.flink.agents.plan.resourceprovider.ResourceProvider; import org.junit.jupiter.api.*; @@ -48,14 +46,17 @@ import static org.junit.jupiter.api.Assertions.assertFalse; import static org.junit.jupiter.api.Assertions.assertInstanceOf; import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertNull; import static org.junit.jupiter.api.Assertions.assertTrue; import static org.junit.jupiter.api.Assumptions.assumeTrue; /** * Tests for MCP server integration with AgentPlan. * - *

This test verifies that MCP servers, tools, and prompts are properly discovered and registered - * in the agent plan, following the pattern from {@link AgentPlanDeclareToolMethodTest}. + *

Verifies that MCP servers are registered in the agent plan at compile time, while tool and + * prompt discovery is deferred to operator startup (runtime). Tool/prompt retrieval tests + * instantiate the MCPServer directly from the plan's provider to simulate what {@code + * JavaMCPResourceDiscovery} does at runtime. * *

Uses the Python MCP server from python/flink_agents/api/tests/mcp/mcp_server.py. */ @@ -185,16 +186,14 @@ void setup() throws Exception { agentPlan = new AgentPlan(new TestMCPAgent()); } - /** Resolves a resource directly from its provider. */ - private Resource resolveResource(String name, ResourceType type) throws Exception { - return agentPlan - .getResourceProviders() - .get(type) - .get(name) - .provide( - (n, t) -> { - throw new UnsupportedOperationException("No dependencies expected"); - }); + /** + * Returns an MCPServer instantiated from the plan's provider, simulating what + * JavaMCPResourceDiscovery does at operator startup. + */ + private MCPServer instantiateMCPServer() throws Exception { + ResourceProvider provider = + agentPlan.getResourceProviders().get(ResourceType.MCP_SERVER).get("testMcpServer"); + return (MCPServer) provider.provide(null); } @AfterAll @@ -211,7 +210,7 @@ static void afterAll() { @Test @DisabledOnJre(JRE.JAVA_11) - @DisplayName("Discover @MCPServer method and register MCP server") + @DisplayName("Discover @MCPServer method and register MCP server provider in plan") void discoverMCPServer() { Map> providers = agentPlan.getResourceProviders(); @@ -222,116 +221,148 @@ void discoverMCPServer() { @Test @DisabledOnJre(JRE.JAVA_11) - @DisplayName("Discover and register tools from MCP server") + @DisplayName("Tools are NOT in AgentPlan providers — discovery is deferred to operator startup") void discoverToolsFromMCPServer() { Map> providers = agentPlan.getResourceProviders(); - assertTrue(providers.containsKey(ResourceType.TOOL)); - - Map toolProviders = providers.get(ResourceType.TOOL); - assertTrue(toolProviders.containsKey("add"), "add tool should be discovered"); - assertEquals(1, toolProviders.size(), "Should have exactly 1 tool from Python server"); + // Tools are discovered at runtime by JavaMCPResourceDiscovery, not during plan construction + assertNull( + providers.get(ResourceType.TOOL), + "TOOL providers should be absent from AgentPlan; discovery is deferred to runtime"); } @Test @DisabledOnJre(JRE.JAVA_11) - @DisplayName("Discover and register prompts from MCP server") + @DisplayName( + "Prompts are NOT in AgentPlan providers — discovery is deferred to operator startup") void discoverPromptsFromMCPServer() { Map> providers = agentPlan.getResourceProviders(); - assertTrue(providers.containsKey(ResourceType.PROMPT)); - - Map promptProviders = providers.get(ResourceType.PROMPT); - assertTrue(promptProviders.containsKey("ask_sum"), "ask_sum prompt should be discovered"); - assertEquals(1, promptProviders.size(), "Should have exactly 1 prompt from Python server"); + // Prompts are discovered at runtime by JavaMCPResourceDiscovery, not during plan + // construction + assertNull( + providers.get(ResourceType.PROMPT), + "PROMPT providers should be absent from AgentPlan; discovery is deferred to runtime"); } @Test @DisabledOnJre(JRE.JAVA_11) - @DisplayName("Retrieve MCP tool from AgentPlan - add tool") + @DisplayName("Retrieve MCP tool at runtime - add tool") void retrieveMCPToolAdd() throws Exception { - Tool tool = (Tool) resolveResource("add", ResourceType.TOOL); - assertNotNull(tool); - assertInstanceOf(MCPTool.class, tool); - - MCPTool mcpTool = (MCPTool) tool; - assertEquals("add", mcpTool.getName()); - // Verify description starts with expected text - assertTrue( - mcpTool.getMetadata() - .getDescription() - .startsWith("Get the detailed information of a specified IP address."), - "Description should start with expected text"); - // Verify input schema contains expected parameters - String schema = mcpTool.getMetadata().getInputSchema(); - assertTrue(schema.contains("a"), "Schema should contain parameter 'a'"); - assertTrue(schema.contains("b"), "Schema should contain parameter 'b'"); + MCPServer server = instantiateMCPServer(); + try { + MCPTool tool = null; + for (MCPTool t : server.listTools()) { + if ("add".equals(t.getName())) { + tool = t; + break; + } + } + assertNotNull(tool, "add tool should be discoverable from MCPServer"); + assertInstanceOf(MCPTool.class, tool); + assertEquals("add", tool.getName()); + assertTrue( + tool.getMetadata() + .getDescription() + .startsWith("Get the detailed information of a specified IP address."), + "Description should start with expected text"); + String schema = tool.getMetadata().getInputSchema(); + assertTrue(schema.contains("a"), "Schema should contain parameter 'a'"); + assertTrue(schema.contains("b"), "Schema should contain parameter 'b'"); + } finally { + server.close(); + } } @Test @DisabledOnJre(JRE.JAVA_11) - @DisplayName("Retrieve MCP prompt from AgentPlan - ask_sum") + @DisplayName("Retrieve MCP prompt at runtime - ask_sum") void retrieveMCPPromptAskSum() throws Exception { - Prompt prompt = (Prompt) resolveResource("ask_sum", ResourceType.PROMPT); - assertNotNull(prompt); - assertInstanceOf(MCPPrompt.class, prompt); - - MCPPrompt mcpPrompt = (MCPPrompt) prompt; - assertEquals("ask_sum", mcpPrompt.getName()); - assertEquals("Prompt of add tool.", mcpPrompt.getDescription()); - // ask_sum prompt should have 'a' and 'b' as arguments - Map args = mcpPrompt.getPromptArguments(); - assertTrue(args.containsKey("a"), "Should have 'a' argument"); - assertTrue(args.containsKey("b"), "Should have 'b' argument"); + MCPServer server = instantiateMCPServer(); + try { + MCPPrompt prompt = null; + for (MCPPrompt p : server.listPrompts()) { + if ("ask_sum".equals(p.getName())) { + prompt = p; + break; + } + } + assertNotNull(prompt, "ask_sum prompt should be discoverable from MCPServer"); + assertInstanceOf(MCPPrompt.class, prompt); + assertEquals("ask_sum", prompt.getName()); + assertEquals("Prompt of add tool.", prompt.getDescription()); + Map args = prompt.getPromptArguments(); + assertTrue(args.containsKey("a"), "Should have 'a' argument"); + assertTrue(args.containsKey("b"), "Should have 'b' argument"); + } finally { + server.close(); + } } @Test @DisabledOnJre(JRE.JAVA_11) - @DisplayName("AgentPlan JSON serialization with MCP resources") + @DisplayName( + "AgentPlan JSON serialization contains MCPServer descriptor, not tool/prompt entries") void testAgentPlanJsonSerializableWithMCP() throws Exception { ObjectMapper mapper = new ObjectMapper(); String json = mapper.writeValueAsString(agentPlan); - // Verify JSON contains MCP resources - assertTrue(json.contains("add"), "JSON should contain add tool"); - assertTrue(json.contains("ask_sum"), "JSON should contain ask_sum prompt"); + // Serialized plan contains the MCPServer configuration assertTrue(json.contains("mcp_server"), "JSON should contain mcp_server type"); + assertTrue(json.contains("testMcpServer"), "JSON should contain the server provider name"); + assertTrue(json.contains(MCP_ENDPOINT), "JSON should contain the endpoint"); + + // Tools and prompts are NOT serialized into the plan (they are runtime-discovered) + assertFalse( + json.contains("\"add\"") && json.contains("java_serializable"), + "JSON should not contain a serialized 'add' tool provider"); - // Verify serialization works without errors + // Verify serialization/deserialization roundtrip works without errors assertNotNull(json); assertFalse(json.isEmpty()); } @Test @DisabledOnJre(JRE.JAVA_11) - @DisplayName("Test MCP server is closed after discovery") - void testMCPServerClosedAfterDiscovery() throws Exception { - // The MCPServer.close() should be called after listTools() and listPrompts() - // We verify this indirectly by checking that the plan was created successfully - assertNotNull(agentPlan); - assertTrue(agentPlan.getResourceProviders().containsKey(ResourceType.MCP_SERVER)); - assertTrue(agentPlan.getResourceProviders().containsKey(ResourceType.TOOL)); - assertTrue(agentPlan.getResourceProviders().containsKey(ResourceType.PROMPT)); + @DisplayName("AgentPlan construction does not make network calls to MCP server") + void testNoNetworkCallsDuringPlanBuild() { + Map> providers = + agentPlan.getResourceProviders(); + assertNull(providers.get(ResourceType.TOOL), "No TOOL providers expected in plan"); + assertNull(providers.get(ResourceType.PROMPT), "No PROMPT providers expected in plan"); + assertTrue( + providers.containsKey(ResourceType.MCP_SERVER), + "MCP_SERVER provider must still be in plan for runtime discovery"); } @Test @DisabledOnJre(JRE.JAVA_11) @DisplayName("Test metadata from MCP tool - add") void testMCPToolMetadata() throws Exception { - Tool tool = (Tool) resolveResource("add", ResourceType.TOOL); - ToolMetadata metadata = tool.getMetadata(); - - assertEquals("add", metadata.getName()); - // Verify description starts with expected text (full docstring includes Args/Returns) - assertTrue( - metadata.getDescription() - .startsWith("Get the detailed information of a specified IP address."), - "Description should start with expected text"); - assertNotNull(metadata.getInputSchema()); - - String schema = metadata.getInputSchema(); - // Verify the tool has expected parameters - assertTrue(schema.contains("a"), "Schema should contain 'a' parameter"); - assertTrue(schema.contains("b"), "Schema should contain 'b' parameter"); + MCPServer server = instantiateMCPServer(); + try { + MCPTool tool = null; + for (MCPTool t : server.listTools()) { + if ("add".equals(t.getName())) { + tool = t; + break; + } + } + assertNotNull(tool); + ToolMetadata metadata = tool.getMetadata(); + + assertEquals("add", metadata.getName()); + assertTrue( + metadata.getDescription() + .startsWith("Get the detailed information of a specified IP address."), + "Description should start with expected text"); + assertNotNull(metadata.getInputSchema()); + + String schema = metadata.getInputSchema(); + assertTrue(schema.contains("a"), "Schema should contain 'a' parameter"); + assertTrue(schema.contains("b"), "Schema should contain 'b' parameter"); + } finally { + server.close(); + } } } diff --git a/runtime/src/main/java/org/apache/flink/agents/runtime/JavaMCPResourceDiscovery.java b/runtime/src/main/java/org/apache/flink/agents/runtime/JavaMCPResourceDiscovery.java new file mode 100644 index 000000000..0b74b244b --- /dev/null +++ b/runtime/src/main/java/org/apache/flink/agents/runtime/JavaMCPResourceDiscovery.java @@ -0,0 +1,89 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.flink.agents.runtime; + +import org.apache.flink.agents.api.resource.Resource; +import org.apache.flink.agents.api.resource.ResourceType; +import org.apache.flink.agents.plan.resourceprovider.JavaResourceProvider; +import org.apache.flink.agents.plan.resourceprovider.ResourceProvider; + +import java.lang.reflect.Method; +import java.util.Map; + +import static org.apache.flink.agents.api.resource.ResourceType.MCP_SERVER; +import static org.apache.flink.agents.api.resource.ResourceType.PROMPT; +import static org.apache.flink.agents.api.resource.ResourceType.TOOL; + +/** + * Discovers tools and prompts from Java MCP servers and registers them in a ResourceCache. + * + *

Called once during operator initialization, immediately after the ResourceCache is created. + * Uses reflection throughout to preserve Java 11 compatibility (MCP classes are conditionally + * compiled for Java 17+). + */ +public class JavaMCPResourceDiscovery { + + /** + * Initializes Java MCP servers from the resource providers, extracts their tools and prompts, + * and registers them in the cache. + * + * @param resourceProviders the resource providers from the agent plan + * @param cache the resource cache to register discovered resources in + * @throws Exception if a Java MCP server fails to initialize or discovery fails + */ + public static void discoverJavaMCPResources( + Map> resourceProviders, ResourceCache cache) + throws Exception { + + Map servers = resourceProviders.get(MCP_SERVER); + if (servers == null) { + return; + } + + for (ResourceProvider rp : servers.values()) { + if (!(rp instanceof JavaResourceProvider)) { + continue; + } + + Object mcpServer = rp.provide(null); + + Method listToolsMethod = mcpServer.getClass().getMethod("listTools"); + @SuppressWarnings("unchecked") + Iterable tools = (Iterable) listToolsMethod.invoke(mcpServer); + for (Resource tool : tools) { + String toolName = (String) tool.getClass().getMethod("getName").invoke(tool); + cache.put(toolName, TOOL, tool); + } + + Method supportsPromptsMethod = mcpServer.getClass().getMethod("supportsPrompts"); + boolean supportsPrompts = (Boolean) supportsPromptsMethod.invoke(mcpServer); + if (supportsPrompts) { + Method listPromptsMethod = mcpServer.getClass().getMethod("listPrompts"); + @SuppressWarnings("unchecked") + Iterable prompts = + (Iterable) listPromptsMethod.invoke(mcpServer); + for (Resource prompt : prompts) { + String promptName = + (String) prompt.getClass().getMethod("getName").invoke(prompt); + cache.put(promptName, PROMPT, prompt); + } + } + } + } +} diff --git a/runtime/src/main/java/org/apache/flink/agents/runtime/operator/ActionExecutionOperator.java b/runtime/src/main/java/org/apache/flink/agents/runtime/operator/ActionExecutionOperator.java index e5015a3a5..b0735c2c4 100644 --- a/runtime/src/main/java/org/apache/flink/agents/runtime/operator/ActionExecutionOperator.java +++ b/runtime/src/main/java/org/apache/flink/agents/runtime/operator/ActionExecutionOperator.java @@ -35,6 +35,7 @@ import org.apache.flink.agents.plan.PythonFunction; import org.apache.flink.agents.plan.actions.Action; import org.apache.flink.agents.plan.resourceprovider.PythonResourceProvider; +import org.apache.flink.agents.runtime.JavaMCPResourceDiscovery; import org.apache.flink.agents.runtime.PythonMCPResourceDiscovery; import org.apache.flink.agents.runtime.ResourceCache; import org.apache.flink.agents.runtime.actionstate.ActionState; @@ -271,6 +272,8 @@ public void open() throws Exception { shortTermMemState = getRuntimeContext().getMapState(shortTermMemStateDescriptor); resourceCache = new ResourceCache(agentPlan.getResourceProviders()); + JavaMCPResourceDiscovery.discoverJavaMCPResources( + agentPlan.getResourceProviders(), resourceCache); metricGroup = new FlinkAgentsMetricGroupImpl(getMetricGroup()); builtInMetrics = new BuiltInMetrics(metricGroup, agentPlan); diff --git a/runtime/src/test/java/org/apache/flink/agents/runtime/JavaMCPResourceDiscoveryTest.java b/runtime/src/test/java/org/apache/flink/agents/runtime/JavaMCPResourceDiscoveryTest.java new file mode 100644 index 000000000..1bdd4087f --- /dev/null +++ b/runtime/src/test/java/org/apache/flink/agents/runtime/JavaMCPResourceDiscoveryTest.java @@ -0,0 +1,243 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.flink.agents.runtime; + +import org.apache.flink.agents.api.resource.Resource; +import org.apache.flink.agents.api.resource.ResourceDescriptor; +import org.apache.flink.agents.api.resource.ResourceType; +import org.apache.flink.agents.plan.resourceprovider.JavaResourceProvider; +import org.apache.flink.agents.plan.resourceprovider.JavaSerializableResourceProvider; +import org.apache.flink.agents.plan.resourceprovider.ResourceProvider; +import org.apache.flink.agents.runtime.ResourceCacheTest.TestTool; +import org.junit.jupiter.api.Test; + +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.function.BiFunction; + +import static org.apache.flink.agents.api.resource.ResourceType.MCP_SERVER; +import static org.apache.flink.agents.api.resource.ResourceType.PROMPT; +import static org.apache.flink.agents.api.resource.ResourceType.TOOL; +import static org.assertj.core.api.Assertions.assertThat; + +/** Tests for {@link JavaMCPResourceDiscovery}. */ +public class JavaMCPResourceDiscoveryTest { + + static class FakeMCPTool extends Resource { + private final String name; + + FakeMCPTool(String name) { + this.name = name; + } + + public String getName() { + return name; + } + + @Override + public ResourceType getResourceType() { + return TOOL; + } + } + + static class FakeMCPPrompt extends Resource { + private final String name; + + FakeMCPPrompt(String name) { + this.name = name; + } + + public String getName() { + return name; + } + + @Override + public ResourceType getResourceType() { + return PROMPT; + } + } + + static class FakeMCPServer extends Resource { + private final List tools; + private final List prompts; + private final boolean supportsPrompts; + + FakeMCPServer( + List tools, List prompts, boolean supportsPrompts) { + this.tools = tools; + this.prompts = prompts; + this.supportsPrompts = supportsPrompts; + } + + public List listTools() { + return tools; + } + + public List listPrompts() { + return prompts; + } + + public boolean supportsPrompts() { + return supportsPrompts; + } + + @Override + public ResourceType getResourceType() { + return MCP_SERVER; + } + } + + /** JavaResourceProvider subclass that returns a pre-built server without reflection. */ + static class StubJavaResourceProvider extends JavaResourceProvider { + private final Resource serverToReturn; + + StubJavaResourceProvider(String name, Resource serverToReturn) { + super(name, MCP_SERVER, new ResourceDescriptor("", "FakeServer", new HashMap<>())); + this.serverToReturn = serverToReturn; + } + + @Override + public Resource provide(BiFunction getResource) { + return serverToReturn; + } + } + + private static Map> buildProviders( + String serverName, ResourceProvider provider) { + Map servers = new HashMap<>(); + servers.put(serverName, provider); + Map> resourceProviders = new HashMap<>(); + resourceProviders.put(MCP_SERVER, servers); + return resourceProviders; + } + + private static ResourceCache emptyCache() { + return new ResourceCache(new HashMap<>()); + } + + // --------------------------------------------------------------------------- + // Tests + // --------------------------------------------------------------------------- + + @Test + void testDiscoverToolsAndPromptsFromJavaMCPServer() throws Exception { + FakeMCPServer server = + new FakeMCPServer( + List.of(new FakeMCPTool("add"), new FakeMCPTool("subtract")), + List.of(new FakeMCPPrompt("ask_sum")), + true); + + ResourceCache cache = emptyCache(); + JavaMCPResourceDiscovery.discoverJavaMCPResources( + buildProviders("myServer", new StubJavaResourceProvider("myServer", server)), + cache); + + assertThat(cache.getResource("add", TOOL)).isInstanceOf(FakeMCPTool.class); + assertThat(cache.getResource("subtract", TOOL)).isInstanceOf(FakeMCPTool.class); + assertThat(cache.getResource("ask_sum", PROMPT)).isInstanceOf(FakeMCPPrompt.class); + } + + @Test + void testSkipsPromptDiscoveryWhenNotSupported() throws Exception { + FakeMCPServer server = + new FakeMCPServer( + List.of(new FakeMCPTool("add")), + List.of(new FakeMCPPrompt("should_not_appear")), + false /* supportsPrompts = false */); + + ResourceCache cache = emptyCache(); + JavaMCPResourceDiscovery.discoverJavaMCPResources( + buildProviders("myServer", new StubJavaResourceProvider("myServer", server)), + cache); + + assertThat(cache.getResource("add", TOOL)).isInstanceOf(FakeMCPTool.class); + + // Prompt must not be in the cache + org.assertj.core.api.Assertions.assertThatThrownBy( + () -> cache.getResource("should_not_appear", PROMPT)) + .isInstanceOf(IllegalArgumentException.class); + } + + @Test + void testSkipsNonJavaResourceProviders() throws Exception { + // Use a JavaSerializableResourceProvider (not JavaResourceProvider) — must be ignored + TestTool dummyTool = new TestTool("dummyTool"); + ResourceProvider nonJavaProvider = + JavaSerializableResourceProvider.createResourceProvider("nonJava", TOOL, dummyTool); + + Map servers = new HashMap<>(); + servers.put("nonJava", nonJavaProvider); + Map> resourceProviders = new HashMap<>(); + resourceProviders.put(MCP_SERVER, servers); + + ResourceCache cache = emptyCache(); + // Should complete without errors and without putting anything in the cache + JavaMCPResourceDiscovery.discoverJavaMCPResources(resourceProviders, cache); + + org.assertj.core.api.Assertions.assertThatThrownBy(() -> cache.getResource("nonJava", TOOL)) + .isInstanceOf(IllegalArgumentException.class); + } + + @Test + void testHandlesNoMCPServersRegistered() throws Exception { + // No MCP_SERVER entry at all + Map> resourceProviders = new HashMap<>(); + + ResourceCache cache = emptyCache(); + // Must complete without throwing + JavaMCPResourceDiscovery.discoverJavaMCPResources(resourceProviders, cache); + } + + @Test + void testHandlesEmptyToolAndPromptLists() throws Exception { + FakeMCPServer server = + new FakeMCPServer(List.of(), List.of(), true /* supportsPrompts but no prompts */); + + ResourceCache cache = emptyCache(); + JavaMCPResourceDiscovery.discoverJavaMCPResources( + buildProviders("myServer", new StubJavaResourceProvider("myServer", server)), + cache); + + // Nothing should be in the cache; no exception should be thrown + } + + @Test + void testMixedJavaAndNonJavaProviders() throws Exception { + FakeMCPServer javaServer = + new FakeMCPServer(List.of(new FakeMCPTool("javaToolA")), List.of(), false); + + TestTool dummyTool = new TestTool("dummyTool"); + ResourceProvider nonJavaProvider = + JavaSerializableResourceProvider.createResourceProvider("nonJava", TOOL, dummyTool); + + Map servers = new HashMap<>(); + servers.put("javaServer", new StubJavaResourceProvider("javaServer", javaServer)); + servers.put("nonJavaServer", nonJavaProvider); + + Map> resourceProviders = new HashMap<>(); + resourceProviders.put(MCP_SERVER, servers); + + ResourceCache cache = emptyCache(); + JavaMCPResourceDiscovery.discoverJavaMCPResources(resourceProviders, cache); + + // Only the Java server's tools should be discoverable + assertThat(cache.getResource("javaToolA", TOOL)).isInstanceOf(FakeMCPTool.class); + } +}