[hotfix] Register all resource types added via Agent.addResource (#671)

Author: GreatEugenius

plan/src/main/java/org/apache/flink/agents/plan/AgentPlan.java

@@ -359,6 +359,18 @@ private void extractActionsFromAgent(Agent agent) throws Exception {
         }
     }
 
+    private static ResourceDescriptor requireResourceDescriptor(
+            String name, ResourceType type, Object value) {
+        if (!(value instanceof ResourceDescriptor)) {
+            throw new IllegalStateException(
+                    String.format(
+                            "Resource '%s' of type %s must be a ResourceDescriptor when added via"
+                                    + " Agent.addResource, but got %s",
+                            name, type, value == null ? "null" : value.getClass().getName()));
+        }
+        return (ResourceDescriptor) value;
+    }
+
     private void extractResource(ResourceType type, Method method) throws Exception {
         extractResource(type, method, null);
     }
@@ -561,25 +573,7 @@ private void extractResourceProvidersFromAgent(Agent agent) throws Exception {
 
         for (Map.Entry<ResourceType, Map<String, Object>> entry : agent.getResources().entrySet()) {
             ResourceType type = entry.getKey();
-            if (type == ResourceType.CHAT_MODEL || type == ResourceType.CHAT_MODEL_CONNECTION) {
-                for (Map.Entry<String, Object> kv : entry.getValue().entrySet()) {
-                    ResourceProvider provider;
-                    if (PythonResourceWrapper.class.isAssignableFrom(
-                            Class.forName(
-                                    ((ResourceDescriptor) kv.getValue()).getClazz(),
-                                    true,
-                                    Thread.currentThread().getContextClassLoader()))) {
-                        provider =
-                                new PythonResourceProvider(
-                                        kv.getKey(), type, (ResourceDescriptor) kv.getValue());
-                    } else {
-                        provider =
-                                new JavaResourceProvider(
-                                        kv.getKey(), type, (ResourceDescriptor) kv.getValue());
-                    }
-                    addResourceProvider(provider);
-                }
-            } else if (type == PROMPT) {
+            if (type == PROMPT) {
                 for (Map.Entry<String, Object> kv : entry.getValue().entrySet()) {
                     JavaSerializableResourceProvider provider =
                             JavaSerializableResourceProvider.createResourceProvider(
@@ -593,6 +587,30 @@ private void extractResourceProvidersFromAgent(Agent agent) throws Exception {
                             ((org.apache.flink.agents.api.tools.FunctionTool) kv.getValue())
                                     .getMethod());
                 }
+            } else if (type == MCP_SERVER) {
+                if (!entry.getValue().isEmpty()) {
+                    throw new UnsupportedOperationException(
+                            "Adding an MCP server via Agent.addResource is not supported."
+                                    + " Declare the MCP server with a @MCPServer-annotated static"
+                                    + " method on your Agent class so its tools and prompts can be"
+                                    + " discovered.");
+                }
+            } else {
+                for (Map.Entry<String, Object> kv : entry.getValue().entrySet()) {
+                    ResourceDescriptor descriptor =
+                            requireResourceDescriptor(kv.getKey(), type, kv.getValue());
+                    ResourceProvider provider;
+                    if (PythonResourceWrapper.class.isAssignableFrom(
+                            Class.forName(
+                                    descriptor.getClazz(),
+                                    true,
+                                    Thread.currentThread().getContextClassLoader()))) {
+                        provider = new PythonResourceProvider(kv.getKey(), type, descriptor);
+                    } else {
+                        provider = new JavaResourceProvider(kv.getKey(), type, descriptor);
+                    }
+                    addResourceProvider(provider);
+                }
             }
         }
     }

plan/src/test/java/org/apache/flink/agents/plan/AgentPlanTest.java

@@ -38,6 +38,7 @@
 import org.apache.flink.agents.api.vectorstores.VectorStoreQuery;
 import org.apache.flink.agents.api.vectorstores.VectorStoreQueryResult;
 import org.apache.flink.agents.plan.actions.Action;
+import org.apache.flink.agents.plan.resourceprovider.JavaResourceProvider;
 import org.apache.flink.agents.plan.resourceprovider.JavaSerializableResourceProvider;
 import org.apache.flink.agents.plan.resourceprovider.PythonResourceProvider;
 import org.apache.flink.agents.plan.resourceprovider.ResourceProvider;
@@ -478,4 +479,72 @@ public void testGetResourceFromResourceProvider() throws Exception {
         Resource myToolAgain = agentPlan.getResource("myTool", ResourceType.TOOL);
         assertThat(myTool).isSameAs(myToolAgain);
     }
+
+    @Test
+    public void testAddResourceRegistersEmbeddingModelProvider() throws Exception {
+        Agent agent = new Agent();
+        ResourceDescriptor descriptor =
+                ResourceDescriptor.Builder.newBuilder(TestPythonResource.class.getName())
+                        .addInitialArgument("pythonClazz", "test.module.EmbeddingClazz")
+                        .build();
+        agent.addResource("myEmbedding", ResourceType.EMBEDDING_MODEL, descriptor);
+
+        AgentPlan plan = new AgentPlan(agent);
+
+        Map<String, ResourceProvider> providers =
+                plan.getResourceProviders().get(ResourceType.EMBEDDING_MODEL);
+        assertThat(providers).isNotNull();
+        assertThat(providers).containsKey("myEmbedding");
+
+        ResourceProvider provider = providers.get("myEmbedding");
+        // TestPythonResource implements PythonResourceWrapper, so a Python provider is expected.
+        assertThat(provider).isInstanceOf(PythonResourceProvider.class);
+        assertThat(provider.getName()).isEqualTo("myEmbedding");
+        assertThat(provider.getType()).isEqualTo(ResourceType.EMBEDDING_MODEL);
+    }
+
+    @Test
+    public void testAddResourceRegistersVectorStoreJavaProvider() throws Exception {
+        Agent agent = new Agent();
+        // A non-Python-wrapper class triggers JavaResourceProvider — String is fine for this
+        // structural check; we never call provide() here.
+        ResourceDescriptor descriptor =
+                ResourceDescriptor.Builder.newBuilder(String.class.getName()).build();
+        agent.addResource("myVectorStore", ResourceType.VECTOR_STORE, descriptor);
+
+        AgentPlan plan = new AgentPlan(agent);
+
+        ResourceProvider provider =
+                plan.getResourceProviders().get(ResourceType.VECTOR_STORE).get("myVectorStore");
+        assertThat(provider).isInstanceOf(JavaResourceProvider.class);
+        assertThat(provider.getType()).isEqualTo(ResourceType.VECTOR_STORE);
+    }
+
+    @Test
+    public void testAddResourceRejectsNonDescriptorForUnsupportedType() {
+        Agent agent = new Agent();
+        // SerializableResource for EMBEDDING_MODEL is illegal — only PROMPT / TOOL accept
+        // non-descriptors. The new code path must reject it with a clear message instead of CCE.
+        agent.addResource(
+                "badEmbedding",
+                ResourceType.EMBEDDING_MODEL,
+                new TestSerializableChatModel("badEmbedding"));
+
+        Assertions.assertThrows(IllegalStateException.class, () -> new AgentPlan(agent));
+    }
+
+    @Test
+    public void testAddResourceMCPServerRejectedWithGuidance() {
+        Agent agent = new Agent();
+        ResourceDescriptor descriptor =
+                ResourceDescriptor.Builder.newBuilder("dummy.MCPServer")
+                        .addInitialArgument("endpoint", "http://127.0.0.1:0/mcp")
+                        .build();
+        agent.addResource("addedMcpServer", ResourceType.MCP_SERVER, descriptor);
+
+        UnsupportedOperationException ex =
+                Assertions.assertThrows(
+                        UnsupportedOperationException.class, () -> new AgentPlan(agent));
+        assertThat(ex.getMessage()).contains("@MCPServer");
+    }
 }