This is an automated email from the ASF dual-hosted git repository.
wenjin272 pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/flink-agents.git
The following commit(s) were added to refs/heads/main by this push:
new 220cbd4e [api][plan][runtime][python] Rebuild registered sub-agent
setups from descriptors (#1132)
220cbd4e is described below
commit 220cbd4e550362361cdba8f9bd7ddb5303c9a9b0
Author: yunfengzhou-hub <[email protected]>
AuthorDate: Tue Sep 29 10:07:36 2026 +0800
[api][plan][runtime][python] Rebuild registered sub-agent setups from
descriptors (#1132)
Generated-by: Qoder 1.29.0 (Qwen3.8-Max)
Co-authored-by: Qoder <[email protected]>
---
.../flink/agents/api/subagent/SubagentSetup.java | 86 +++++++++---
.../chat/model/BaseChatModelSetupSubagentTest.java | 25 +++-
.../agents/api/subagent/SubagentRegisterTest.java | 12 +-
.../agents/api/subagent/SubagentSetupTest.java | 91 ++++++++++---
.../agents/api/subagent/TestSubagentSetup.java | 28 ++--
.../org/apache/flink/agents/plan/AgentPlan.java | 9 +-
.../agents/plan/AgentPlanSubagentResourceTest.java | 2 +-
.../plan/AgentPlanSubagentSerializationTest.java | 148 +++++++++++++++++++++
.../plan/actions/ToolCallActionSubagentTest.java | 30 ++++-
python/flink_agents/plan/agent_plan.py | 16 ++-
.../tests/test_agent_plan_subagent_resources.py | 33 +++++
.../runtime/subagent/BaseAsyncSubagentSetup.java | 13 +-
.../subagent/BaseDeferredSubagentSetup.java | 8 ++
.../agents/runtime/subagent/BaseSubagentSetup.java | 10 ++
.../flink/agents/runtime/ResourceCacheTest.java | 11 ++
.../subagent/BaseAsyncSubagentSetupTest.java | 13 +-
.../runtime/subagent/BaseSubagentSetupTest.java | 24 ++++
.../subagent/DeferredSubagentSubmitTest.java | 18 ++-
.../runtime/subagent/MockAsyncSubagentSetup.java | 64 ++++++---
.../subagent/MockDeferredSubagentSetup.java | 16 +++
.../external/ExternalAgentSubagentSetupTest.java | 7 +
.../external/ExternalAsyncSubagentSetup.java | 38 +++++-
.../external/ExternalDeferredSubagentSetup.java | 37 +++++-
23 files changed, 628 insertions(+), 111 deletions(-)
diff --git
a/api/src/main/java/org/apache/flink/agents/api/subagent/SubagentSetup.java
b/api/src/main/java/org/apache/flink/agents/api/subagent/SubagentSetup.java
index 364db659..130ea034 100644
--- a/api/src/main/java/org/apache/flink/agents/api/subagent/SubagentSetup.java
+++ b/api/src/main/java/org/apache/flink/agents/api/subagent/SubagentSetup.java
@@ -19,13 +19,16 @@
package org.apache.flink.agents.api.subagent;
import com.fasterxml.jackson.annotation.JsonIgnore;
-import com.fasterxml.jackson.annotation.JsonProperty;
import org.apache.flink.agents.api.context.RunnerContext;
+import org.apache.flink.agents.api.resource.ResourceContext;
+import org.apache.flink.agents.api.resource.ResourceDescriptor;
import org.apache.flink.agents.api.resource.ResourceType;
import org.apache.flink.agents.api.resource.SerializableResource;
import javax.annotation.Nullable;
+import java.util.Objects;
+
/**
* Caller-facing definition of a sub-agent, registered in the agent plan as an
{@code AGENT}
* resource.
@@ -40,34 +43,77 @@ public abstract class SubagentSetup extends
SerializableResource {
public static final String CALLABLE_NAME_PREFIX = "_subagent_";
/**
- * Tells a caller what this sub-agent is for, so that it can decide
whether to delegate to it.
- * This is routing information for the caller, not an instruction for the
sub-agent itself.
+ * Descriptor argument key carrying the caller-facing {@link
#getDescription()}. It is the
+ * single source of truth shared by the constructor that reads it back and
the subclasses whose
+ * descriptors write it, and the cross-language wire key the Python side
reads.
*/
- @JsonProperty("description")
- private String description;
+ public static final String FIELD_DESCRIPTION = "description";
/**
- * JSON Schema of the arguments this sub-agent accepts, as declared
explicitly. Null when it was
- * not, in which case {@link #getInputSchema()} derives it from {@link
#getInputType()}.
+ * Descriptor argument key carrying the explicitly declared {@link
#getInputSchema()}, absent
+ * when it was not declared. Shared like {@link #FIELD_DESCRIPTION} across
the write and read
+ * ends so they cannot drift apart.
*/
- @JsonProperty("input_schema")
- @Nullable
- private String inputSchema;
+ public static final String FIELD_INPUT_SCHEMA = "input_schema";
- protected SubagentSetup() {
- this("");
- }
+ /**
+ * The descriptor capturing this setup's construction configuration. A
compiled plan carries
+ * this descriptor across the JobManager to TaskManager transfer, and a
remote task rebuilds an
+ * equivalent setup from it through the {@code (ResourceDescriptor,
ResourceContext)}
+ * constructor. Every setup carries one, so a registered sub-agent is
always rebuildable.
+ */
+ private final ResourceDescriptor descriptor;
- protected SubagentSetup(String description) {
- this(description, null);
- }
+ /**
+ * Tells a caller what this sub-agent is for, so that it can decide
whether to delegate to it.
+ * This is routing information for the caller, not an instruction for the
sub-agent itself. It
+ * travels as the descriptor's {@value #FIELD_DESCRIPTION} argument, so a
remote task reads it
+ * back instead of losing it with the live object.
+ */
+ private final String description;
- protected SubagentSetup(String description, @Nullable String inputSchema) {
- this.description = description == null ? "" : description;
- if (inputSchema != null && inputSchema.isBlank()) {
+ /**
+ * JSON Schema of the arguments this sub-agent accepts, as declared
explicitly. Null when it was
+ * not, in which case {@link #getInputSchema()} derives it from {@link
#getInputType()}. It
+ * travels as the descriptor's {@value #FIELD_INPUT_SCHEMA} argument.
+ */
+ @Nullable private final String inputSchema;
+
+ /**
+ * Constructs the setup from the descriptor carrying its configuration.
This is the only
+ * construction path: concrete subclasses expose a public form of it so
the framework can
+ * rebuild them from a descriptor on a remote task. The descriptor must
name this setup's own
+ * concrete type as its clazz, because that name is what the remote
rebuild reflects over; a
+ * mismatch is rejected here rather than surfacing as a wrong-class
rebuild on a far task. The
+ * caller-facing metadata travels as descriptor arguments, so it is read
back here rather than
+ * passed alongside the descriptor.
+ */
+ protected SubagentSetup(ResourceDescriptor descriptor, ResourceContext
resourceContext) {
+ this.descriptor =
+ Objects.requireNonNull(
+ descriptor,
+ "A SubagentSetup must carry a ResourceDescriptor so it
can be rebuilt on a"
+ + " remote task.");
+ if (!getClass().getName().equals(this.descriptor.getClazz())) {
+ throw new IllegalArgumentException(
+ String.format(
+ "A %s must carry a descriptor naming its own type,
but the descriptor"
+ + " names %s; a remote task would rebuild
the wrong class.",
+ getClass().getName(), this.descriptor.getClazz()));
+ }
+ String declaredDescription =
this.descriptor.getArgument(FIELD_DESCRIPTION);
+ this.description = declaredDescription == null ? "" :
declaredDescription;
+ String declaredInputSchema =
this.descriptor.getArgument(FIELD_INPUT_SCHEMA);
+ if (declaredInputSchema != null && declaredInputSchema.isBlank()) {
throw new IllegalArgumentException("Sub-agent input schema must
not be blank.");
}
- this.inputSchema = inputSchema;
+ this.inputSchema = declaredInputSchema;
+ }
+
+ /** The descriptor this setup is rebuilt from on a remote task. */
+ @JsonIgnore
+ public ResourceDescriptor getDescriptor() {
+ return descriptor;
}
@Override
diff --git
a/api/src/test/java/org/apache/flink/agents/api/chat/model/BaseChatModelSetupSubagentTest.java
b/api/src/test/java/org/apache/flink/agents/api/chat/model/BaseChatModelSetupSubagentTest.java
index 54e8e50d..abbc9aab 100644
---
a/api/src/test/java/org/apache/flink/agents/api/chat/model/BaseChatModelSetupSubagentTest.java
+++
b/api/src/test/java/org/apache/flink/agents/api/chat/model/BaseChatModelSetupSubagentTest.java
@@ -138,11 +138,30 @@ class BaseChatModelSetupSubagentTest {
private static final long serialVersionUID = 1L;
StubSubagentSetup(String description) {
- super(description);
+ this(description, null);
}
StubSubagentSetup(String description, String inputSchema) {
- super(description, inputSchema);
+ this(metadataDescriptor(StubSubagentSetup.class, description,
inputSchema), null);
+ }
+
+ StubSubagentSetup(ResourceDescriptor descriptor, ResourceContext
resourceContext) {
+ super(descriptor, resourceContext);
+ }
+
+ /**
+ * Builds the metadata-only descriptor naming {@code concreteClass},
so a subclass that adds
+ * no configuration of its own still carries a descriptor naming its
own concrete type.
+ */
+ static ResourceDescriptor metadataDescriptor(
+ Class<?> concreteClass, String description, String
inputSchema) {
+ ResourceDescriptor.Builder builder =
+
ResourceDescriptor.Builder.newBuilder(concreteClass.getName())
+ .addInitialArgument(FIELD_DESCRIPTION,
description);
+ if (inputSchema != null) {
+ builder.addInitialArgument(FIELD_INPUT_SCHEMA, inputSchema);
+ }
+ return builder.build();
}
@Override
@@ -168,7 +187,7 @@ class BaseChatModelSetupSubagentTest {
private static final long serialVersionUID = 1L;
TypedStubSubagentSetup(String description) {
- super(description);
+ super(metadataDescriptor(TypedStubSubagentSetup.class,
description, null), null);
}
@Override
diff --git
a/api/src/test/java/org/apache/flink/agents/api/subagent/SubagentRegisterTest.java
b/api/src/test/java/org/apache/flink/agents/api/subagent/SubagentRegisterTest.java
index a0dfb7a2..4b4056f3 100644
---
a/api/src/test/java/org/apache/flink/agents/api/subagent/SubagentRegisterTest.java
+++
b/api/src/test/java/org/apache/flink/agents/api/subagent/SubagentRegisterTest.java
@@ -34,7 +34,7 @@ class SubagentRegisterTest {
@Test
void registerSubagentSetupAsResource() {
Agent agent = new Agent();
- TestSubagentSetup setup = new TestSubagentSetup();
+ TestSubagentSetup setup = new TestSubagentSetup(null);
agent.addResource("reviewer", ResourceType.AGENT, setup);
Map<String, Object> agentResources =
agent.getResources().get(ResourceType.AGENT);
@@ -46,17 +46,19 @@ class SubagentRegisterTest {
@Test
void duplicateNameThrows() {
Agent agent = new Agent();
- agent.addResource("reviewer", ResourceType.AGENT, new
TestSubagentSetup());
+ agent.addResource("reviewer", ResourceType.AGENT, new
TestSubagentSetup(null));
assertThrows(
IllegalArgumentException.class,
- () -> agent.addResource("reviewer", ResourceType.AGENT, new
TestSubagentSetup()));
+ () ->
+ agent.addResource(
+ "reviewer", ResourceType.AGENT, new
TestSubagentSetup(null)));
}
@Test
void multipleSubagentsRegistered() {
Agent agent = new Agent();
- TestSubagentSetup reviewer = new TestSubagentSetup();
- TestSubagentSetup coder = new TestSubagentSetup();
+ TestSubagentSetup reviewer = new TestSubagentSetup(null);
+ TestSubagentSetup coder = new TestSubagentSetup(null);
agent.addResource("reviewer", ResourceType.AGENT, reviewer);
agent.addResource("coder", ResourceType.AGENT, coder);
diff --git
a/api/src/test/java/org/apache/flink/agents/api/subagent/SubagentSetupTest.java
b/api/src/test/java/org/apache/flink/agents/api/subagent/SubagentSetupTest.java
index 1c342ba4..525abbaf 100644
---
a/api/src/test/java/org/apache/flink/agents/api/subagent/SubagentSetupTest.java
+++
b/api/src/test/java/org/apache/flink/agents/api/subagent/SubagentSetupTest.java
@@ -22,6 +22,8 @@ import com.fasterxml.jackson.annotation.JsonProperty;
import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import org.apache.flink.agents.api.context.RunnerContext;
+import org.apache.flink.agents.api.resource.ResourceContext;
+import org.apache.flink.agents.api.resource.ResourceDescriptor;
import org.junit.jupiter.api.Test;
import javax.annotation.Nullable;
@@ -33,7 +35,13 @@ import java.util.Map;
import static org.assertj.core.api.Assertions.assertThat;
import static org.assertj.core.api.Assertions.assertThatThrownBy;
-/** Pins the routing metadata {@link SubagentSetup} carries for a caller. */
+/**
+ * Pins the construction contract of {@link SubagentSetup} and the routing
metadata it carries for a
+ * caller. Every setup carries the descriptor a remote task rebuilds it from,
and that descriptor
+ * names the setup's own type, so a registered sub-agent is always rebuildable
into the right class.
+ * The metadata travels as descriptor arguments, the single wire both a remote
task and the Python
+ * side read, so it is pinned there rather than on a serialized object.
+ */
public class SubagentSetupTest {
private static final ObjectMapper MAPPER = new ObjectMapper();
@@ -44,15 +52,35 @@ public class SubagentSetupTest {
private static final long serialVersionUID = 1L;
MetadataOnlySetup() {
- super();
+ this("");
}
MetadataOnlySetup(String description) {
- super(description);
+ this(description, null);
}
MetadataOnlySetup(String description, @Nullable String inputSchema) {
- super(description, inputSchema);
+ this(metadataDescriptor(MetadataOnlySetup.class, description,
inputSchema), null);
+ }
+
+ /** Descriptor-based construction, the path a remote task rebuilds the
setup through. */
+ MetadataOnlySetup(ResourceDescriptor descriptor, ResourceContext
resourceContext) {
+ super(descriptor, resourceContext);
+ }
+
+ /**
+ * Builds the metadata-only descriptor naming {@code concreteClass},
so a subclass that adds
+ * no configuration of its own still carries a descriptor naming its
own concrete type.
+ */
+ static ResourceDescriptor metadataDescriptor(
+ Class<?> concreteClass, String description, @Nullable String
inputSchema) {
+ ResourceDescriptor.Builder builder =
+
ResourceDescriptor.Builder.newBuilder(concreteClass.getName())
+ .addInitialArgument(FIELD_DESCRIPTION,
description);
+ if (inputSchema != null) {
+ builder.addInitialArgument(FIELD_INPUT_SCHEMA, inputSchema);
+ }
+ return builder.build();
}
@Override
@@ -81,7 +109,7 @@ public class SubagentSetupTest {
private final Class<?> resultType;
TypedSetup(Class<?> inputType, Class<?> resultType) {
- super("Reviews a file.");
+ super(metadataDescriptor(TypedSetup.class, "Reviews a file.",
null), null);
this.inputType = inputType;
this.resultType = resultType;
}
@@ -310,32 +338,55 @@ public class SubagentSetupTest {
.hasMessageContaining("is self-referential");
}
- /** The declared types drive behavior, so they must not leak into the
cross-language plan. */
+ /**
+ * The declared types drive behavior, so they are not descriptor state:
they never travel to a
+ * remote task or across to Python, which reads only the descriptor's
arguments.
+ */
@Test
- void theDeclaredTypesStayOutOfThePlanJson() throws Exception {
- String json = MAPPER.writeValueAsString(new TypedSetup(Review.class,
Review.class));
+ void theDeclaredTypesStayOutOfTheDescriptor() {
+ Map<String, Object> arguments =
+ new TypedSetup(Review.class,
Review.class).getDescriptor().getInitialArguments();
-
assertThat(json).doesNotContain("inputType").doesNotContain("resultType");
+ assertThat(arguments).doesNotContainKeys("inputType", "resultType");
}
/**
- * The plan JSON is a cross-language contract: these two keys are what the
Python side reads, so
- * they are pinned literally rather than through the getters.
+ * The descriptor's arguments are the cross-language wire the Python side
reads, so the metadata
+ * keys are pinned literally there rather than through the getters, and
the descriptor names the
+ * concrete type a remote task reflects over to rebuild the setup.
*/
@Test
- void theMetadataSerializesUnderTheCrossLanguageKeys() throws Exception {
+ void theMetadataTravelsUnderTheCrossLanguageDescriptorKeys() {
String customSchema =
"{\"type\":\"object\",\"properties\":{\"path\":{\"type\":\"string\"}}}";
- String json =
- MAPPER.writeValueAsString(new MetadataOnlySetup("Reviews a
file.", customSchema));
+ MetadataOnlySetup setup = new MetadataOnlySetup("Reviews a file.",
customSchema);
+ Map<String, Object> arguments =
setup.getDescriptor().getInitialArguments();
- assertThat(json)
- .contains("\"description\":\"Reviews a file.\"")
- .contains("\"input_schema\"");
- assertThat(json).doesNotContain("inputSchema");
- Map<String, Object> parsed = MAPPER.readValue(json, Map.class);
- assertThat(parsed).containsEntry("input_schema", customSchema);
+ assertThat(arguments).containsEntry("description", "Reviews a file.");
+ assertThat(arguments).containsEntry("input_schema", customSchema);
+
assertThat(setup.getDescriptor().getClazz()).isEqualTo(MetadataOnlySetup.class.getName());
+ }
+
+ @Test
+ void nullDescriptorIsRejectedAtConstruction() {
+ // The descriptor is what a remote task rebuilds the setup from, so
construction requires
+ // one and reports a missing descriptor at the point of the mistake.
+ assertThatThrownBy(() -> new TestSubagentSetup((ResourceDescriptor)
null, null))
+ .isInstanceOf(NullPointerException.class)
+ .hasMessageContaining("must carry a ResourceDescriptor");
+ }
+
+ @Test
+ void mismatchedDescriptorClazzIsRejectedAtConstruction() {
+ // The descriptor names the class a remote task reflects over to
rebuild the setup, so it
+ // must name this setup's own type; a copy-paste or aliasing mistake
is caught here rather
+ // than rebuilding the wrong class on a far task.
+ ResourceDescriptor mismatched =
+
ResourceDescriptor.Builder.newBuilder("com.example.SomeOtherSubagent").build();
+ assertThatThrownBy(() -> new TestSubagentSetup(mismatched, null))
+ .isInstanceOf(IllegalArgumentException.class)
+ .hasMessageContaining("must carry a descriptor naming its own
type");
}
private static List<String> textValues(JsonNode array) {
diff --git
a/api/src/test/java/org/apache/flink/agents/api/subagent/TestSubagentSetup.java
b/api/src/test/java/org/apache/flink/agents/api/subagent/TestSubagentSetup.java
index c9ee099f..5cbb290e 100644
---
a/api/src/test/java/org/apache/flink/agents/api/subagent/TestSubagentSetup.java
+++
b/api/src/test/java/org/apache/flink/agents/api/subagent/TestSubagentSetup.java
@@ -18,8 +18,6 @@
package org.apache.flink.agents.api.subagent;
-import com.fasterxml.jackson.annotation.JsonCreator;
-import com.fasterxml.jackson.annotation.JsonProperty;
import org.apache.flink.agents.api.context.RunnerContext;
import org.apache.flink.agents.api.resource.ResourceContext;
import org.apache.flink.agents.api.resource.ResourceDescriptor;
@@ -35,30 +33,30 @@ public class TestSubagentSetup extends SubagentSetup {
private static final long serialVersionUID = 1L;
+ private static final String FIELD_ENDPOINT = "endpoint";
+ private static final String FIELD_FAIL_ON_CALL = "fail_on_call";
+
@Nullable private final String endpoint;
private final boolean failOnCall;
- public TestSubagentSetup() {
- this(null, false);
- }
-
public TestSubagentSetup(@Nullable String endpoint) {
this(endpoint, false);
}
- @JsonCreator
- public TestSubagentSetup(
- @JsonProperty("endpoint") @Nullable String endpoint,
- @JsonProperty("failOnCall") boolean failOnCall) {
- this.endpoint = endpoint;
- this.failOnCall = failOnCall;
+ public TestSubagentSetup(@Nullable String endpoint, boolean failOnCall) {
+ this(
+
ResourceDescriptor.Builder.newBuilder(TestSubagentSetup.class.getName())
+ .addInitialArgument(FIELD_ENDPOINT, endpoint)
+ .addInitialArgument(FIELD_FAIL_ON_CALL, failOnCall)
+ .build(),
+ null);
}
/** Descriptor-based construction, as used by YAML-declared {@code
subagents:} entries. */
public TestSubagentSetup(ResourceDescriptor descriptor, ResourceContext
resourceContext) {
- this(
- (String) descriptor.getArgument("endpoint"),
- Boolean.TRUE.equals(descriptor.getArgument("fail_on_call")));
+ super(descriptor, resourceContext);
+ this.endpoint = descriptor.getArgument(FIELD_ENDPOINT);
+ this.failOnCall =
Boolean.TRUE.equals(descriptor.getArgument(FIELD_FAIL_ON_CALL));
}
@Nullable
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 89fb9c2c..90fe54a9 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
@@ -599,9 +599,14 @@ public class AgentPlan implements Serializable {
String name = kv.getKey();
Object value = kv.getValue();
if (value instanceof SubagentSetup) {
+ // A live setup registered programmatically travels as
its descriptor, so a
+ // remote task rebuilds it through the
(ResourceDescriptor, ResourceContext)
+ // constructor.
addResourceProvider(
-
JavaSerializableResourceProvider.createResourceProvider(
- name, ResourceType.AGENT,
(SubagentSetup) value));
+ createDescriptorResourceProvider(
+ name,
+ ResourceType.AGENT,
+ ((SubagentSetup)
value).getDescriptor()));
} else if (value instanceof ResourceDescriptor) {
// Declared via YAML: the descriptor names a
SubagentSetup subclass that is
// instantiated when the resource is first resolved.
diff --git
a/plan/src/test/java/org/apache/flink/agents/plan/AgentPlanSubagentResourceTest.java
b/plan/src/test/java/org/apache/flink/agents/plan/AgentPlanSubagentResourceTest.java
index d7be8ad4..e943c408 100644
---
a/plan/src/test/java/org/apache/flink/agents/plan/AgentPlanSubagentResourceTest.java
+++
b/plan/src/test/java/org/apache/flink/agents/plan/AgentPlanSubagentResourceTest.java
@@ -42,7 +42,7 @@ public class AgentPlanSubagentResourceTest {
@Test
void subagentSetupInstanceCompilesIntoAgentProvider() throws Exception {
Agent agent = new Agent();
- agent.addResource("reviewer", ResourceType.AGENT, new
TestSubagentSetup());
+ agent.addResource("reviewer", ResourceType.AGENT, new
TestSubagentSetup(null));
AgentPlan plan = new AgentPlan(agent);
diff --git
a/plan/src/test/java/org/apache/flink/agents/plan/AgentPlanSubagentSerializationTest.java
b/plan/src/test/java/org/apache/flink/agents/plan/AgentPlanSubagentSerializationTest.java
new file mode 100644
index 00000000..6665a59f
--- /dev/null
+++
b/plan/src/test/java/org/apache/flink/agents/plan/AgentPlanSubagentSerializationTest.java
@@ -0,0 +1,148 @@
+/*
+ * 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.plan;
+
+import org.apache.flink.agents.api.agents.Agent;
+import org.apache.flink.agents.api.context.RunnerContext;
+import org.apache.flink.agents.api.resource.Resource;
+import org.apache.flink.agents.api.resource.ResourceContext;
+import org.apache.flink.agents.api.resource.ResourceDescriptor;
+import org.apache.flink.agents.api.resource.ResourceType;
+import org.apache.flink.agents.api.subagent.SubagentFuture;
+import org.apache.flink.agents.api.subagent.SubagentSetup;
+import org.apache.flink.agents.plan.resourceprovider.ResourceProvider;
+import org.junit.jupiter.api.Test;
+
+import java.io.ByteArrayInputStream;
+import java.io.ByteArrayOutputStream;
+import java.io.IOException;
+import java.io.ObjectInputStream;
+import java.io.ObjectOutputStream;
+import java.util.Map;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
+/**
+ * An external sub-agent registered programmatically as a live {@link
SubagentSetup} instance must
+ * survive the plan's serialization round-trip — the transfer that carries a
compiled plan from the
+ * JobManager to a remote TaskManager — and resolve there to an equivalent
setup carrying the same
+ * configuration. In-process the plan resolves to the live object, so only the
round-trip exercises
+ * the descriptor-based rebuild on the far side.
+ */
+public class AgentPlanSubagentSerializationTest {
+
+ @Test
+ void liveExternalSubagentIsRebuiltAfterPlanSerializationRoundTrip() throws
Exception {
+ Agent agent = new Agent();
+ agent.addResource(
+ "reviewer",
+ ResourceType.AGENT,
+ new ExternalReviewSubagent("http://review.internal:8080";, 3));
+
+ AgentPlan plan = new AgentPlan(agent);
+
+ // The JobManager -> TaskManager transfer: AgentPlan crosses the
process boundary through
+ // its Serializable hooks, so the live setup object does not travel
with it and must be
+ // rebuilt from what does.
+ AgentPlan restored = serializeRoundTrip(plan);
+
+ Map<String, ResourceProvider> agentProviders =
+ restored.getResourceProviders().get(ResourceType.AGENT);
+ assertThat(agentProviders).containsKey("reviewer");
+
+ Resource resolved = agentProviders.get("reviewer").provide(null);
+
+ assertThat(resolved).isInstanceOf(ExternalReviewSubagent.class);
+ ExternalReviewSubagent reviewer = (ExternalReviewSubagent) resolved;
+
assertThat(reviewer.getEndpoint()).isEqualTo("http://review.internal:8080";);
+ assertThat(reviewer.getMaxRetries()).isEqualTo(3);
+ assertThat(resolved.getResourceType()).isEqualTo(ResourceType.AGENT);
+ }
+
+ /**
+ * Round-trips the plan through Java serialization, the transfer that
carries it to a remote
+ * TaskManager. {@link AgentPlan} exposes no {@code serialize()} method:
its public contract is
+ * {@code Serializable}, whose private {@code writeObject}/{@code
readObject} hooks Java invokes
+ * automatically, so writing the plan to an {@link ObjectOutputStream} is
that entry point.
+ */
+ private static AgentPlan serializeRoundTrip(AgentPlan plan)
+ throws IOException, ClassNotFoundException {
+ ByteArrayOutputStream bytes = new ByteArrayOutputStream();
+ try (ObjectOutputStream out = new ObjectOutputStream(bytes)) {
+ out.writeObject(plan);
+ }
+ try (ObjectInputStream in =
+ new ObjectInputStream(new
ByteArrayInputStream(bytes.toByteArray()))) {
+ return (AgentPlan) in.readObject();
+ }
+ }
+
+ /**
+ * A user-defined external sub-agent in the documented shape: construction
configuration held in
+ * final fields and set through a parameterized constructor, with no
Jackson creator.
+ */
+ public static class ExternalReviewSubagent extends SubagentSetup {
+
+ private static final String FIELD_ENDPOINT = "endpoint";
+ private static final String FIELD_MAX_RETRIES = "max_retries";
+
+ private final String endpoint;
+ private final int maxRetries;
+
+ public ExternalReviewSubagent(String endpoint, int maxRetries) {
+ this(
+
ResourceDescriptor.Builder.newBuilder(ExternalReviewSubagent.class.getName())
+ .addInitialArgument(FIELD_ENDPOINT, endpoint)
+ .addInitialArgument(FIELD_MAX_RETRIES, maxRetries)
+ .build(),
+ null);
+ }
+
+ public ExternalReviewSubagent(
+ ResourceDescriptor descriptor, ResourceContext
resourceContext) {
+ super(descriptor, resourceContext);
+ this.endpoint = descriptor.getArgument(FIELD_ENDPOINT);
+ this.maxRetries = descriptor.getArgument(FIELD_MAX_RETRIES, 0);
+ }
+
+ public String getEndpoint() {
+ return endpoint;
+ }
+
+ public int getMaxRetries() {
+ return maxRetries;
+ }
+
+ @Override
+ public SubagentFuture submit(RunnerContext ctx, Object prompt) {
+ throw new UnsupportedOperationException("serialization is under
test, not invocation");
+ }
+
+ @Override
+ public SubagentFuture submit(RunnerContext ctx, Object prompt, String
sessionId) {
+ throw new UnsupportedOperationException("serialization is under
test, not invocation");
+ }
+
+ @Override
+ public SubagentFuture submit(
+ RunnerContext ctx, Object prompt, String sessionId, String
callId) {
+ throw new UnsupportedOperationException("serialization is under
test, not invocation");
+ }
+ }
+}
diff --git
a/plan/src/test/java/org/apache/flink/agents/plan/actions/ToolCallActionSubagentTest.java
b/plan/src/test/java/org/apache/flink/agents/plan/actions/ToolCallActionSubagentTest.java
index 16858efa..d4ec380a 100644
---
a/plan/src/test/java/org/apache/flink/agents/plan/actions/ToolCallActionSubagentTest.java
+++
b/plan/src/test/java/org/apache/flink/agents/plan/actions/ToolCallActionSubagentTest.java
@@ -30,6 +30,7 @@ import org.apache.flink.agents.api.event.ToolResponseEvent;
import org.apache.flink.agents.api.memory.BaseLongTermMemory;
import org.apache.flink.agents.api.metrics.FlinkAgentsMetricGroup;
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.api.subagent.SubagentFuture;
import org.apache.flink.agents.api.subagent.SubagentFutures;
@@ -497,6 +498,17 @@ class ToolCallActionSubagentTest {
Map.of("name", secondTool, "arguments",
Map.of("q", "y")))));
}
+ /**
+ * Builds a metadata-only descriptor naming {@code concreteClass}, so each
in-process double
+ * carries a descriptor naming its own type as the {@link SubagentSetup}
constructor requires.
+ */
+ private static ResourceDescriptor subagentDescriptor(
+ Class<?> concreteClass, String description) {
+ return ResourceDescriptor.Builder.newBuilder(concreteClass.getName())
+ .addInitialArgument(SubagentSetup.FIELD_DESCRIPTION,
description)
+ .build();
+ }
+
/** Captures every prompt it is handed and resolves to a preset outcome. */
private static class RecordingSubagentSetup extends SubagentSetup {
private final SubagentResult outcome;
@@ -504,7 +516,11 @@ class ToolCallActionSubagentTest {
private Exception submitFailure;
RecordingSubagentSetup(SubagentResult outcome) {
- super("Reviews a diff.");
+ this(subagentDescriptor(RecordingSubagentSetup.class, "Reviews a
diff."), outcome);
+ }
+
+ RecordingSubagentSetup(ResourceDescriptor descriptor, SubagentResult
outcome) {
+ super(descriptor, null);
this.outcome = outcome;
}
@@ -534,7 +550,9 @@ class ToolCallActionSubagentTest {
/** Declares a result type, so its result is read through it. */
private static class TypedRecordingSubagentSetup extends
RecordingSubagentSetup {
TypedRecordingSubagentSetup(SubagentResult outcome) {
- super(outcome);
+ super(
+ subagentDescriptor(TypedRecordingSubagentSetup.class,
"Reviews a diff."),
+ outcome);
}
@Override
@@ -602,7 +620,7 @@ class ToolCallActionSubagentTest {
private final List<String> ops;
OrderRecordingSubagentSetup(String label, List<String> ops) {
- super("Orders a diff.");
+ super(subagentDescriptor(OrderRecordingSubagentSetup.class,
"Orders a diff."), null);
this.label = label;
this.ops = ops;
}
@@ -664,7 +682,9 @@ class ToolCallActionSubagentTest {
private final List<String> ops;
InterruptingSubagentSetup(String label, List<String> ops) {
- super("Interrupts on await.");
+ super(
+ subagentDescriptor(InterruptingSubagentSetup.class,
"Interrupts on await."),
+ null);
this.label = label;
this.ops = ops;
}
@@ -723,7 +743,7 @@ class ToolCallActionSubagentTest {
/** Its await overflows the stack, to drive the {@link StackOverflowError}
absorption path. */
private static class OverflowingSubagentSetup extends SubagentSetup {
OverflowingSubagentSetup() {
- super("Overflows on await.");
+ super(subagentDescriptor(OverflowingSubagentSetup.class,
"Overflows on await."), null);
}
@Override
diff --git a/python/flink_agents/plan/agent_plan.py
b/python/flink_agents/plan/agent_plan.py
index 41d80c8a..a8034b16 100644
--- a/python/flink_agents/plan/agent_plan.py
+++ b/python/flink_agents/plan/agent_plan.py
@@ -409,23 +409,25 @@ def _get_resource_providers(
for name, value in agent.resources[ResourceType.AGENT].items():
if isinstance(value, SubagentSetup):
- resource_providers.append(
- PythonSerializableResourceProvider.from_resource(
- name=name, resource=value
- )
+ # A live setup travels as its descriptor, so a remote task rebuilds
+ # it from the class name and configuration.
+ descriptor = ResourceDescriptor(
+
clazz=f"{value.__class__.__module__}.{value.__class__.__name__}",
+ arguments=value.model_dump(),
)
elif isinstance(value, ResourceDescriptor):
# Declared via YAML: the descriptor names a SubagentSetup subclass
# that is instantiated when the resource is first resolved.
- resource_providers.append(
- PythonResourceProvider.get(name=name, descriptor=value)
- )
+ descriptor = value
else:
msg = (
f"AGENT resource '{name}' must be a SubagentSetup or a "
f"ResourceDescriptor, but got {type(value).__name__}."
)
raise TypeError(msg)
+ resource_providers.append(
+ PythonResourceProvider.get(name=name, descriptor=descriptor)
+ )
for resource_type in [
ResourceType.CHAT_MODEL,
diff --git
a/python/flink_agents/plan/tests/test_agent_plan_subagent_resources.py
b/python/flink_agents/plan/tests/test_agent_plan_subagent_resources.py
index 818e109c..6e255352 100644
--- a/python/flink_agents/plan/tests/test_agent_plan_subagent_resources.py
+++ b/python/flink_agents/plan/tests/test_agent_plan_subagent_resources.py
@@ -24,6 +24,7 @@ from flink_agents.api.subagent import SubagentSetup
from flink_agents.api.tests.subagent_test_utils import TestSubagentSetup
from flink_agents.plan.agent_plan import AgentPlan
from flink_agents.plan.configuration import AgentConfiguration
+from flink_agents.plan.resource_provider import PythonResourceProvider
def test_subagent_setup_compiles_into_agent_provider() -> None:
@@ -70,3 +71,35 @@ def test_non_setup_agent_resource_is_rejected() -> None:
with pytest.raises(TypeError, match="must be a SubagentSetup"):
AgentPlan.from_agent(agent, AgentConfiguration())
+
+
+def test_live_external_subagent_is_rebuilt_after_plan_json_round_trip() ->
None:
+ """A live setup crosses the plan JSON boundary as its descriptor and is
+ rebuilt with its configuration on the other side.
+ """
+ agent = Agent()
+ agent.add_resource(
+ "reviewer",
+ ResourceType.AGENT,
+ TestSubagentSetup(
+ endpoint_url="http://review.internal:8080";, fail_on_call=True
+ ),
+ )
+
+ plan = AgentPlan.from_agent(agent, AgentConfiguration())
+
+ # The transfer that carries a compiled plan between processes: the live
+ # setup does not travel with it and must be rebuilt from what does.
+ restored = AgentPlan.model_validate_json(
+ plan.model_dump_json(serialize_as_any=True)
+ )
+
+ provider = restored.resource_providers[ResourceType.AGENT]["reviewer"]
+ assert isinstance(provider, PythonResourceProvider)
+ assert provider.descriptor.clazz is TestSubagentSetup
+
+ resolved = provider.provide(resource_context=None,
config=AgentConfiguration())
+ assert isinstance(resolved, TestSubagentSetup)
+ assert resolved.endpoint_url == "http://review.internal:8080";
+ assert resolved.fail_on_call is True
+ assert resolved.resource_type() == ResourceType.AGENT
diff --git
a/runtime/src/main/java/org/apache/flink/agents/runtime/subagent/BaseAsyncSubagentSetup.java
b/runtime/src/main/java/org/apache/flink/agents/runtime/subagent/BaseAsyncSubagentSetup.java
index 4bd66dd8..a4eab648 100644
---
a/runtime/src/main/java/org/apache/flink/agents/runtime/subagent/BaseAsyncSubagentSetup.java
+++
b/runtime/src/main/java/org/apache/flink/agents/runtime/subagent/BaseAsyncSubagentSetup.java
@@ -36,6 +36,14 @@ import java.util.concurrent.Callable;
*/
public abstract class BaseAsyncSubagentSetup extends BaseSubagentSetup {
+ /**
+ * Descriptor argument key for the status-poll interval, and the single
source of truth for it:
+ * the descriptor-based constructor below reads this key, while subclasses
that capture a poll
+ * interval into their descriptor write it. Sharing the constant keeps the
read and write ends
+ * from drifting apart.
+ */
+ protected static final String FIELD_STATUS_POLL_INTERVAL_MILLIS =
"status_poll_interval_millis";
+
/**
* Delay between status probes while waiting for the run to reach a
terminal state. Defaults to
* {@code 500}. The descriptor-based constructor reads the optional {@code
@@ -43,8 +51,6 @@ public abstract class BaseAsyncSubagentSetup extends
BaseSubagentSetup {
*/
protected long statusPollIntervalMillis = 500;
- protected BaseAsyncSubagentSetup() {}
-
/**
* Descriptor-based construction, as used by YAML-declared {@code
subagents:} entries: reads the
* optional {@code status_poll_interval_millis} argument, falling back to
the default of {@code
@@ -52,7 +58,8 @@ public abstract class BaseAsyncSubagentSetup extends
BaseSubagentSetup {
*/
protected BaseAsyncSubagentSetup(
ResourceDescriptor descriptor, ResourceContext resourceContext) {
- Number statusPollInterval =
descriptor.getArgument("status_poll_interval_millis");
+ super(descriptor, resourceContext);
+ Number statusPollInterval =
descriptor.getArgument(FIELD_STATUS_POLL_INTERVAL_MILLIS);
if (statusPollInterval != null) {
this.statusPollIntervalMillis = statusPollInterval.longValue();
}
diff --git
a/runtime/src/main/java/org/apache/flink/agents/runtime/subagent/BaseDeferredSubagentSetup.java
b/runtime/src/main/java/org/apache/flink/agents/runtime/subagent/BaseDeferredSubagentSetup.java
index 666898eb..56b7cedf 100644
---
a/runtime/src/main/java/org/apache/flink/agents/runtime/subagent/BaseDeferredSubagentSetup.java
+++
b/runtime/src/main/java/org/apache/flink/agents/runtime/subagent/BaseDeferredSubagentSetup.java
@@ -20,6 +20,8 @@ package org.apache.flink.agents.runtime.subagent;
import org.apache.flink.agents.api.context.DurableCallable;
import org.apache.flink.agents.api.context.RunnerContext;
+import org.apache.flink.agents.api.resource.ResourceContext;
+import org.apache.flink.agents.api.resource.ResourceDescriptor;
import org.apache.flink.agents.api.subagent.SubagentFuture;
import org.apache.flink.agents.api.subagent.SubagentResult;
@@ -31,6 +33,12 @@ import org.apache.flink.agents.api.subagent.SubagentResult;
*/
public abstract class BaseDeferredSubagentSetup extends BaseSubagentSetup {
+ /** Descriptor-based construction, forwarding to the sub-agent base. */
+ protected BaseDeferredSubagentSetup(
+ ResourceDescriptor descriptor, ResourceContext resourceContext) {
+ super(descriptor, resourceContext);
+ }
+
/** Registers the invocation and returns its deferred handle. */
@Override
public SubagentFuture submit(RunnerContext ctx, Object prompt, String
sessionId, String callId)
diff --git
a/runtime/src/main/java/org/apache/flink/agents/runtime/subagent/BaseSubagentSetup.java
b/runtime/src/main/java/org/apache/flink/agents/runtime/subagent/BaseSubagentSetup.java
index a3b07c3c..9dba697c 100644
---
a/runtime/src/main/java/org/apache/flink/agents/runtime/subagent/BaseSubagentSetup.java
+++
b/runtime/src/main/java/org/apache/flink/agents/runtime/subagent/BaseSubagentSetup.java
@@ -19,6 +19,8 @@
package org.apache.flink.agents.runtime.subagent;
import org.apache.flink.agents.api.context.RunnerContext;
+import org.apache.flink.agents.api.resource.ResourceContext;
+import org.apache.flink.agents.api.resource.ResourceDescriptor;
import org.apache.flink.agents.api.subagent.SubagentFuture;
import org.apache.flink.agents.api.subagent.SubagentSetup;
import org.apache.flink.agents.runtime.lifecycle.TaskLifecycleListener;
@@ -42,6 +44,14 @@ public abstract class BaseSubagentSetup extends
SubagentSetup implements TaskLif
/** The task whose execution is currently issuing calls. */
@Nullable private ActionTask currentTask;
+ /**
+ * Descriptor-based construction for a sub-agent rebuilt on a remote task
from the configuration
+ * its descriptor carries.
+ */
+ protected BaseSubagentSetup(ResourceDescriptor descriptor, ResourceContext
resourceContext) {
+ super(descriptor, resourceContext);
+ }
+
@Override
public void onActionPrepared(ActionTask task) {
currentTask = task;
diff --git
a/runtime/src/test/java/org/apache/flink/agents/runtime/ResourceCacheTest.java
b/runtime/src/test/java/org/apache/flink/agents/runtime/ResourceCacheTest.java
index c39498d2..d82cc3ec 100644
---
a/runtime/src/test/java/org/apache/flink/agents/runtime/ResourceCacheTest.java
+++
b/runtime/src/test/java/org/apache/flink/agents/runtime/ResourceCacheTest.java
@@ -528,6 +528,17 @@ public class ResourceCacheTest {
/** Test Java sub-agent setup, registered as an AGENT resource. */
public static class TestAgentSetup extends BaseSubagentSetup {
+
+ public TestAgentSetup() {
+ this(
+
ResourceDescriptor.Builder.newBuilder(TestAgentSetup.class.getName()).build(),
+ null);
+ }
+
+ public TestAgentSetup(ResourceDescriptor descriptor, ResourceContext
resourceContext) {
+ super(descriptor, resourceContext);
+ }
+
@Override
public SubagentFuture submit(
RunnerContext ctx, Object prompt, String sessionId, String
callId) {
diff --git
a/runtime/src/test/java/org/apache/flink/agents/runtime/subagent/BaseAsyncSubagentSetupTest.java
b/runtime/src/test/java/org/apache/flink/agents/runtime/subagent/BaseAsyncSubagentSetupTest.java
index 3bc24138..411f9e88 100644
---
a/runtime/src/test/java/org/apache/flink/agents/runtime/subagent/BaseAsyncSubagentSetupTest.java
+++
b/runtime/src/test/java/org/apache/flink/agents/runtime/subagent/BaseAsyncSubagentSetupTest.java
@@ -38,6 +38,7 @@ import org.apache.flink.api.java.functions.KeySelector;
import org.apache.flink.api.java.typeutils.TypeExtractor;
import org.apache.flink.streaming.runtime.streamrecord.StreamRecord;
import org.apache.flink.streaming.util.KeyedOneInputStreamOperatorTestHarness;
+import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;
import java.util.List;
@@ -57,6 +58,11 @@ public class BaseAsyncSubagentSetupTest {
private static final String RESOURCE_NAME = "ext-agent";
+ @BeforeEach
+ void resetCounters() {
+ MockAsyncSubagentSetup.reset();
+ }
+
//
------------------------------------------------------------------------------------------
// The pub: one durable POST, issued immediately
//
------------------------------------------------------------------------------------------
@@ -181,7 +187,9 @@ public class BaseAsyncSubagentSetupTest {
assertThat(original.statusQueryCount()).isEqualTo(3);
// Replay: the run has already reached a terminal state, so the same
await takes a
- // shorter path — fewer probes — to the same result.
+ // shorter path — fewer probes — to the same result. The endpoint
counters are shared
+ // with the mock service, so restart them to measure the replay on its
own.
+ MockAsyncSubagentSetup.reset();
MockAsyncSubagentSetup replay = new MockAsyncSubagentSetup(2, false);
replay.seedRun("sid-1", "call-1", "done:ping", null, 0);
SubagentResult after = replay.awaitResultForTest("sid-1",
"call-1").call();
@@ -467,6 +475,9 @@ public class BaseAsyncSubagentSetupTest {
@Test
void descriptorArgumentSetsTheStatusPollInterval() {
+ // The key is spelled out rather than referencing the base's protected
constant on purpose:
+ // this asserts the wire contract (the exact argument name a
descriptor-built setup reads),
+ // so it stays a literal to catch any change to that name.
DescriptorAsyncSetup setup =
new DescriptorAsyncSetup(
new ResourceDescriptor(
diff --git
a/runtime/src/test/java/org/apache/flink/agents/runtime/subagent/BaseSubagentSetupTest.java
b/runtime/src/test/java/org/apache/flink/agents/runtime/subagent/BaseSubagentSetupTest.java
index 8cca1acf..326a4e10 100644
---
a/runtime/src/test/java/org/apache/flink/agents/runtime/subagent/BaseSubagentSetupTest.java
+++
b/runtime/src/test/java/org/apache/flink/agents/runtime/subagent/BaseSubagentSetupTest.java
@@ -23,6 +23,8 @@ import org.apache.flink.agents.api.InputEvent;
import org.apache.flink.agents.api.OutputEvent;
import org.apache.flink.agents.api.agents.Agent;
import org.apache.flink.agents.api.context.RunnerContext;
+import org.apache.flink.agents.api.resource.ResourceContext;
+import org.apache.flink.agents.api.resource.ResourceDescriptor;
import org.apache.flink.agents.api.resource.ResourceType;
import org.apache.flink.agents.api.subagent.SubagentFuture;
import org.apache.flink.agents.api.subagent.SubagentFutures;
@@ -68,6 +70,18 @@ public class BaseSubagentSetupTest {
/** Captures every assigned {@code (sessionId, callId)} pair like a
collecting setup. */
public static class AllocatingCaptureSetup extends BaseSubagentSetup {
+ public AllocatingCaptureSetup() {
+ this(
+
ResourceDescriptor.Builder.newBuilder(AllocatingCaptureSetup.class.getName())
+ .build(),
+ null);
+ }
+
+ public AllocatingCaptureSetup(
+ ResourceDescriptor descriptor, ResourceContext
resourceContext) {
+ super(descriptor, resourceContext);
+ }
+
private static final List<String[]> CAPTURES =
Collections.synchronizedList(new ArrayList<>());
@@ -231,6 +245,16 @@ public class BaseSubagentSetupTest {
/** Capture setup exposing the current task's registry for direct
tracking. */
public static class RegistryExposingSetup extends AllocatingCaptureSetup {
+ // Names its own type in the descriptor: SubagentSetup's construction
check requires the
+ // descriptor clazz to match the concrete type, which the inherited
convenience constructor
+ // would otherwise name as the superclass.
+ public RegistryExposingSetup() {
+ super(
+
ResourceDescriptor.Builder.newBuilder(RegistryExposingSetup.class.getName())
+ .build(),
+ null);
+ }
+
PendingSubagentCallRegistry exposedRegistry() {
return currentTaskRegistry();
}
diff --git
a/runtime/src/test/java/org/apache/flink/agents/runtime/subagent/DeferredSubagentSubmitTest.java
b/runtime/src/test/java/org/apache/flink/agents/runtime/subagent/DeferredSubagentSubmitTest.java
index cfc46c50..b274e90e 100644
---
a/runtime/src/test/java/org/apache/flink/agents/runtime/subagent/DeferredSubagentSubmitTest.java
+++
b/runtime/src/test/java/org/apache/flink/agents/runtime/subagent/DeferredSubagentSubmitTest.java
@@ -24,6 +24,8 @@ import org.apache.flink.agents.api.OutputEvent;
import org.apache.flink.agents.api.agents.Agent;
import org.apache.flink.agents.api.context.DurableCallable;
import org.apache.flink.agents.api.context.RunnerContext;
+import org.apache.flink.agents.api.resource.ResourceContext;
+import org.apache.flink.agents.api.resource.ResourceDescriptor;
import org.apache.flink.agents.api.resource.ResourceType;
import org.apache.flink.agents.api.subagent.SubagentFuture;
import org.apache.flink.agents.api.subagent.SubagentResult;
@@ -269,7 +271,21 @@ public class DeferredSubagentSubmitTest {
}
/** Deferred setup whose prepared callable throws a system-level failure
instead of folding. */
- private static final class ThrowingDeferredSubagentSetup extends
BaseDeferredSubagentSetup {
+ public static final class ThrowingDeferredSubagentSetup extends
BaseDeferredSubagentSetup {
+
+ public ThrowingDeferredSubagentSetup() {
+ this(
+ ResourceDescriptor.Builder.newBuilder(
+
ThrowingDeferredSubagentSetup.class.getName())
+ .build(),
+ null);
+ }
+
+ public ThrowingDeferredSubagentSetup(
+ ResourceDescriptor descriptor, ResourceContext
resourceContext) {
+ super(descriptor, resourceContext);
+ }
+
@Override
protected DurableCallable<SubagentResult> prepare(
RunnerContext ctx, Object prompt, String sessionId, String
callId) {
diff --git
a/runtime/src/test/java/org/apache/flink/agents/runtime/subagent/MockAsyncSubagentSetup.java
b/runtime/src/test/java/org/apache/flink/agents/runtime/subagent/MockAsyncSubagentSetup.java
index f5d9818e..d52b7866 100644
---
a/runtime/src/test/java/org/apache/flink/agents/runtime/subagent/MockAsyncSubagentSetup.java
+++
b/runtime/src/test/java/org/apache/flink/agents/runtime/subagent/MockAsyncSubagentSetup.java
@@ -19,12 +19,15 @@
package org.apache.flink.agents.runtime.subagent;
import org.apache.flink.agents.api.context.DurableCallable;
+import org.apache.flink.agents.api.resource.ResourceContext;
+import org.apache.flink.agents.api.resource.ResourceDescriptor;
import org.apache.flink.agents.api.subagent.SubagentResult;
import javax.annotation.Nullable;
import java.util.HashMap;
import java.util.Map;
+import java.util.concurrent.atomic.AtomicInteger;
/**
* Example integration of {@link BaseAsyncSubagentSetup}: an in-memory
asynchronous agent service.
@@ -35,6 +38,9 @@ import java.util.Map;
*/
public class MockAsyncSubagentSetup extends BaseAsyncSubagentSetup {
+ private static final String FIELD_QUERIES_UNTIL_COMPLETE =
"queries_until_complete";
+ private static final String FIELD_FAIL_ON_POST = "fail_on_post";
+
/** One recorded remote run, keyed by {@code sessionId#callId}. */
private static final class Run {
private final Object result;
@@ -52,10 +58,22 @@ public class MockAsyncSubagentSetup extends
BaseAsyncSubagentSetup {
private final int queriesUntilComplete;
private final boolean failOnPost;
- private int postCount;
- private int statusQueryCount;
- private int fetchCount;
- private int cancelCount;
+ // Counters are static because the operator materializes a rebuilt
instance from the
+ // descriptor, so the registered setup a test holds is not the one that
runs; call reset()
+ // before each independent scenario. They are atomic because the framework
bumps them on
+ // async pool threads, where batched calls run concurrently.
+ private static final AtomicInteger POST_COUNT = new AtomicInteger();
+ private static final AtomicInteger STATUS_QUERY_COUNT = new
AtomicInteger();
+ private static final AtomicInteger FETCH_COUNT = new AtomicInteger();
+ private static final AtomicInteger CANCEL_COUNT = new AtomicInteger();
+
+ /** Clears every endpoint counter. Call before each independent scenario.
*/
+ public static void reset() {
+ POST_COUNT.set(0);
+ STATUS_QUERY_COUNT.set(0);
+ FETCH_COUNT.set(0);
+ CANCEL_COUNT.set(0);
+ }
public MockAsyncSubagentSetup() {
this(2, false);
@@ -66,16 +84,28 @@ public class MockAsyncSubagentSetup extends
BaseAsyncSubagentSetup {
* terminal; {@code failOnPost} makes every submission fail.
*/
public MockAsyncSubagentSetup(int queriesUntilComplete, boolean
failOnPost) {
- this.queriesUntilComplete = queriesUntilComplete;
- this.failOnPost = failOnPost;
- // Runs turn terminal after a fixed number of probes rather than after
elapsed time, so
- // probing without a delay keeps the counts identical and the tests
fast.
- this.statusPollIntervalMillis = 0;
+ this(
+
ResourceDescriptor.Builder.newBuilder(MockAsyncSubagentSetup.class.getName())
+ .addInitialArgument(FIELD_QUERIES_UNTIL_COMPLETE,
queriesUntilComplete)
+ .addInitialArgument(FIELD_FAIL_ON_POST, failOnPost)
+ // Runs turn terminal after a fixed number of probes
rather than after
+ // elapsed time, so probing without a delay keeps the
counts identical and
+ // the tests fast.
+ .addInitialArgument(FIELD_STATUS_POLL_INTERVAL_MILLIS,
0L)
+ .build(),
+ null);
+ }
+
+ /** Descriptor-based construction, reading the config the convenience
constructor captured. */
+ public MockAsyncSubagentSetup(ResourceDescriptor descriptor,
ResourceContext resourceContext) {
+ super(descriptor, resourceContext);
+ this.queriesUntilComplete =
descriptor.getArgument(FIELD_QUERIES_UNTIL_COMPLETE, 2);
+ this.failOnPost = descriptor.getArgument(FIELD_FAIL_ON_POST, false);
}
@Override
protected void callSubmitRequest(String sessionId, String callId, Object
prompt) {
- postCount++;
+ POST_COUNT.incrementAndGet();
if (failOnPost) {
throw new IllegalStateException("post failed");
}
@@ -84,7 +114,7 @@ public class MockAsyncSubagentSetup extends
BaseAsyncSubagentSetup {
@Override
protected RunStatus callQueryStatus(String sessionId, String callId) {
- statusQueryCount++;
+ STATUS_QUERY_COUNT.incrementAndGet();
Run run = runs.get(sessionId + "#" + callId);
if (run == null) {
return RunStatus.notStarted();
@@ -98,7 +128,7 @@ public class MockAsyncSubagentSetup extends
BaseAsyncSubagentSetup {
@Override
protected SubagentResult callFetchResult(String sessionId, String callId) {
- fetchCount++;
+ FETCH_COUNT.incrementAndGet();
Run run = runs.get(sessionId + "#" + callId);
if (run == null) {
return SubagentResult.error("no run on record");
@@ -108,7 +138,7 @@ public class MockAsyncSubagentSetup extends
BaseAsyncSubagentSetup {
@Override
protected void callCancelRequest(String sessionId, String callId) {
- cancelCount++;
+ CANCEL_COUNT.incrementAndGet();
}
/** Test hook: injects a run that already exists remotely, exercising
reconciler reuse. */
@@ -123,22 +153,22 @@ public class MockAsyncSubagentSetup extends
BaseAsyncSubagentSetup {
/** Number of times the POST endpoint has been hit. */
public int postCount() {
- return postCount;
+ return POST_COUNT.get();
}
/** Number of times the status endpoint has been probed. */
public int statusQueryCount() {
- return statusQueryCount;
+ return STATUS_QUERY_COUNT.get();
}
/** Number of times the result endpoint has been fetched. */
public int fetchCount() {
- return fetchCount;
+ return FETCH_COUNT.get();
}
/** Number of times the cancel hook has been invoked. */
public int cancelCount() {
- return cancelCount;
+ return CANCEL_COUNT.get();
}
/** Exposes the pub durable call for unit-style POST and reconciler
assertions. */
diff --git
a/runtime/src/test/java/org/apache/flink/agents/runtime/subagent/MockDeferredSubagentSetup.java
b/runtime/src/test/java/org/apache/flink/agents/runtime/subagent/MockDeferredSubagentSetup.java
index 0a8b69a7..087f95c2 100644
---
a/runtime/src/test/java/org/apache/flink/agents/runtime/subagent/MockDeferredSubagentSetup.java
+++
b/runtime/src/test/java/org/apache/flink/agents/runtime/subagent/MockDeferredSubagentSetup.java
@@ -20,6 +20,8 @@ package org.apache.flink.agents.runtime.subagent;
import org.apache.flink.agents.api.context.DurableCallable;
import org.apache.flink.agents.api.context.RunnerContext;
+import org.apache.flink.agents.api.resource.ResourceContext;
+import org.apache.flink.agents.api.resource.ResourceDescriptor;
import org.apache.flink.agents.api.subagent.SubagentResult;
import java.util.ArrayList;
@@ -38,6 +40,18 @@ import java.util.concurrent.atomic.AtomicInteger;
*/
public class MockDeferredSubagentSetup extends BaseDeferredSubagentSetup {
+ public MockDeferredSubagentSetup() {
+ this(
+
ResourceDescriptor.Builder.newBuilder(MockDeferredSubagentSetup.class.getName())
+ .build(),
+ null);
+ }
+
+ public MockDeferredSubagentSetup(
+ ResourceDescriptor descriptor, ResourceContext resourceContext) {
+ super(descriptor, resourceContext);
+ }
+
/** One {@code (sessionId, callId)} assignment captured at prepare time. */
public static final class Capture {
public final String sessionId;
@@ -62,6 +76,8 @@ public class MockDeferredSubagentSetup extends
BaseDeferredSubagentSetup {
}
}
+ // Static observation state (see class doc), reset() each scenario.
EXECUTION_COUNT is atomic
+ // because call bodies run on async pool threads, where batched calls run
concurrently.
private static final List<Capture> CAPTURES =
Collections.synchronizedList(new ArrayList<>());
private static final AtomicInteger EXECUTION_COUNT = new AtomicInteger();
diff --git
a/runtime/src/test/java/org/apache/flink/agents/runtime/subagent/external/ExternalAgentSubagentSetupTest.java
b/runtime/src/test/java/org/apache/flink/agents/runtime/subagent/external/ExternalAgentSubagentSetupTest.java
index e42b9a2f..4e652384 100644
---
a/runtime/src/test/java/org/apache/flink/agents/runtime/subagent/external/ExternalAgentSubagentSetupTest.java
+++
b/runtime/src/test/java/org/apache/flink/agents/runtime/subagent/external/ExternalAgentSubagentSetupTest.java
@@ -39,6 +39,7 @@ import
org.apache.flink.streaming.util.KeyedOneInputStreamOperatorTestHarness;
import org.junit.jupiter.api.AfterAll;
import org.junit.jupiter.api.Assumptions;
import org.junit.jupiter.api.BeforeAll;
+import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
@@ -135,6 +136,12 @@ public class ExternalAgentSubagentSetupTest {
}
}
+ @BeforeEach
+ void resetObservationState() {
+ ExternalAsyncSubagentSetup.reset();
+ ExternalDeferredSubagentSetup.reset();
+ }
+
//
------------------------------------------------------------------------------------------
// Mode 1: async pub/sub through BaseAsyncSubagentSetup
//
------------------------------------------------------------------------------------------
diff --git
a/runtime/src/test/java/org/apache/flink/agents/runtime/subagent/external/ExternalAsyncSubagentSetup.java
b/runtime/src/test/java/org/apache/flink/agents/runtime/subagent/external/ExternalAsyncSubagentSetup.java
index 29c25d7d..3991bdf5 100644
---
a/runtime/src/test/java/org/apache/flink/agents/runtime/subagent/external/ExternalAsyncSubagentSetup.java
+++
b/runtime/src/test/java/org/apache/flink/agents/runtime/subagent/external/ExternalAsyncSubagentSetup.java
@@ -18,6 +18,8 @@
package org.apache.flink.agents.runtime.subagent.external;
+import org.apache.flink.agents.api.resource.ResourceContext;
+import org.apache.flink.agents.api.resource.ResourceDescriptor;
import org.apache.flink.agents.api.subagent.SubagentResult;
import org.apache.flink.agents.runtime.subagent.BaseAsyncSubagentSetup;
import org.slf4j.Logger;
@@ -47,14 +49,38 @@ public class ExternalAsyncSubagentSetup extends
BaseAsyncSubagentSetup {
private static final Logger LOG =
LoggerFactory.getLogger(ExternalAsyncSubagentSetup.class);
+ private static final String FIELD_BASE_URL = "base_url";
+
private final String baseUrl;
- private final List<Long> probeTimestamps =
Collections.synchronizedList(new ArrayList<>());
+
+ // Static because the operator materializes a rebuilt instance from the
descriptor; the
+ // registered setup a test holds is not the one that probes. Call reset()
before each scenario.
+ // Synchronized because probes are recorded on async pool threads, where
batched calls run
+ // concurrently.
+ private static final List<Long> PROBE_TIMESTAMPS =
+ Collections.synchronizedList(new ArrayList<>());
@Nullable private transient ExternalAgentClient client;
+ /** Clears the recorded probe timestamps. Call before each independent
scenario. */
+ public static void reset() {
+ PROBE_TIMESTAMPS.clear();
+ }
+
public ExternalAsyncSubagentSetup(String baseUrl, long
probeIntervalMillis) {
- this.baseUrl = baseUrl;
- this.statusPollIntervalMillis = probeIntervalMillis;
+ this(
+
ResourceDescriptor.Builder.newBuilder(ExternalAsyncSubagentSetup.class.getName())
+ .addInitialArgument(FIELD_BASE_URL, baseUrl)
+ .addInitialArgument(FIELD_STATUS_POLL_INTERVAL_MILLIS,
probeIntervalMillis)
+ .build(),
+ null);
+ }
+
+ /** Descriptor-based construction, reading the config the convenience
constructor captured. */
+ public ExternalAsyncSubagentSetup(
+ ResourceDescriptor descriptor, ResourceContext resourceContext) {
+ super(descriptor, resourceContext);
+ this.baseUrl = descriptor.getArgument(FIELD_BASE_URL);
}
private ExternalAgentClient client() {
@@ -73,7 +99,7 @@ public class ExternalAsyncSubagentSetup extends
BaseAsyncSubagentSetup {
@Override
protected RunStatus callQueryStatus(String sessionId, String callId) {
- probeTimestamps.add(System.currentTimeMillis());
+ PROBE_TIMESTAMPS.add(System.currentTimeMillis());
ExternalAgentClient.TaskStatus probe;
try {
probe = client().status(taskId(sessionId, callId));
@@ -116,8 +142,8 @@ public class ExternalAsyncSubagentSetup extends
BaseAsyncSubagentSetup {
/** Timestamps of every status probe, for asserting the probe pacing in
tests. */
public List<Long> probeTimestamps() {
- synchronized (probeTimestamps) {
- return new ArrayList<>(probeTimestamps);
+ synchronized (PROBE_TIMESTAMPS) {
+ return new ArrayList<>(PROBE_TIMESTAMPS);
}
}
}
diff --git
a/runtime/src/test/java/org/apache/flink/agents/runtime/subagent/external/ExternalDeferredSubagentSetup.java
b/runtime/src/test/java/org/apache/flink/agents/runtime/subagent/external/ExternalDeferredSubagentSetup.java
index 439b86a6..3d6037a9 100644
---
a/runtime/src/test/java/org/apache/flink/agents/runtime/subagent/external/ExternalDeferredSubagentSetup.java
+++
b/runtime/src/test/java/org/apache/flink/agents/runtime/subagent/external/ExternalDeferredSubagentSetup.java
@@ -20,6 +20,8 @@ package org.apache.flink.agents.runtime.subagent.external;
import org.apache.flink.agents.api.context.DurableCallable;
import org.apache.flink.agents.api.context.RunnerContext;
+import org.apache.flink.agents.api.resource.ResourceContext;
+import org.apache.flink.agents.api.resource.ResourceDescriptor;
import org.apache.flink.agents.api.subagent.SubagentResult;
import org.apache.flink.agents.runtime.subagent.BaseDeferredSubagentSetup;
import org.slf4j.Logger;
@@ -43,15 +45,40 @@ public class ExternalDeferredSubagentSetup extends
BaseDeferredSubagentSetup {
private static final Logger LOG =
LoggerFactory.getLogger(ExternalDeferredSubagentSetup.class);
+ private static final String FIELD_BASE_URL = "base_url";
+ private static final String FIELD_POLL_INTERVAL_MILLIS =
"poll_interval_millis";
+
private final String baseUrl;
private final long pollIntervalMillis;
- private final AtomicInteger pollCount = new AtomicInteger();
+
+ // Static because the operator materializes a rebuilt instance from the
descriptor; the
+ // registered setup a test holds is not the one that polls. Call reset()
before each scenario.
+ // Atomic because polls are counted on async pool threads, where batched
calls run
+ // concurrently.
+ private static final AtomicInteger POLL_COUNT = new AtomicInteger();
@Nullable private transient ExternalAgentClient client;
+ /** Clears the poll counter. Call before each independent scenario. */
+ public static void reset() {
+ POLL_COUNT.set(0);
+ }
+
public ExternalDeferredSubagentSetup(String baseUrl, long
pollIntervalMillis) {
- this.baseUrl = baseUrl;
- this.pollIntervalMillis = pollIntervalMillis;
+ this(
+
ResourceDescriptor.Builder.newBuilder(ExternalDeferredSubagentSetup.class.getName())
+ .addInitialArgument(FIELD_BASE_URL, baseUrl)
+ .addInitialArgument(FIELD_POLL_INTERVAL_MILLIS,
pollIntervalMillis)
+ .build(),
+ null);
+ }
+
+ /** Descriptor-based construction, reading the config the convenience
constructor captured. */
+ public ExternalDeferredSubagentSetup(
+ ResourceDescriptor descriptor, ResourceContext resourceContext) {
+ super(descriptor, resourceContext);
+ this.baseUrl = descriptor.getArgument(FIELD_BASE_URL);
+ this.pollIntervalMillis =
descriptor.getArgument(FIELD_POLL_INTERVAL_MILLIS, 0L);
}
private ExternalAgentClient client() {
@@ -90,7 +117,7 @@ public class ExternalDeferredSubagentSetup extends
BaseDeferredSubagentSetup {
callId,
taskId);
while (true) {
- int check = pollCount.incrementAndGet();
+ int check = POLL_COUNT.incrementAndGet();
ExternalAgentClient.TaskStatus probe =
client().status(taskId);
LOG.info(
"[deferred] {}#{} check #{} -> {}",
@@ -125,6 +152,6 @@ public class ExternalDeferredSubagentSetup extends
BaseDeferredSubagentSetup {
/** Total number of status polls across all invocations, for test pacing
assertions. */
public int pollCount() {
- return pollCount.get();
+ return POLL_COUNT.get();
}
}
