You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@flink.apache.org by we...@apache.org on 2026/09/09 08:49:29 UTC
(flink-agents) branch main updated: [api][plan][runtime] Introduce AGENT resource type and sub-agent invocation API (#938)
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 223b20280 [api][plan][runtime] Introduce AGENT resource type and sub-agent invocation API (#938)
223b20280 is described below
commit 223b2028074b7647a06b3fd87da61a3afa85ded6
Author: Gen Luo <86...@users.noreply.github.com>
AuthorDate: Wed Sep 9 16:49:22 2026 +0800
[api][plan][runtime] Introduce AGENT resource type and sub-agent invocation API (#938)
* [runtime] Consolidate per-task execution contexts into a single transferable contexts record
* [python] Add optional durable id keying for durable execution
* [runtime] Record the input sequence number on action tasks
* [runtime] Broadcast record/action lifecycle events to registered listeners
* [runtime] Materialize resources of a type through their owning runtime
* [api] Add sub-agent resource definitions
* [plan] Register sub-agent setups as AGENT resources in the planner
* [runtime][python] Add sub-agent setup base with deterministic id assignment
* [runtime][python] Add deferred execution mode for external sub-agent setups
* [runtime][python] Add async external sub-agent base running in durable pub/sub mode
* [runtime][e2e] Add sub-agent integration tests
---
api/pom.xml | 17 +
.../flink/agents/api/resource/ResourceType.java | 3 +-
.../flink/agents/api/subagent/SubagentFuture.java | 57 ++
.../flink/agents/api/subagent/SubagentFutures.java | 44 ++
.../flink/agents/api/subagent/SubagentResult.java | 119 +++++
.../flink/agents/api/subagent/SubagentSetup.java | 58 +++
.../apache/flink/agents/api/yaml/YamlLoader.java | 2 +
.../flink/agents/api/yaml/spec/AgentSpec.java | 9 +-
.../agents/api/yaml/spec/YamlAgentsDocument.java | 9 +-
.../agents/api/subagent/SubagentRegisterTest.java | 68 +++
.../agents/api/subagent/SubagentResultTest.java | 75 +++
.../agents/api/subagent/TestSubagentSetup.java | 91 ++++
.../flink-agents-dev/assets/yaml-contracts.yaml | 2 +-
.../flink-agents-dev/assets/yaml-schema.json | 14 +
docs/yaml-schema.json | 14 +
plan/pom.xml | 7 +
.../org/apache/flink/agents/plan/AgentPlan.java | 25 +
.../plan/resourceprovider/ResourceProvider.java | 9 +
.../agents/plan/AgentPlanSubagentResourceTest.java | 87 ++++
python/flink_agents/api/resource.py | 3 +-
python/flink_agents/api/runner_context.py | 14 +
python/flink_agents/api/subagent.py | 186 +++++++
.../flink_agents/api/tests/subagent_test_utils.py | 49 ++
python/flink_agents/api/tests/test_subagent.py | 60 +++
python/flink_agents/api/yaml/loader.py | 1 +
python/flink_agents/api/yaml/specs.py | 2 +
.../subagent_external_integration_agent.py | 148 ++++++
.../subagent_external_integration_test.py | 96 ++++
.../subagent_integration_agent.py | 87 ++++
.../subagent_integration_test.py | 89 ++++
python/flink_agents/plan/agent_plan.py | 21 +
python/flink_agents/plan/resource_provider.py | 11 +
.../tests/test_agent_plan_subagent_resources.py | 72 +++
python/flink_agents/runtime/async_subagent.py | 347 +++++++++++++
python/flink_agents/runtime/base_subagent.py | 294 +++++++++++
python/flink_agents/runtime/deferred_subagent.py | 180 +++++++
python/flink_agents/runtime/durable_execution.py | 37 +-
.../flink_agents/runtime/flink_runner_context.py | 147 ++++++
python/flink_agents/runtime/resource_cache.py | 7 +
python/flink_agents/runtime/subagent_handles.py | 144 ++++++
.../runtime/task_lifecycle_listener.py | 96 ++++
.../runtime/tests/test_async_subagent.py | 571 +++++++++++++++++++++
.../runtime/tests/test_base_subagent.py | 305 +++++++++++
.../runtime/tests/test_deferred_subagent.py | 452 ++++++++++++++++
.../runtime/tests/test_durable_execution.py | 21 +
.../runtime/tests/test_eager_materialize.py | 129 +++++
.../runtime/tests/test_task_lifecycle_bridge.py | 177 +++++++
.../apache/flink/agents/runtime/ResourceCache.java | 68 +++
.../agents/runtime/context/RunnerContextImpl.java | 108 ++--
.../lifecycle/ComponentExecutionListener.java | 59 +++
.../lifecycle/PythonTaskLifecycleListener.java | 82 +++
.../runtime/lifecycle/TaskLifecycleListener.java | 130 +++++
.../runtime/operator/ActionExecutionOperator.java | 221 ++++++--
.../flink/agents/runtime/operator/ActionTask.java | 38 +-
.../runtime/operator/ActionTaskContextManager.java | 189 ++++---
.../agents/runtime/operator/JavaActionTask.java | 12 +-
.../runtime/operator/PythonBridgeManager.java | 30 +-
.../runtime/python/operator/PythonActionTask.java | 22 +-
.../python/operator/PythonGeneratorActionTask.java | 19 +-
.../python/resource/PythonRuntimeResource.java | 57 ++
.../runtime/python/utils/PythonActionExecutor.java | 110 ++++
.../runtime/subagent/AsyncSubagentFuture.java | 125 +++++
.../runtime/subagent/BaseAsyncSubagentSetup.java | 287 +++++++++++
.../subagent/BaseDeferredSubagentSetup.java | 66 +++
.../agents/runtime/subagent/BaseSubagentSetup.java | 153 ++++++
.../runtime/subagent/CompletedSubagentFuture.java | 49 ++
.../runtime/subagent/DeferredSubagentFuture.java | 129 +++++
.../subagent/PendingSubagentCallRegistry.java | 67 +++
.../runtime/subagent/SubagentFutureGroup.java | 93 ++++
.../runtime/subagent/SubagentIdAllocator.java | 143 ++++++
.../trace/EventLogComponentExecutionListener.java | 86 ++++
.../trace/EventLogTaskLifecycleListener.java | 64 +++
.../flink/agents/runtime/ResourceCacheTest.java | 86 ++++
.../RunnerContextImplExecutionReporterTest.java | 262 ++++++----
.../RunnerContextPendingEventsContractTest.java | 86 ++++
.../context/TestMemoryObservationFlush.java | 17 +-
.../TaskLifecycleListenerNotificationTest.java | 532 +++++++++++++++++++
.../operator/ActionTaskContextManagerTest.java | 185 ++++---
.../agents/runtime/operator/ActionTaskTest.java | 16 +-
.../flink/agents/runtime/operator/TestActions.java | 4 +-
.../python/resource/PythonRuntimeResourceTest.java | 55 ++
.../subagent/BaseAsyncSubagentSetupTest.java | 521 +++++++++++++++++++
.../runtime/subagent/BaseSubagentSetupTest.java | 292 +++++++++++
.../subagent/DeferredSubagentSubmitTest.java | 322 ++++++++++++
.../runtime/subagent/MockAsyncSubagentSetup.java | 154 ++++++
.../subagent/MockDeferredSubagentSetup.java | 108 ++++
.../runtime/subagent/SubagentIdAllocatorTest.java | 251 +++++++++
.../subagent/external/ExternalAgentClient.java | 289 +++++++++++
.../external/ExternalAgentStubService.java | 249 +++++++++
.../external/ExternalAgentSubagentSetupTest.java | 480 +++++++++++++++++
.../external/ExternalAsyncSubagentSetup.java | 123 +++++
.../external/ExternalDeferredSubagentSetup.java | 130 +++++
.../EventLogComponentExecutionListenerTest.java | 232 +++++++++
.../trace/EventLogTaskLifecycleListenerTest.java | 150 ++++++
94 files changed, 10991 insertions(+), 416 deletions(-)
diff --git a/api/pom.xml b/api/pom.xml
index 170740a54..02bcb7035 100644
--- a/api/pom.xml
+++ b/api/pom.xml
@@ -68,4 +68,21 @@ under the License.
</dependency>
</dependencies>
+ <build>
+ <plugins>
+ <plugin>
+ <groupId>org.apache.maven.plugins</groupId>
+ <artifactId>maven-jar-plugin</artifactId>
+ <version>3.4.2</version>
+ <executions>
+ <execution>
+ <goals>
+ <goal>test-jar</goal>
+ </goals>
+ </execution>
+ </executions>
+ </plugin>
+ </plugins>
+ </build>
+
</project>
\ No newline at end of file
diff --git a/api/src/main/java/org/apache/flink/agents/api/resource/ResourceType.java b/api/src/main/java/org/apache/flink/agents/api/resource/ResourceType.java
index dc8b866b0..9ea0d1256 100644
--- a/api/src/main/java/org/apache/flink/agents/api/resource/ResourceType.java
+++ b/api/src/main/java/org/apache/flink/agents/api/resource/ResourceType.java
@@ -33,7 +33,8 @@ public enum ResourceType {
TOOL("tool"),
MCP_SERVER("mcp_server"),
SKILLS("skills"),
- MODEL_ROUTER("model_router");
+ MODEL_ROUTER("model_router"),
+ AGENT("agent");
private final String value;
diff --git a/api/src/main/java/org/apache/flink/agents/api/subagent/SubagentFuture.java b/api/src/main/java/org/apache/flink/agents/api/subagent/SubagentFuture.java
new file mode 100644
index 000000000..38e2eb8cd
--- /dev/null
+++ b/api/src/main/java/org/apache/flink/agents/api/subagent/SubagentFuture.java
@@ -0,0 +1,57 @@
+/*
+ * 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.api.subagent;
+
+/**
+ * Handle for one sub-agent invocation, identified by the {@code (sessionId, callId)} pair that keys
+ * the invocation.
+ */
+public abstract class SubagentFuture {
+
+ private final String sessionId;
+ private final String callId;
+
+ protected SubagentFuture(String sessionId, String callId) {
+ this.sessionId = sessionId;
+ this.callId = callId;
+ }
+
+ public String getSessionId() {
+ return sessionId;
+ }
+
+ public String getCallId() {
+ return callId;
+ }
+
+ /** Whether the invocation has reached a terminal state. */
+ public abstract boolean isDone();
+
+ /**
+ * Resolves the invocation, waiting until it reaches a terminal state. Failures converge into a
+ * failed {@link SubagentResult} rather than a separately reported exceptional completion.
+ */
+ public abstract SubagentResult await() throws Exception;
+
+ /** Requests cancellation of the invocation. */
+ public void cancel() {}
+
+ /** Groups this handle with others to be resolved together through {@link SubagentFutures}. */
+ public abstract SubagentFutures combine(SubagentFuture... others);
+}
diff --git a/api/src/main/java/org/apache/flink/agents/api/subagent/SubagentFutures.java b/api/src/main/java/org/apache/flink/agents/api/subagent/SubagentFutures.java
new file mode 100644
index 000000000..9c3202de2
--- /dev/null
+++ b/api/src/main/java/org/apache/flink/agents/api/subagent/SubagentFutures.java
@@ -0,0 +1,44 @@
+/*
+ * 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.api.subagent;
+
+import java.util.List;
+
+/**
+ * A group of sub-agent handles to be resolved together. A group is not itself an invocation and
+ * carries no {@code (sessionId, callId)} identity.
+ */
+public abstract class SubagentFutures {
+
+ /** Whether every handle in the group has reached a terminal state. */
+ public abstract boolean isDone();
+
+ /**
+ * Waits for every handle in the group and returns their outcomes in the order the handles were
+ * added. Like {@link SubagentFuture#await()}, failures surface through failed {@link
+ * SubagentResult}s.
+ */
+ public abstract List<SubagentResult> awaitAll() throws Exception;
+
+ /** Requests cancellation of every handle in the group. */
+ public void cancel() {}
+
+ /** Adds more handles to the group. */
+ public abstract SubagentFutures combine(SubagentFuture... others);
+}
diff --git a/api/src/main/java/org/apache/flink/agents/api/subagent/SubagentResult.java b/api/src/main/java/org/apache/flink/agents/api/subagent/SubagentResult.java
new file mode 100644
index 000000000..b9322eebf
--- /dev/null
+++ b/api/src/main/java/org/apache/flink/agents/api/subagent/SubagentResult.java
@@ -0,0 +1,119 @@
+/*
+ * 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.api.subagent;
+
+import com.fasterxml.jackson.annotation.JsonCreator;
+import com.fasterxml.jackson.annotation.JsonIgnore;
+import com.fasterxml.jackson.annotation.JsonProperty;
+import com.fasterxml.jackson.databind.ObjectMapper;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.io.Serializable;
+
+/**
+ * Outcome of a sub-agent call issued through {@link SubagentSetup}. A successful outcome carries a
+ * JSON-serializable payload, and a failed one carries a serializable error message.
+ *
+ * <p>Implementations capture their internal failures into a result through {@link #error} instead
+ * of throwing, so callers inspect {@link #isSuccess()} rather than catching. Because the failure is
+ * carried as a message rather than a live exception, the whole result can be persisted through
+ * durable execution and survive a failover.
+ */
+public class SubagentResult implements Serializable {
+
+ private static final Logger LOG = LoggerFactory.getLogger(SubagentResult.class);
+
+ private static final long serialVersionUID = 1L;
+
+ private static final ObjectMapper OBJECT_MAPPER = new ObjectMapper();
+
+ private final boolean success;
+ private final Object result;
+ private final String errorMessage;
+
+ @JsonCreator
+ public SubagentResult(
+ @JsonProperty("success") boolean success,
+ @JsonProperty("result") Object result,
+ @JsonProperty("errorMessage") String errorMessage) {
+ this.success = success;
+ this.result = result;
+ this.errorMessage = errorMessage;
+ }
+
+ /** Creates a successful result carrying the given value. */
+ public static SubagentResult ok(Object result) {
+ return new SubagentResult(true, result, null);
+ }
+
+ /**
+ * Creates a failed result carrying the exception's type and message. The full stack trace is
+ * logged here rather than persisted, keeping the durable payload bounded.
+ */
+ public static SubagentResult error(Exception exception) {
+ if (exception == null) {
+ return new SubagentResult(false, null, null);
+ }
+ LOG.warn("Sub-agent call failed; persisting the exception summary.", exception);
+ return new SubagentResult(false, null, summaryOf(exception));
+ }
+
+ /** Creates a failed result carrying the given message. */
+ public static SubagentResult error(String errorMessage) {
+ return new SubagentResult(false, null, errorMessage);
+ }
+
+ private static String summaryOf(Exception exception) {
+ return exception.getClass().getName() + ": " + exception.getMessage();
+ }
+
+ public boolean isSuccess() {
+ return success;
+ }
+
+ public Object getResult() {
+ return result;
+ }
+
+ /**
+ * Returns the payload converted to {@code resultClass}.
+ *
+ * <p>Durable recovery re-binds the persisted payload through a plain {@link ObjectMapper}
+ * without polymorphic typing, so after a failover replay {@link #getResult()} hands back a
+ * {@code LinkedHashMap} instead of the caller's type. This accessor converts the payload to the
+ * expected class uniformly on both the first execution and a replay.
+ */
+ public <T> T getResult(Class<T> resultClass) {
+ return OBJECT_MAPPER.convertValue(result, resultClass);
+ }
+
+ public String getErrorMessage() {
+ return errorMessage;
+ }
+
+ /**
+ * Reconstructs an exception carrying the stored summary as its message, or null if this result
+ * is successful.
+ */
+ @JsonIgnore
+ public Exception getException() {
+ return success ? null : new RuntimeException(errorMessage);
+ }
+}
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
new file mode 100644
index 000000000..6357e1e67
--- /dev/null
+++ b/api/src/main/java/org/apache/flink/agents/api/subagent/SubagentSetup.java
@@ -0,0 +1,58 @@
+/*
+ * 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.api.subagent;
+
+import com.fasterxml.jackson.annotation.JsonIgnore;
+import org.apache.flink.agents.api.context.RunnerContext;
+import org.apache.flink.agents.api.resource.ResourceType;
+import org.apache.flink.agents.api.resource.SerializableResource;
+
+/**
+ * Caller-facing definition of a sub-agent, registered in the agent plan as an {@code AGENT}
+ * resource.
+ */
+public abstract class SubagentSetup extends SerializableResource {
+
+ @Override
+ @JsonIgnore
+ public ResourceType getResourceType() {
+ return ResourceType.AGENT;
+ }
+
+ /**
+ * Issues a new invocation with an implementation-assigned identity. This is the preferred form.
+ */
+ public abstract SubagentFuture submit(RunnerContext ctx, Object prompt) throws Exception;
+
+ /**
+ * Issues an invocation that continues the conversation of an earlier invocation. Pass the
+ * {@code sessionId} of the earlier invocation to continue it. The session id is available on
+ * the handle returned by that invocation. Whether a conversation can be continued across
+ * actions is up to the concrete implementation.
+ */
+ public abstract SubagentFuture submit(RunnerContext ctx, Object prompt, String sessionId)
+ throws Exception;
+
+ /**
+ * Issues an invocation under the given {@code (sessionId, callId)} identity. This form is
+ * reserved for implementation use.
+ */
+ public abstract SubagentFuture submit(
+ RunnerContext ctx, Object prompt, String sessionId, String callId) throws Exception;
+}
diff --git a/api/src/main/java/org/apache/flink/agents/api/yaml/YamlLoader.java b/api/src/main/java/org/apache/flink/agents/api/yaml/YamlLoader.java
index 7b27e2711..d5c3d30fb 100644
--- a/api/src/main/java/org/apache/flink/agents/api/yaml/YamlLoader.java
+++ b/api/src/main/java/org/apache/flink/agents/api/yaml/YamlLoader.java
@@ -292,6 +292,7 @@ public final class YamlLoader {
addSharedDescriptors(
sharedResources, ResourceType.VECTOR_STORE, doc.getVectorStores(), path);
addSharedDescriptors(sharedResources, ResourceType.MCP_SERVER, doc.getMcpServers(), path);
+ addSharedDescriptors(sharedResources, ResourceType.AGENT, doc.getSubagents(), path);
for (ToolSpec t : doc.getTools()) {
if (sharedResources.get(ResourceType.TOOL).put(t.getName(), buildTool(t)) != null) {
@@ -411,6 +412,7 @@ public final class YamlLoader {
addAgentDescriptors(agent, ResourceType.EMBEDDING_MODEL, spec.getEmbeddingModelSetups());
addAgentDescriptors(agent, ResourceType.VECTOR_STORE, spec.getVectorStores());
addAgentDescriptors(agent, ResourceType.MCP_SERVER, spec.getMcpServers());
+ addAgentDescriptors(agent, ResourceType.AGENT, spec.getSubagents());
for (ToolSpec t : spec.getTools()) {
agent.addResource(t.getName(), ResourceType.TOOL, buildTool(t));
diff --git a/api/src/main/java/org/apache/flink/agents/api/yaml/spec/AgentSpec.java b/api/src/main/java/org/apache/flink/agents/api/yaml/spec/AgentSpec.java
index 0e37a8aad..1826a61a0 100644
--- a/api/src/main/java/org/apache/flink/agents/api/yaml/spec/AgentSpec.java
+++ b/api/src/main/java/org/apache/flink/agents/api/yaml/spec/AgentSpec.java
@@ -40,6 +40,7 @@ public final class AgentSpec {
private final List<DescriptorSpec> embeddingModelSetups;
private final List<DescriptorSpec> vectorStores;
private final List<DescriptorSpec> mcpServers;
+ private final List<DescriptorSpec> subagents;
@JsonCreator
public AgentSpec(
@@ -55,7 +56,8 @@ public final class AgentSpec {
List<DescriptorSpec> embeddingModelConnections,
@JsonProperty("embedding_model_setups") List<DescriptorSpec> embeddingModelSetups,
@JsonProperty("vector_stores") List<DescriptorSpec> vectorStores,
- @JsonProperty("mcp_servers") List<DescriptorSpec> mcpServers) {
+ @JsonProperty("mcp_servers") List<DescriptorSpec> mcpServers,
+ @JsonProperty("subagents") List<DescriptorSpec> subagents) {
this.name = name;
this.description = description;
this.prompts = orEmpty(prompts);
@@ -68,6 +70,7 @@ public final class AgentSpec {
this.embeddingModelSetups = orEmpty(embeddingModelSetups);
this.vectorStores = orEmpty(vectorStores);
this.mcpServers = orEmpty(mcpServers);
+ this.subagents = orEmpty(subagents);
}
private static <T> List<T> orEmpty(List<T> list) {
@@ -121,4 +124,8 @@ public final class AgentSpec {
public List<DescriptorSpec> getMcpServers() {
return mcpServers;
}
+
+ public List<DescriptorSpec> getSubagents() {
+ return subagents;
+ }
}
diff --git a/api/src/main/java/org/apache/flink/agents/api/yaml/spec/YamlAgentsDocument.java b/api/src/main/java/org/apache/flink/agents/api/yaml/spec/YamlAgentsDocument.java
index 1b0fbfa99..019c06d4e 100644
--- a/api/src/main/java/org/apache/flink/agents/api/yaml/spec/YamlAgentsDocument.java
+++ b/api/src/main/java/org/apache/flink/agents/api/yaml/spec/YamlAgentsDocument.java
@@ -39,6 +39,7 @@ public final class YamlAgentsDocument {
private final List<DescriptorSpec> embeddingModelSetups;
private final List<DescriptorSpec> vectorStores;
private final List<DescriptorSpec> mcpServers;
+ private final List<DescriptorSpec> subagents;
@JsonCreator
public YamlAgentsDocument(
@@ -53,7 +54,8 @@ public final class YamlAgentsDocument {
List<DescriptorSpec> embeddingModelConnections,
@JsonProperty("embedding_model_setups") List<DescriptorSpec> embeddingModelSetups,
@JsonProperty("vector_stores") List<DescriptorSpec> vectorStores,
- @JsonProperty("mcp_servers") List<DescriptorSpec> mcpServers) {
+ @JsonProperty("mcp_servers") List<DescriptorSpec> mcpServers,
+ @JsonProperty("subagents") List<DescriptorSpec> subagents) {
this.agents = orEmpty(agents);
this.prompts = orEmpty(prompts);
this.tools = orEmpty(tools);
@@ -65,6 +67,7 @@ public final class YamlAgentsDocument {
this.embeddingModelSetups = orEmpty(embeddingModelSetups);
this.vectorStores = orEmpty(vectorStores);
this.mcpServers = orEmpty(mcpServers);
+ this.subagents = orEmpty(subagents);
}
private static <T> List<T> orEmpty(List<T> list) {
@@ -114,4 +117,8 @@ public final class YamlAgentsDocument {
public List<DescriptorSpec> getMcpServers() {
return mcpServers;
}
+
+ public List<DescriptorSpec> getSubagents() {
+ return subagents;
+ }
}
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
new file mode 100644
index 000000000..a0dfb7a29
--- /dev/null
+++ b/api/src/test/java/org/apache/flink/agents/api/subagent/SubagentRegisterTest.java
@@ -0,0 +1,68 @@
+/*
+ * 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.api.subagent;
+
+import org.apache.flink.agents.api.agents.Agent;
+import org.apache.flink.agents.api.resource.ResourceType;
+import org.junit.jupiter.api.Test;
+
+import java.util.Map;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertSame;
+import static org.junit.jupiter.api.Assertions.assertThrows;
+
+/** Tests registering sub-agents as AGENT resources. */
+class SubagentRegisterTest {
+
+ @Test
+ void registerSubagentSetupAsResource() {
+ Agent agent = new Agent();
+ TestSubagentSetup setup = new TestSubagentSetup();
+ agent.addResource("reviewer", ResourceType.AGENT, setup);
+
+ Map<String, Object> agentResources = agent.getResources().get(ResourceType.AGENT);
+ assertEquals(1, agentResources.size());
+ assertSame(setup, agentResources.get("reviewer"));
+ assertEquals(ResourceType.AGENT, setup.getResourceType());
+ }
+
+ @Test
+ void duplicateNameThrows() {
+ Agent agent = new Agent();
+ agent.addResource("reviewer", ResourceType.AGENT, new TestSubagentSetup());
+ assertThrows(
+ IllegalArgumentException.class,
+ () -> agent.addResource("reviewer", ResourceType.AGENT, new TestSubagentSetup()));
+ }
+
+ @Test
+ void multipleSubagentsRegistered() {
+ Agent agent = new Agent();
+ TestSubagentSetup reviewer = new TestSubagentSetup();
+ TestSubagentSetup coder = new TestSubagentSetup();
+ agent.addResource("reviewer", ResourceType.AGENT, reviewer);
+ agent.addResource("coder", ResourceType.AGENT, coder);
+
+ Map<String, Object> agentResources = agent.getResources().get(ResourceType.AGENT);
+ assertEquals(2, agentResources.size());
+ assertSame(reviewer, agentResources.get("reviewer"));
+ assertSame(coder, agentResources.get("coder"));
+ }
+}
diff --git a/api/src/test/java/org/apache/flink/agents/api/subagent/SubagentResultTest.java b/api/src/test/java/org/apache/flink/agents/api/subagent/SubagentResultTest.java
new file mode 100644
index 000000000..3a2f3ce4d
--- /dev/null
+++ b/api/src/test/java/org/apache/flink/agents/api/subagent/SubagentResultTest.java
@@ -0,0 +1,75 @@
+/*
+ * 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.api.subagent;
+
+import com.fasterxml.jackson.databind.ObjectMapper;
+import org.junit.jupiter.api.Test;
+
+import java.util.Map;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
+/**
+ * Pins the payload-typing behavior of {@link SubagentResult} across a durable-style JSON
+ * round-trip.
+ */
+public class SubagentResultTest {
+
+ /** Plain mapper with no polymorphic typing, as used to re-bind durable results on recovery. */
+ private static final ObjectMapper RECOVERY_MAPPER = new ObjectMapper();
+
+ /** A POJO payload, standing in for a record returned by an external sub-agent. */
+ public static class Review {
+ public String verdict;
+ public int score;
+
+ public Review() {}
+
+ Review(String verdict, int score) {
+ this.verdict = verdict;
+ this.score = score;
+ }
+ }
+
+ @Test
+ void typedAccessorConvertsThePayloadOnFirstExecution() {
+ SubagentResult result = SubagentResult.ok(new Review("approve", 7));
+
+ Review review = result.getResult(Review.class);
+
+ assertThat(review.verdict).isEqualTo("approve");
+ assertThat(review.score).isEqualTo(7);
+ }
+
+ @Test
+ void typedAccessorRecoversThePayloadTypeAfterAJsonRoundTrip() throws Exception {
+ SubagentResult original = SubagentResult.ok(new Review("approve", 7));
+ String serialized = RECOVERY_MAPPER.writeValueAsString(original);
+
+ // Recovery re-binds through a plain mapper: the payload degrades to a LinkedHashMap.
+ SubagentResult replayed = RECOVERY_MAPPER.readValue(serialized, SubagentResult.class);
+
+ assertThat(replayed.getResult()).isInstanceOf(Map.class);
+
+ Review review = replayed.getResult(Review.class);
+
+ assertThat(review.verdict).isEqualTo("approve");
+ assertThat(review.score).isEqualTo(7);
+ }
+}
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
new file mode 100644
index 000000000..c9ee099fa
--- /dev/null
+++ b/api/src/test/java/org/apache/flink/agents/api/subagent/TestSubagentSetup.java
@@ -0,0 +1,91 @@
+/*
+ * 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.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;
+
+import javax.annotation.Nullable;
+
+/**
+ * Shared {@link SubagentSetup} test double, constructible directly or from a {@link
+ * ResourceDescriptor} (the YAML shape). A pure api-layer descriptor: invocation behavior lives in
+ * the runtime layer, so the {@code submit} forms throw.
+ */
+public class TestSubagentSetup extends SubagentSetup {
+
+ private static final long serialVersionUID = 1L;
+
+ @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;
+ }
+
+ /** 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")));
+ }
+
+ @Nullable
+ public String getEndpoint() {
+ return endpoint;
+ }
+
+ public boolean isFailOnCall() {
+ return failOnCall;
+ }
+
+ @Override
+ public SubagentFuture submit(
+ RunnerContext ctx, Object prompt, String sessionId, String callId) {
+ throw new UnsupportedOperationException(
+ "Descriptor-only sub-agent setup; invocation lives in the runtime layer.");
+ }
+
+ @Override
+ public SubagentFuture submit(RunnerContext ctx, Object prompt, String sessionId) {
+ throw new UnsupportedOperationException(
+ "Descriptor-only sub-agent setup; invocation lives in the runtime layer.");
+ }
+
+ @Override
+ public SubagentFuture submit(RunnerContext ctx, Object prompt) {
+ throw new UnsupportedOperationException(
+ "Descriptor-only sub-agent setup; invocation lives in the runtime layer.");
+ }
+}
diff --git a/dev/agent-skills/flink-agents-dev/assets/yaml-contracts.yaml b/dev/agent-skills/flink-agents-dev/assets/yaml-contracts.yaml
index abcc37890..fc00ea8af 100644
--- a/dev/agent-skills/flink-agents-dev/assets/yaml-contracts.yaml
+++ b/dev/agent-skills/flink-agents-dev/assets/yaml-contracts.yaml
@@ -13,7 +13,7 @@ contracts:
repository: apache/flink-agents
ref: main
path: docs/yaml-schema.json
- blob_sha: 78629a46d42d96d6fe177250f6c91ef95d5a9d3a
+ blob_sha: 77634cc58e41d4a62640c2964e5634197b6991e3
versions_without_yaml_api:
- "0.2.1"
diff --git a/dev/agent-skills/flink-agents-dev/assets/yaml-schema.json b/dev/agent-skills/flink-agents-dev/assets/yaml-schema.json
index 78629a46d..77634cc58 100644
--- a/dev/agent-skills/flink-agents-dev/assets/yaml-schema.json
+++ b/dev/agent-skills/flink-agents-dev/assets/yaml-schema.json
@@ -149,6 +149,13 @@
"title": "Skills",
"type": "array"
},
+ "subagents": {
+ "items": {
+ "$ref": "#/$defs/DescriptorSpec"
+ },
+ "title": "Subagents",
+ "type": "array"
+ },
"tools": {
"items": {
"$ref": "#/$defs/ToolSpec"
@@ -553,6 +560,13 @@
"title": "Skills",
"type": "array"
},
+ "subagents": {
+ "items": {
+ "$ref": "#/$defs/DescriptorSpec"
+ },
+ "title": "Subagents",
+ "type": "array"
+ },
"tools": {
"items": {
"$ref": "#/$defs/ToolSpec"
diff --git a/docs/yaml-schema.json b/docs/yaml-schema.json
index 78629a46d..77634cc58 100644
--- a/docs/yaml-schema.json
+++ b/docs/yaml-schema.json
@@ -149,6 +149,13 @@
"title": "Skills",
"type": "array"
},
+ "subagents": {
+ "items": {
+ "$ref": "#/$defs/DescriptorSpec"
+ },
+ "title": "Subagents",
+ "type": "array"
+ },
"tools": {
"items": {
"$ref": "#/$defs/ToolSpec"
@@ -553,6 +560,13 @@
"title": "Skills",
"type": "array"
},
+ "subagents": {
+ "items": {
+ "$ref": "#/$defs/DescriptorSpec"
+ },
+ "title": "Subagents",
+ "type": "array"
+ },
"tools": {
"items": {
"$ref": "#/$defs/ToolSpec"
diff --git a/plan/pom.xml b/plan/pom.xml
index 9e12d3bbb..cc1f7dd65 100644
--- a/plan/pom.xml
+++ b/plan/pom.xml
@@ -40,6 +40,13 @@ under the License.
<artifactId>flink-agents-api</artifactId>
<version>${project.version}</version>
</dependency>
+ <dependency>
+ <groupId>org.apache.flink</groupId>
+ <artifactId>flink-agents-api</artifactId>
+ <version>${project.version}</version>
+ <type>test-jar</type>
+ <scope>test</scope>
+ </dependency>
<dependency>
<groupId>org.apache.flink</groupId>
<artifactId>flink-agents-integrations-mcp</artifactId>
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 c865c72ce..051f1fff8 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
@@ -37,6 +37,7 @@ import org.apache.flink.agents.api.resource.ResourceType;
import org.apache.flink.agents.api.resource.SerializableResource;
import org.apache.flink.agents.api.skills.SkillSourceSpec;
import org.apache.flink.agents.api.skills.Skills;
+import org.apache.flink.agents.api.subagent.SubagentSetup;
import org.apache.flink.agents.api.tools.ToolMetadata;
import org.apache.flink.agents.api.tools.ToolParameterInjection;
import org.apache.flink.agents.api.tools.ToolParameterInjectionValidator;
@@ -585,6 +586,30 @@ public class AgentPlan implements Serializable {
+ " method on your Agent class so its tools and prompts can be"
+ " discovered.");
}
+ } else if (type == ResourceType.AGENT) {
+ for (Map.Entry<String, Object> kv : entry.getValue().entrySet()) {
+ String name = kv.getKey();
+ Object value = kv.getValue();
+ if (value instanceof SubagentSetup) {
+ addResourceProvider(
+ JavaSerializableResourceProvider.createResourceProvider(
+ name, ResourceType.AGENT, (SubagentSetup) value));
+ } else if (value instanceof ResourceDescriptor) {
+ // Declared via YAML: the descriptor names a SubagentSetup subclass that is
+ // instantiated when the resource is first resolved.
+ addResourceProvider(
+ createDescriptorResourceProvider(
+ name, ResourceType.AGENT, (ResourceDescriptor) value));
+ } else {
+ throw new IllegalArgumentException(
+ "AGENT resource '"
+ + name
+ + "' must be a SubagentSetup or a ResourceDescriptor, but"
+ + " got "
+ + value.getClass().getName()
+ + ".");
+ }
+ }
} else {
for (Map.Entry<String, Object> kv : entry.getValue().entrySet()) {
ResourceDescriptor descriptor =
diff --git a/plan/src/main/java/org/apache/flink/agents/plan/resourceprovider/ResourceProvider.java b/plan/src/main/java/org/apache/flink/agents/plan/resourceprovider/ResourceProvider.java
index a90bba58b..49c34ab59 100644
--- a/plan/src/main/java/org/apache/flink/agents/plan/resourceprovider/ResourceProvider.java
+++ b/plan/src/main/java/org/apache/flink/agents/plan/resourceprovider/ResourceProvider.java
@@ -62,6 +62,15 @@ public abstract class ResourceProvider implements java.io.Serializable {
return type;
}
+ /**
+ * Whether the given provider materializes a resource owned by the Python runtime, so the
+ * runtime must ask that runtime to build it instead of resolving it on the Java side.
+ */
+ public static boolean isPythonOwned(ResourceProvider provider) {
+ return provider instanceof PythonResourceProvider
+ || provider instanceof PythonSerializableResourceProvider;
+ }
+
/**
* Create resource at runtime.
*
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
new file mode 100644
index 000000000..8c801ae43
--- /dev/null
+++ b/plan/src/test/java/org/apache/flink/agents/plan/AgentPlanSubagentResourceTest.java
@@ -0,0 +1,87 @@
+/*
+ * 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.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.SubagentSetup;
+import org.apache.flink.agents.api.subagent.TestSubagentSetup;
+import org.apache.flink.agents.plan.resourceprovider.ResourceProvider;
+import org.junit.jupiter.api.Test;
+
+import java.util.Map;
+
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
+
+/**
+ * Tests compiling AGENT resources into the agent plan, for both registration shapes: a {@link
+ * SubagentSetup} instance (programmatic) and a {@link ResourceDescriptor} (the YAML shape).
+ */
+public class AgentPlanSubagentResourceTest {
+
+ @Test
+ void subagentSetupInstanceCompilesIntoAgentProvider() throws Exception {
+ Agent agent = new Agent();
+ agent.addResource("reviewer", ResourceType.AGENT, new TestSubagentSetup());
+
+ AgentPlan plan = new AgentPlan(agent);
+
+ Map<String, ResourceProvider> agentProviders =
+ plan.getResourceProviders().get(ResourceType.AGENT);
+ assertThat(agentProviders).containsKey("reviewer");
+ Resource resolved = agentProviders.get("reviewer").provide(null);
+ assertThat(resolved).isInstanceOf(SubagentSetup.class);
+ }
+
+ @Test
+ void agentDescriptorCompilesAndResolvesToSubagentSetup() throws Exception {
+ Agent agent = new Agent();
+ agent.addResource(
+ "summarizer",
+ ResourceType.AGENT,
+ ResourceDescriptor.Builder.newBuilder(TestSubagentSetup.class.getName())
+ .addInitialArgument("endpoint", "http://summarizer:8080")
+ .build());
+
+ AgentPlan plan = new AgentPlan(agent);
+
+ Map<String, ResourceProvider> agentProviders =
+ plan.getResourceProviders().get(ResourceType.AGENT);
+ assertThat(agentProviders).containsKey("summarizer");
+
+ Resource resolved = agentProviders.get("summarizer").provide(null);
+ assertThat(resolved).isInstanceOf(TestSubagentSetup.class);
+ assertThat(((TestSubagentSetup) resolved).getEndpoint())
+ .isEqualTo("http://summarizer:8080");
+ assertThat(resolved.getResourceType()).isEqualTo(ResourceType.AGENT);
+ }
+
+ @Test
+ void nonSubagentAgentResourceIsRejected() {
+ Agent agent = new Agent();
+ agent.getResources().get(ResourceType.AGENT).put("bad", "not-a-subagent");
+
+ assertThatThrownBy(() -> new AgentPlan(agent))
+ .isInstanceOf(IllegalArgumentException.class)
+ .hasMessageContaining("must be a SubagentSetup or a ResourceDescriptor");
+ }
+}
diff --git a/python/flink_agents/api/resource.py b/python/flink_agents/api/resource.py
index 4dcc904f0..934b6c42e 100644
--- a/python/flink_agents/api/resource.py
+++ b/python/flink_agents/api/resource.py
@@ -32,7 +32,7 @@ class ResourceType(Enum):
"""Type enum of resource.
Currently, support chat_model, chat_model_server, tool, embedding_model,
- vector_store, prompt, mcp_server, skills, model_router.
+ vector_store, prompt, mcp_server, skills, model_router, agent.
"""
CHAT_MODEL = "chat_model"
@@ -49,6 +49,7 @@ class ResourceType(Enum):
# Python side: mixed jobs (Java router + Python actions) must not fail at
# operator open with a ValidationError.
MODEL_ROUTER = "model_router"
+ AGENT = "agent"
class Resource(BaseModel, ABC):
diff --git a/python/flink_agents/api/runner_context.py b/python/flink_agents/api/runner_context.py
index 58990a4f9..b84bc7226 100644
--- a/python/flink_agents/api/runner_context.py
+++ b/python/flink_agents/api/runner_context.py
@@ -240,6 +240,7 @@ class RunnerContext(ABC):
func: Callable[[Any], Any],
*args: Any,
reconciler: Callable[[], Any] | None = None,
+ durable_id: str | None = None,
**kwargs: Any,
) -> Any:
"""Synchronously execute the provided function with durable execution support.
@@ -282,6 +283,12 @@ class RunnerContext(ABC):
Optional zero-argument reconciler callable used only during recovery.
This is a reserved keyword-only parameter and is not forwarded to
`func`.
+ durable_id : str | None
+ Optional stable identity keying this call's persisted state. Supply
+ it when the caller owns an identity that survives failover;
+ otherwise the identity is derived from the callable and its
+ arguments. Reserved keyword-only parameter, not forwarded to
+ `func`.
**kwargs : Any
Keyword arguments to pass to the function.
@@ -297,6 +304,7 @@ class RunnerContext(ABC):
func: Callable[[Any], Any],
*args: Any,
reconciler: Callable[[], Any] | None = None,
+ durable_id: str | None = None,
**kwargs: Any,
) -> "AsyncExecutionResult":
"""Asynchronously execute the provided function with durable execution support.
@@ -340,6 +348,12 @@ class RunnerContext(ABC):
Optional zero-argument reconciler callable used only during recovery.
This is a reserved keyword-only parameter and is not forwarded to
`func`.
+ durable_id : str | None
+ Optional stable identity keying this call's persisted state. Supply
+ it when the caller owns an identity that survives failover;
+ otherwise the identity is derived from the callable and its
+ arguments. Reserved keyword-only parameter, not forwarded to
+ `func`.
**kwargs : Any
Keyword arguments to pass to the function.
diff --git a/python/flink_agents/api/subagent.py b/python/flink_agents/api/subagent.py
new file mode 100644
index 000000000..0b15a0557
--- /dev/null
+++ b/python/flink_agents/api/subagent.py
@@ -0,0 +1,186 @@
+################################################################################
+# 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.
+#################################################################################
+import logging
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Any
+
+from flink_agents.api.resource import ResourceType, SerializableResource
+
+if TYPE_CHECKING:
+ from flink_agents.api.runner_context import RunnerContext
+
+_LOG = logging.getLogger(__name__)
+
+
+@dataclass
+class SubagentResult:
+ """Outcome of a sub-agent call issued through :class:`SubagentSetup`.
+
+ A successful outcome carries a JSON-serializable ``result``, and a failed
+ one carries a serializable ``error_message``.
+
+ Implementations capture their internal failures into a result through
+ :meth:`error` instead of raising, so callers inspect ``success`` rather
+ than catching. Because the failure is carried as a message rather than a
+ live exception, the whole result can be persisted through durable
+ execution and survive a failover.
+ """
+
+ success: bool
+ result: Any = None
+ error_message: str | None = None
+
+ @staticmethod
+ def ok(result: Any) -> "SubagentResult":
+ """Create a successful result carrying ``result``."""
+ return SubagentResult(success=True, result=result)
+
+ @staticmethod
+ def error(error: BaseException | str) -> "SubagentResult":
+ """Create a failed result from an exception or a plain message.
+
+ For an exception the exception's type and message are stored as a
+ serializable string so the result can survive durable execution. The
+ full stack trace is logged here rather than persisted, keeping the
+ durable payload bounded.
+ """
+ if isinstance(error, BaseException):
+ _LOG.warning(
+ "Sub-agent call failed; persisting the exception summary.",
+ exc_info=(
+ type(error),
+ error,
+ error.__traceback__,
+ ),
+ )
+ message = f"{type(error).__name__}: {error}"
+ else:
+ message = error
+ return SubagentResult(success=False, error_message=message)
+
+ @property
+ def exception(self) -> Exception | None:
+ """Reconstruct an exception carrying the stored summary; None on success."""
+ return None if self.success else RuntimeError(self.error_message)
+
+
+class SubagentFuture(ABC):
+ """Handle for one sub-agent invocation, identified by the
+ ``(session_id, call_id)`` pair that keys the invocation.
+ """
+
+ def __init__(self, session_id: str, call_id: str) -> None:
+ """Initialize with the invocation identity."""
+ self._session_id = session_id
+ self._call_id = call_id
+
+ @property
+ def session_id(self) -> str:
+ """The session this invocation belongs to."""
+ return self._session_id
+
+ @property
+ def call_id(self) -> str:
+ """The id of this invocation within its session."""
+ return self._call_id
+
+ @property
+ def identity(self) -> str:
+ """The ``session_id#call_id`` string keying this invocation."""
+ return f"{self._session_id}#{self._call_id}"
+
+ @abstractmethod
+ def done(self) -> bool:
+ """Whether the invocation has been resolved."""
+
+ def cancel(self) -> None: # noqa: B027 - deliberate no-op default
+ """Request cancellation of the invocation."""
+
+ @abstractmethod
+ def combine(self, *others: "SubagentFuture") -> "SubagentFutures":
+ """Group this handle with others to be resolved together."""
+
+ @abstractmethod
+ def __await__(self) -> Any:
+ """Resolve the invocation, waiting until it reaches a terminal state.
+
+ Failures converge into a failed :class:`SubagentResult` rather than a
+ separately raised exception.
+ """
+
+
+class SubagentFutures(ABC):
+ """A group of sub-agent handles to be resolved together.
+
+ A group is not itself an invocation and carries no
+ ``(session_id, call_id)`` identity.
+ """
+
+ @abstractmethod
+ def done(self) -> bool:
+ """Whether every handle in the group has been resolved."""
+
+ def cancel(self) -> None: # noqa: B027 - deliberate no-op default
+ """Propagate the cancellation request to every handle in the group."""
+
+ @abstractmethod
+ def combine(self, *others: "SubagentFuture") -> "SubagentFutures":
+ """Add more handles to the group."""
+
+ @abstractmethod
+ def __await__(self) -> Any:
+ """Resolve every handle in the group and return their outcomes in the
+ order the handles were added. Like awaiting a single handle, failures
+ surface through failed :class:`SubagentResult`s.
+ """
+
+
+class SubagentSetup(SerializableResource):
+ """Caller-facing definition of a sub-agent, registered as an AGENT resource."""
+
+ @classmethod
+ def resource_type(cls) -> ResourceType:
+ """Return resource type of class."""
+ return ResourceType.AGENT
+
+ @abstractmethod
+ async def submit(
+ self,
+ ctx: "RunnerContext",
+ prompt: Any,
+ session_id: str | None = None,
+ call_id: str | None = None,
+ ) -> SubagentFuture:
+ """Issue one invocation and return its handle.
+
+ Declared ``async`` to reserve the ability to await while the request
+ is being issued, so that one calling form holds whether or not an
+ implementation has anything to await.
+
+ Without ids, the implementation assigns the identity and starts a
+ fresh conversation. This is the preferred form.
+
+ Pass ``session_id`` to continue the conversation of an earlier
+ invocation. The session id is available on the handle returned by
+ that invocation. Whether a conversation can be continued across
+ actions is up to the concrete implementation.
+
+ The complete ``(session_id, call_id)`` identity is reserved for
+ implementation use.
+ """
diff --git a/python/flink_agents/api/tests/subagent_test_utils.py b/python/flink_agents/api/tests/subagent_test_utils.py
new file mode 100644
index 000000000..ee1b25897
--- /dev/null
+++ b/python/flink_agents/api/tests/subagent_test_utils.py
@@ -0,0 +1,49 @@
+################################################################################
+# 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.
+################################################################################
+"""Shared sub-agent test doubles."""
+
+from typing import TYPE_CHECKING, Any
+
+from flink_agents.api.subagent import SubagentSetup
+
+if TYPE_CHECKING:
+ from flink_agents.api.runner_context import RunnerContext
+ from flink_agents.api.subagent import SubagentFuture
+
+
+class TestSubagentSetup(SubagentSetup):
+ """Shared ``SubagentSetup`` test double, constructible directly or from a
+ resource descriptor (the YAML shape).
+
+ A pure api-layer descriptor: invocation behavior lives in the runtime
+ layer, so the ``submit`` forms raise.
+ """
+
+ endpoint_url: str | None = None
+ fail_on_call: bool = False
+
+ def submit(
+ self,
+ ctx: "RunnerContext",
+ prompt: Any,
+ session_id: str | None = None,
+ call_id: str | None = None,
+ ) -> "SubagentFuture":
+ """Descriptor-only double; invocation lives in the runtime layer."""
+ msg = "Descriptor-only sub-agent setup; invocation lives in the runtime layer."
+ raise NotImplementedError(msg)
diff --git a/python/flink_agents/api/tests/test_subagent.py b/python/flink_agents/api/tests/test_subagent.py
new file mode 100644
index 000000000..a6cebfa8b
--- /dev/null
+++ b/python/flink_agents/api/tests/test_subagent.py
@@ -0,0 +1,60 @@
+################################################################################
+# 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.
+################################################################################
+"""Tests registering sub-agents as AGENT resources."""
+import pytest
+
+from flink_agents.api.agents.agent import Agent
+from flink_agents.api.resource import ResourceType
+from flink_agents.api.tests.subagent_test_utils import TestSubagentSetup
+
+
+def test_register_subagent_setup_as_resource() -> None:
+ """A ``SubagentSetup`` registers under the AGENT resource map."""
+ agent = Agent()
+ setup = TestSubagentSetup()
+
+ agent.add_resource("reviewer", ResourceType.AGENT, setup)
+
+ agent_resources = agent.resources[ResourceType.AGENT]
+ assert len(agent_resources) == 1
+ assert agent_resources["reviewer"] is setup
+ assert setup.resource_type() == ResourceType.AGENT
+
+
+def test_duplicate_name_throws() -> None:
+ """Registering a duplicate AGENT name raises."""
+ agent = Agent()
+ agent.add_resource("reviewer", ResourceType.AGENT, TestSubagentSetup())
+
+ with pytest.raises(ValueError):
+ agent.add_resource("reviewer", ResourceType.AGENT, TestSubagentSetup())
+
+
+def test_multiple_subagents_registered() -> None:
+ """Multiple distinct AGENT resources coexist."""
+ agent = Agent()
+ reviewer = TestSubagentSetup()
+ coder = TestSubagentSetup()
+
+ agent.add_resource("reviewer", ResourceType.AGENT, reviewer)
+ agent.add_resource("coder", ResourceType.AGENT, coder)
+
+ agent_resources = agent.resources[ResourceType.AGENT]
+ assert len(agent_resources) == 2
+ assert agent_resources["reviewer"] is reviewer
+ assert agent_resources["coder"] is coder
diff --git a/python/flink_agents/api/yaml/loader.py b/python/flink_agents/api/yaml/loader.py
index 5cfce8577..30e754851 100644
--- a/python/flink_agents/api/yaml/loader.py
+++ b/python/flink_agents/api/yaml/loader.py
@@ -62,6 +62,7 @@ _DESCRIPTOR_TYPES: Dict[str, ResourceType] = {
"embedding_model_setups": ResourceType.EMBEDDING_MODEL,
"vector_stores": ResourceType.VECTOR_STORE,
"mcp_servers": ResourceType.MCP_SERVER,
+ "subagents": ResourceType.AGENT,
}
diff --git a/python/flink_agents/api/yaml/specs.py b/python/flink_agents/api/yaml/specs.py
index ee5f58ad9..11cadf5b1 100644
--- a/python/flink_agents/api/yaml/specs.py
+++ b/python/flink_agents/api/yaml/specs.py
@@ -260,6 +260,7 @@ class AgentSpec(BaseModel):
embedding_model_setups: List[DescriptorSpec] = Field(default_factory=list)
vector_stores: List[DescriptorSpec] = Field(default_factory=list)
mcp_servers: List[DescriptorSpec] = Field(default_factory=list)
+ subagents: List[DescriptorSpec] = Field(default_factory=list)
class YamlAgentsDocument(BaseModel):
@@ -289,6 +290,7 @@ class YamlAgentsDocument(BaseModel):
embedding_model_setups: List[DescriptorSpec] = Field(default_factory=list)
vector_stores: List[DescriptorSpec] = Field(default_factory=list)
mcp_servers: List[DescriptorSpec] = Field(default_factory=list)
+ subagents: List[DescriptorSpec] = Field(default_factory=list)
def export() -> str:
diff --git a/python/flink_agents/e2e_tests/e2e_tests_integration/subagent_external_integration_agent.py b/python/flink_agents/e2e_tests/e2e_tests_integration/subagent_external_integration_agent.py
new file mode 100644
index 000000000..c7818574b
--- /dev/null
+++ b/python/flink_agents/e2e_tests/e2e_tests_integration/subagent_external_integration_agent.py
@@ -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.
+################################################################################
+"""Agents exercising the Python external sub-agent modes.
+
+An async (durable pub/sub) setup and a deferred setup, each driven by a Python
+action running on the Java runtime over pemja. The backend is an in-memory run
+store held on the setup instance (pemja's Python runs in the MiniCluster JVM,
+not the test process), so the test needs no external service while still
+exercising the submit / poll / fetch sequence of each mode. The backend
+completes on the first probe, so the multi-probe pacing of the poll loop is
+covered at unit level rather than here.
+"""
+
+from typing import Any
+
+from pydantic import PrivateAttr
+from typing_extensions import override
+
+from flink_agents.api.agents.agent import Agent
+from flink_agents.api.decorators import action
+from flink_agents.api.events.event import Event, InputEvent, OutputEvent
+from flink_agents.api.events.event_type import EventType
+from flink_agents.api.resource import ResourceType
+from flink_agents.api.runner_context import RunnerContext
+from flink_agents.api.subagent import SubagentResult
+from flink_agents.runtime.async_subagent import BaseAsyncSubagentSetup, RunStatus
+from flink_agents.runtime.deferred_subagent import (
+ DeferredSubagentSetup,
+ PreparedTriple,
+)
+
+
+def _outcome(result: SubagentResult) -> str:
+ """Render a sub-agent result as the string emitted downstream."""
+ return result.result if result.success else f"ERR:{result.error_message}"
+
+
+class InMemoryAsyncSubagentSetup(BaseAsyncSubagentSetup):
+ """External async setup backed by an in-memory run store.
+
+ A prompt containing ``fail`` produces a failed run; any other prompt
+ completes and echoes back, tagged with the injected sub-agent name.
+ """
+
+ _runs: dict = PrivateAttr(default_factory=dict)
+
+ @override
+ def call_submit_request(self, session_id: str, call_id: str, prompt: Any) -> None:
+ """Record the run under its (session_id, call_id) identity."""
+ self._runs[(session_id, call_id)] = prompt
+
+ @override
+ def call_query_status(self, session_id: str, call_id: str) -> RunStatus:
+ """Report the run as terminal immediately (completed or failed)."""
+ if (session_id, call_id) not in self._runs:
+ return RunStatus.not_started()
+ prompt = self._runs[(session_id, call_id)]
+ if "fail" in str(prompt):
+ return RunStatus.failed("async run failed on demand")
+ return RunStatus.completed()
+
+ @override
+ def call_fetch_result(self, session_id: str, call_id: str) -> SubagentResult:
+ """Fetch the completed run's echoed answer."""
+ prompt = self._runs[(session_id, call_id)]
+ return SubagentResult.ok(f"async[{self.subagent_name}]:{prompt}")
+
+
+class InMemoryDeferredSubagentSetup(DeferredSubagentSetup):
+ """External deferred setup that runs the whole invocation on resolve."""
+
+ @override
+ def prepare(
+ self,
+ ctx: RunnerContext,
+ prompt: Any,
+ session_id: str,
+ call_id: str,
+ ) -> PreparedTriple:
+ """Return a durable call echoing the prompt (or failing on demand)."""
+ name = self.subagent_name
+
+ def call() -> SubagentResult:
+ if "fail" in str(prompt):
+ return SubagentResult.error("deferred run failed on demand")
+ return SubagentResult.ok(f"deferred[{name}]:{prompt}")
+
+ return (f"{session_id}#{call_id}", call, None)
+
+
+class AsyncExternalAgent(Agent):
+ """Agent whose action submits to the async external sub-agent and awaits it."""
+
+ @action(EventType.InputEvent)
+ @staticmethod
+ async def process(event: Event, ctx: RunnerContext) -> None:
+ """Submit, await, and emit the async sub-agent outcome."""
+ prompt = InputEvent.from_event(event).input
+ reviewer = ctx.get_resource("reviewer", ResourceType.AGENT)
+ # Awaiting the submit hands back the handle once the durable POST has
+ # landed.
+ future = await reviewer.submit(ctx, prompt)
+ result = await future
+ ctx.send_event(OutputEvent(output=_outcome(result)))
+
+
+class DeferredExternalAgent(Agent):
+ """Agent whose action submits to the deferred external sub-agent and awaits."""
+
+ @action(EventType.InputEvent)
+ @staticmethod
+ async def process(event: Event, ctx: RunnerContext) -> None:
+ """Submit, await, and emit the deferred sub-agent outcome."""
+ prompt = InputEvent.from_event(event).input
+ reviewer = ctx.get_resource("reviewer", ResourceType.AGENT)
+ # This mode sends nothing until the handle is awaited.
+ future = await reviewer.submit(ctx, prompt)
+ result = await future
+ ctx.send_event(OutputEvent(output=_outcome(result)))
+
+
+def build_async_agent() -> Agent:
+ """Build the agent registering the async external sub-agent."""
+ agent = AsyncExternalAgent()
+ agent.add_resource("reviewer", ResourceType.AGENT, InMemoryAsyncSubagentSetup())
+ return agent
+
+
+def build_deferred_agent() -> Agent:
+ """Build the agent registering the deferred external sub-agent."""
+ agent = DeferredExternalAgent()
+ agent.add_resource("reviewer", ResourceType.AGENT, InMemoryDeferredSubagentSetup())
+ return agent
diff --git a/python/flink_agents/e2e_tests/e2e_tests_integration/subagent_external_integration_test.py b/python/flink_agents/e2e_tests/e2e_tests_integration/subagent_external_integration_test.py
new file mode 100644
index 000000000..a2e531904
--- /dev/null
+++ b/python/flink_agents/e2e_tests/e2e_tests_integration/subagent_external_integration_test.py
@@ -0,0 +1,96 @@
+################################################################################
+# 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.
+################################################################################
+"""Integration tests for the Python external sub-agent modes.
+
+Each mode runs a real embedded Flink job (MiniCluster) so the Python action
+executes on the Java runtime, exercising the async (durable pub/sub) and
+deferred external setups end to end, including their success and failure
+outcomes.
+"""
+
+import json
+import os
+import sysconfig
+from collections.abc import Callable
+from pathlib import Path
+
+from pyflink.common import Configuration, Encoder
+from pyflink.common.typeinfo import Types
+from pyflink.datastream import RuntimeExecutionMode, StreamExecutionEnvironment
+
+from flink_agents.api.agents.agent import Agent
+from flink_agents.api.execution_environment import AgentsExecutionEnvironment
+from flink_agents.e2e_tests.e2e_tests_integration.subagent_external_integration_agent import (
+ build_async_agent,
+ build_deferred_agent,
+)
+
+os.environ["PYTHONPATH"] = sysconfig.get_paths()["purelib"]
+
+
+def _run_agent(agent_factory: Callable[[], Agent], result_dir: Path) -> list[str]:
+ from pyflink.datastream.connectors.file_system import StreamingFileSink
+
+ config = Configuration()
+ config.set_string("state.backend.type", "rocksdb")
+ config.set_string("execution.checkpointing.interval", "1s")
+ config.set_string("restart-strategy.type", "disable")
+ env = StreamExecutionEnvironment.get_execution_environment(config)
+ env.set_runtime_mode(RuntimeExecutionMode.STREAMING)
+ env.set_parallelism(1)
+
+ input_stream = env.from_collection(["ok-input", "please-fail"])
+
+ agents_env = AgentsExecutionEnvironment.get_execution_environment(env=env)
+ output_datastream = (
+ agents_env.from_datastream(input=input_stream, key_selector=lambda x: x)
+ .apply(agent_factory())
+ .to_datastream()
+ )
+
+ result_dir.mkdir(parents=True, exist_ok=True)
+ output_datastream.map(lambda x: json.dumps(x), Types.STRING()).add_sink(
+ StreamingFileSink.for_row_format(
+ base_path=str(result_dir.absolute()),
+ encoder=Encoder.simple_string_encoder(),
+ ).build()
+ )
+ agents_env.execute()
+
+ lines: list[str] = []
+ for file in result_dir.rglob("*"):
+ if file.is_file():
+ with file.open() as f:
+ lines.extend(line.strip() for line in f if line.strip())
+ return lines
+
+
+def test_async_external_subagent(tmp_path: Path) -> None:
+ """The async external sub-agent completes and fails on the Java runtime."""
+ results = _run_agent(build_async_agent, tmp_path / "results")
+ assert sorted(results) == sorted(
+ ['"async[reviewer]:ok-input"', '"ERR:async run failed on demand"']
+ )
+
+
+def test_deferred_external_subagent(tmp_path: Path) -> None:
+ """The deferred external sub-agent completes and fails on the Java runtime."""
+ results = _run_agent(build_deferred_agent, tmp_path / "results")
+ assert sorted(results) == sorted(
+ ['"deferred[reviewer]:ok-input"', '"ERR:deferred run failed on demand"']
+ )
diff --git a/python/flink_agents/e2e_tests/e2e_tests_integration/subagent_integration_agent.py b/python/flink_agents/e2e_tests/e2e_tests_integration/subagent_integration_agent.py
new file mode 100644
index 000000000..7eb252c0b
--- /dev/null
+++ b/python/flink_agents/e2e_tests/e2e_tests_integration/subagent_integration_agent.py
@@ -0,0 +1,87 @@
+################################################################################
+# 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.
+################################################################################
+"""Python integration agent that uses a Python sub-agent.
+
+A Python action submits to a Python deferred sub-agent and awaits it. As for
+any Python agent, the action runs on the Java runtime (the operator drives it
+over pemja); the sub-agent's id allocation and unresolved-handle enforcement
+follow from that.
+"""
+
+from typing import Any
+
+from typing_extensions import override
+
+from flink_agents.api.agents.agent import Agent
+from flink_agents.api.decorators import action
+from flink_agents.api.events.event import Event, InputEvent, OutputEvent
+from flink_agents.api.events.event_type import EventType
+from flink_agents.api.resource import ResourceType
+from flink_agents.api.runner_context import RunnerContext
+from flink_agents.api.subagent import SubagentResult
+from flink_agents.runtime.deferred_subagent import (
+ DeferredSubagentSetup,
+ PreparedTriple,
+)
+
+
+class EchoSubagentSetup(DeferredSubagentSetup):
+ """In-process deferred sub-agent that echoes the prompt back.
+
+ The subagent name it reports must be the resource name the framework
+ injects; the action asserts on it to prove name injection at runtime.
+ """
+
+ @override
+ def prepare(
+ self,
+ ctx: RunnerContext,
+ prompt: Any,
+ session_id: str,
+ call_id: str,
+ ) -> PreparedTriple:
+ """Return a durable call echoing the prompt, keyed by the identity."""
+ name = self.subagent_name
+
+ def call() -> SubagentResult:
+ return SubagentResult.ok(f"reviewed[{name}]:{prompt}")
+
+ return (f"{session_id}#{call_id}", call, None)
+
+
+class SubagentIntegrationAgent(Agent):
+ """Python agent whose action calls a Python sub-agent and awaits it."""
+
+ @action(EventType.InputEvent)
+ @staticmethod
+ async def process(event: Event, ctx: RunnerContext) -> None:
+ """Submit the input to the sub-agent, await it, and emit its result."""
+ prompt = InputEvent.from_event(event).input
+ reviewer = ctx.get_resource("reviewer", ResourceType.AGENT)
+ # Short-form submit: ids are allocated from the executing task, which
+ # only works when the operator forwarded on_action_prepared over pemja.
+ future = await reviewer.submit(ctx, prompt)
+ result = await future
+ ctx.send_event(OutputEvent(output=result.result))
+
+
+def build_agent() -> Agent:
+ """Build the agent with the sub-agent registered as an AGENT resource."""
+ agent = SubagentIntegrationAgent()
+ agent.add_resource("reviewer", ResourceType.AGENT, EchoSubagentSetup())
+ return agent
diff --git a/python/flink_agents/e2e_tests/e2e_tests_integration/subagent_integration_test.py b/python/flink_agents/e2e_tests/e2e_tests_integration/subagent_integration_test.py
new file mode 100644
index 000000000..3988f31a6
--- /dev/null
+++ b/python/flink_agents/e2e_tests/e2e_tests_integration/subagent_integration_test.py
@@ -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.
+################################################################################
+"""Integration test: a Python agent using a Python sub-agent.
+
+Runs a real embedded Flink job (MiniCluster) so the Python action executes on
+the Java runtime, exercising the sub-agent path end to end: name injection,
+deterministic id allocation on submit, the durable await, and the emitted
+result flowing back downstream.
+"""
+
+import json
+import os
+import sysconfig
+from pathlib import Path
+
+from pyflink.common import Configuration, Encoder
+from pyflink.common.typeinfo import Types
+from pyflink.datastream import RuntimeExecutionMode, StreamExecutionEnvironment
+from pyflink.datastream.connectors.file_system import StreamingFileSink
+
+from flink_agents.api.execution_environment import AgentsExecutionEnvironment
+from flink_agents.e2e_tests.e2e_tests_integration.subagent_integration_agent import (
+ build_agent,
+)
+
+os.environ["PYTHONPATH"] = sysconfig.get_paths()["purelib"]
+
+
+def test_python_agent_uses_python_subagent(tmp_path: Path) -> None:
+ """The Python sub-agent runs end to end on the Java runtime."""
+ config = Configuration()
+ config.set_string("state.backend.type", "rocksdb")
+ config.set_string("execution.checkpointing.interval", "1s")
+ config.set_string("restart-strategy.type", "disable")
+ env = StreamExecutionEnvironment.get_execution_environment(config)
+ env.set_runtime_mode(RuntimeExecutionMode.STREAMING)
+ env.set_parallelism(1)
+
+ input_stream = env.from_collection(["alpha", "beta"])
+
+ agents_env = AgentsExecutionEnvironment.get_execution_environment(env=env)
+ output_datastream = (
+ agents_env.from_datastream(input=input_stream, key_selector=lambda x: x)
+ .apply(build_agent())
+ .to_datastream()
+ )
+
+ result_dir = tmp_path / "results"
+ result_dir.mkdir(parents=True, exist_ok=True)
+ output_datastream.map(lambda x: json.dumps(x), Types.STRING()).add_sink(
+ StreamingFileSink.for_row_format(
+ base_path=str(result_dir.absolute()),
+ encoder=Encoder.simple_string_encoder(),
+ ).build()
+ )
+
+ agents_env.execute()
+
+ results = _read_results(result_dir)
+ # The sub-agent echoes the prompt, tagged with the injected resource name,
+ # proving name injection + deterministic id allocation happened on the
+ # runtime (short-form submit would fail otherwise).
+ assert sorted(results) == sorted(
+ ['"reviewed[reviewer]:alpha"', '"reviewed[reviewer]:beta"']
+ )
+
+
+def _read_results(result_dir: Path) -> list[str]:
+ lines: list[str] = []
+ for file in result_dir.rglob("*"):
+ if file.is_file():
+ with file.open() as f:
+ lines.extend(line.strip() for line in f if line.strip())
+ return lines
diff --git a/python/flink_agents/plan/agent_plan.py b/python/flink_agents/plan/agent_plan.py
index 8c365b02f..5a56c5dc5 100644
--- a/python/flink_agents/plan/agent_plan.py
+++ b/python/flink_agents/plan/agent_plan.py
@@ -33,6 +33,7 @@ from flink_agents.api.skills import (
LOAD_SKILL_TOOL,
Skills,
)
+from flink_agents.api.subagent import SubagentSetup
from flink_agents.api.tools.function_tool import FunctionTool as ApiFunctionTool
from flink_agents.api.tools.tool import Tool
from flink_agents.plan.actions.action import Action
@@ -388,6 +389,26 @@ def _get_resource_providers(
)
_add_skills(all_skills, 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
+ )
+ )
+ 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)
+ )
+ else:
+ msg = (
+ f"AGENT resource '{name}' must be a SubagentSetup or a "
+ f"ResourceDescriptor, but got {type(value).__name__}."
+ )
+ raise TypeError(msg)
+
for resource_type in [
ResourceType.CHAT_MODEL,
ResourceType.CHAT_MODEL_CONNECTION,
diff --git a/python/flink_agents/plan/resource_provider.py b/python/flink_agents/plan/resource_provider.py
index 360e5fda7..ef6fb2f23 100644
--- a/python/flink_agents/plan/resource_provider.py
+++ b/python/flink_agents/plan/resource_provider.py
@@ -242,3 +242,14 @@ class JavaSerializableResourceProvider(SerializableResourceProvider):
"by JavaSerializableResourceProvider in python."
)
raise NotImplementedError(err_msg)
+
+
+def is_python_owned(provider: ResourceProvider) -> bool:
+ """Whether the provider materializes a resource owned by the Python runtime.
+
+ The runtime must ask that runtime to build such a resource instead of
+ resolving it on the Java side. Mirrors Java ``ResourceProvider.isPythonOwned``.
+ """
+ return isinstance(
+ provider, PythonResourceProvider | PythonSerializableResourceProvider
+ )
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
new file mode 100644
index 000000000..818e109c2
--- /dev/null
+++ b/python/flink_agents/plan/tests/test_agent_plan_subagent_resources.py
@@ -0,0 +1,72 @@
+################################################################################
+# 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.
+################################################################################
+"""Tests for compiling AGENT resources (SubagentSetup) into the agent plan."""
+import pytest
+
+from flink_agents.api.agents.agent import Agent
+from flink_agents.api.resource import ResourceDescriptor, ResourceType
+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
+
+
+def test_subagent_setup_compiles_into_agent_provider() -> None:
+ """A registered SubagentSetup lands in the AGENT provider map and resolves."""
+ setup = TestSubagentSetup()
+ agent = Agent()
+ agent.add_resource("reviewer", ResourceType.AGENT, setup)
+
+ plan = AgentPlan.from_agent(agent, AgentConfiguration())
+
+ agents = plan.resource_providers[ResourceType.AGENT]
+ assert agents is not None
+ assert "reviewer" in agents
+ resolved = agents["reviewer"].provide(
+ resource_context=None, config=AgentConfiguration()
+ )
+ assert isinstance(resolved, SubagentSetup)
+ assert resolved.resource_type() == ResourceType.AGENT
+
+
+def test_agent_descriptor_compiles_into_agent_provider() -> None:
+ """Descriptor-shaped AGENT resources (the YAML path) compile into providers."""
+ agent = Agent()
+ agent.add_resource(
+ "summarizer",
+ ResourceType.AGENT,
+ ResourceDescriptor(
+ clazz=f"{TestSubagentSetup.__module__}.{TestSubagentSetup.__name__}",
+ endpoint_url="http://summarizer:8080",
+ ),
+ )
+
+ plan = AgentPlan.from_agent(agent, AgentConfiguration())
+
+ agents = plan.resource_providers[ResourceType.AGENT]
+ assert agents is not None
+ assert "summarizer" in agents
+
+
+def test_non_setup_agent_resource_is_rejected() -> None:
+ """A bare object registered under AGENT fails plan compilation."""
+ agent = Agent()
+ agent.resources[ResourceType.AGENT]["bad"] = object()
+
+ with pytest.raises(TypeError, match="must be a SubagentSetup"):
+ AgentPlan.from_agent(agent, AgentConfiguration())
diff --git a/python/flink_agents/runtime/async_subagent.py b/python/flink_agents/runtime/async_subagent.py
new file mode 100644
index 000000000..585de4919
--- /dev/null
+++ b/python/flink_agents/runtime/async_subagent.py
@@ -0,0 +1,347 @@
+################################################################################
+# 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.
+################################################################################
+"""The async-job execution mode running in durable pub/sub mode."""
+
+import time
+from abc import ABC, abstractmethod
+from concurrent.futures import CancelledError
+from enum import Enum
+from typing import Any
+
+from flink_agents.api.runner_context import RunnerContext
+from flink_agents.api.subagent import (
+ SubagentFuture,
+ SubagentFutures,
+ SubagentResult,
+)
+from flink_agents.runtime.base_subagent import BaseSubagentSetup
+from flink_agents.runtime.subagent_handles import (
+ PendingSubagentCallRegistry,
+ SubagentFutureGroup,
+)
+
+
+class RunStatus:
+ """State snapshot of a remote run reported by the ``call_query_status``
+ probe.
+
+ A state other than ``NOT_STARTED`` means the submission landed on the
+ service, which is the sole basis for ``reconcile_submit_request``
+ deciding between re-posting and polling. The snapshot never carries the
+ result payload.
+ """
+
+ class State(Enum):
+ """Lifecycle of the remote run."""
+
+ NOT_STARTED = "not_started"
+ RUNNING = "running"
+ COMPLETED = "completed"
+ FAILED = "failed"
+
+ def __init__(self, state: "RunStatus.State", error: str | None = None) -> None:
+ """Initialize with the lifecycle state and the optional error."""
+ self._state = state
+ self._error = error
+
+ @staticmethod
+ def not_started() -> "RunStatus":
+ """The service has no record of the run: the POST never landed (or
+ the id mismatches).
+ """
+ return RunStatus(RunStatus.State.NOT_STARTED)
+
+ @staticmethod
+ def running() -> "RunStatus":
+ """The run is in progress."""
+ return RunStatus(RunStatus.State.RUNNING)
+
+ @staticmethod
+ def completed() -> "RunStatus":
+ """The run finished successfully."""
+ return RunStatus(RunStatus.State.COMPLETED)
+
+ @staticmethod
+ def failed(error: str) -> "RunStatus":
+ """The run failed, carrying the error message."""
+ return RunStatus(RunStatus.State.FAILED, error)
+
+ @property
+ def state(self) -> "RunStatus.State":
+ """The lifecycle state of the remote run."""
+ return self._state
+
+ @property
+ def error(self) -> str | None:
+ """The error message of a failed run; None otherwise."""
+ return self._error
+
+
+class AsyncSubagentFuture(SubagentFuture):
+ """The sub side of an async-job invocation.
+
+ The run was already started by the durable POST of ``submit``, so the
+ handle only subscribes to it.
+ """
+
+ def __init__(
+ self,
+ setup: "BaseAsyncSubagentSetup",
+ ctx: RunnerContext,
+ session_id: str,
+ call_id: str,
+ registry: PendingSubagentCallRegistry | None = None,
+ ) -> None:
+ """Initialize with the owning setup, the context, and the identity."""
+ super().__init__(session_id, call_id)
+ self._setup = setup
+ self._ctx = ctx
+ self._registry = registry
+ self._consumed = False
+ self._cancelled = False
+ self._value: SubagentResult | None = None
+ if registry is not None:
+ registry.track_pending_subagent_call(self.identity)
+
+ def done(self) -> bool:
+ """Probe the remote status directly.
+
+ The probe runs outside durable execution, so a failover replay may
+ probe a different number of times than the original execution. A
+ probe failure propagates and fails the action.
+ """
+ if self._consumed or self._cancelled:
+ return True
+ probe = self._setup.query_status(self.session_id, self.call_id)
+ return probe.state in (
+ RunStatus.State.COMPLETED,
+ RunStatus.State.FAILED,
+ )
+
+ def __await__(self) -> Any:
+ """Wait for the run through the durable await composition, releasing
+ the mailbox while waiting.
+
+ A cancelled handle raises :class:`CancelledError`.
+ """
+ if self._cancelled:
+ msg = f"Sub-agent call cancelled: {self.identity}"
+ raise CancelledError(msg)
+ if not self._consumed:
+ # Build the durable awaitable first, then yield from it, so the
+ # await composition and its durable execution cannot be misread
+ # as one serial call.
+ awaitable = self._ctx.durable_execute_async(
+ self._setup._await_until_terminal,
+ self.session_id,
+ self.call_id,
+ durable_id=f"{self.identity}#await",
+ )
+ self._value = yield from awaitable.__await__()
+ self._consumed = True
+ if self._registry is not None:
+ self._registry.untrack_pending_subagent_call(self.identity)
+ return self._value
+
+ def cancel(self) -> None:
+ """Propagate the cancellation through the setup's
+ ``call_cancel_request`` hook.
+
+ The propagation runs synchronously through the hook and is replayed
+ with the enclosing action, so a failover may propagate the same
+ cancellation again. A repeated cancel on the same handle and a
+ cancel after the resolve are local no-ops. A hook failure
+ propagates and fails the action.
+ """
+ if self._consumed or self._cancelled:
+ return
+ self._setup.cancel_request(self._ctx, self.session_id, self.call_id)
+ self._cancelled = True
+ if self._registry is not None:
+ self._registry.untrack_pending_subagent_call(self.identity)
+
+ def combine(self, *others: SubagentFuture) -> SubagentFutures:
+ """Group this handle with others for a batched resolve."""
+ return SubagentFutureGroup((self, *others))
+
+
+class BaseAsyncSubagentSetup(BaseSubagentSetup, ABC):
+ """Production base for sub-agents whose protocol is an asynchronous job,
+ run in durable pub/sub mode.
+
+ ``submit`` publishes the run through one durable POST, the returned
+ handle subscribes to it. The shape matches LangGraph runs, OpenAI
+ Assistants runs, and A2A long-running tasks.
+ """
+
+ #: Delay between status probes while waiting for the run to reach a
+ #: terminal state, declared in YAML as ``status_poll_interval_millis``,
+ #: the same argument the Java side reads from the descriptor. Both default
+ #: to 500, and subclasses override the attribute directly.
+ status_poll_interval_millis: int = 500
+
+ async def submit_with_identity(
+ self,
+ ctx: RunnerContext,
+ prompt: Any,
+ session_id: str,
+ call_id: str,
+ ) -> SubagentFuture:
+ """Start the remote run through the durable POST and return its
+ handle.
+
+ The POST runs through async durable execution and lands before the
+ handle is returned; a POST failure raises and fails the action.
+ """
+ await self.submit_request(ctx, session_id, call_id, prompt)
+ return AsyncSubagentFuture(
+ self, ctx, session_id, call_id, self.pending_call_registry()
+ )
+
+ # --------------------------------------------------------------------------------
+ # Framework wrappers: defaults composing the primitives, overridable
+ # --------------------------------------------------------------------------------
+
+ async def submit_request(
+ self,
+ ctx: RunnerContext,
+ session_id: str,
+ call_id: str,
+ prompt: Any,
+ ) -> None:
+ """Run the durable POST of one invocation. It is the only wrapper
+ wired to a reconciler.
+ """
+ await ctx.durable_execute_async(
+ self._post_submit_request,
+ session_id,
+ call_id,
+ prompt,
+ durable_id=f"{session_id}#{call_id}",
+ reconciler=lambda: self.reconcile_submit_request(
+ session_id, call_id, prompt
+ ),
+ )
+
+ def query_status(self, session_id: str, call_id: str) -> RunStatus:
+ """Probe the remote status. The probe is a direct read-only query,
+ so durable execution does not record it and a failover replay
+ probes again.
+ """
+ return self.call_query_status(session_id, call_id)
+
+ def cancel_request(self, ctx: RunnerContext, session_id: str, call_id: str) -> None:
+ """Propagate the cancellation. The wrapper calls the hook
+ synchronously, so durable execution does not record the
+ propagation and a failover replay propagates it again.
+ """
+ self.call_cancel_request(session_id, call_id)
+
+ def _await_until_terminal(self, session_id: str, call_id: str) -> SubagentResult:
+ """Poll the status until the run reaches a terminal state, then fetch
+ the result.
+
+ The body of the durable await composition keyed by
+ ``session_id#call_id#await``. A probe or fetch failure that escapes
+ the body is a system-level failure: it propagates instead of being
+ folded into an error result.
+ """
+ while True:
+ probe = self.call_query_status(session_id, call_id)
+ if probe.state == RunStatus.State.COMPLETED:
+ return self.call_fetch_result(session_id, call_id)
+ if probe.state == RunStatus.State.FAILED:
+ return SubagentResult.error(probe.error or "run failed")
+ # NOT_STARTED or RUNNING: keep probing. A NOT_STARTED run after a
+ # durable POST means the remote session expired; the replay then
+ # observes the fresh state instead of the original probe path.
+ time.sleep(self.status_poll_interval_millis / 1000)
+
+ def _post_submit_request(self, session_id: str, call_id: str, prompt: Any) -> None:
+ """Run the body of the durable POST by delegating to the transport
+ primitive.
+ """
+ self.call_submit_request(session_id, call_id, prompt)
+
+ # --------------------------------------------------------------------------------
+ # Transport primitives provided by the integration
+ # --------------------------------------------------------------------------------
+
+ @abstractmethod
+ def call_submit_request(self, session_id: str, call_id: str, prompt: Any) -> None:
+ """Start the run remotely. A raised exception fails the enclosing
+ action.
+ """
+
+ @abstractmethod
+ def call_query_status(self, session_id: str, call_id: str) -> RunStatus:
+ """Probe the run's current state read-only; must not alter the remote
+ run.
+
+ The status never carries the result payload — the result is fetched
+ separately through :meth:`call_fetch_result`.
+
+ Implementations must report comprehensible failures (an expired
+ endpoint, expired credentials, a rejected run) as a FAILED status
+ rather than raising; an exception escaping this probe is treated as
+ a system-level failure, propagates, and triggers a job failover.
+ """
+
+ @abstractmethod
+ def call_fetch_result(self, session_id: str, call_id: str) -> SubagentResult:
+ """Fetch the result of a run that reached a terminal state;
+ comprehensible failures go into the :class:`SubagentResult`, while a
+ raised exception is a system-level failure that propagates.
+
+ The fetch must be an idempotent read: a failover re-executes it when
+ the crash hit the fetch in flight.
+ """
+
+ def reconcile_submit_request(
+ self, session_id: str, call_id: str, prompt: Any
+ ) -> None:
+ """The crash-window recovery of the POST: probes the status and
+ handles every state explicitly, so a landed POST is never
+ duplicated. A probe failure propagates and fails the recovery.
+ """
+ probe = self.call_query_status(session_id, call_id)
+ state = probe.state
+ if state == RunStatus.State.NOT_STARTED:
+ # The service has no record of the run: the POST never landed.
+ self.call_submit_request(session_id, call_id, prompt)
+ elif state == RunStatus.State.RUNNING:
+ # The POST landed and the run is in flight; the resolve keeps
+ # polling it. Nothing to repair.
+ pass
+ elif state in (RunStatus.State.COMPLETED, RunStatus.State.FAILED):
+ # The run reached a terminal state while the caller was down;
+ # the resolve picks up the outcome — the fetch or the reported
+ # error. Nothing to repair.
+ pass
+ else:
+ # Fail loudly instead of silently skipping an unknown state.
+ msg = f"Unknown run state: {state}"
+ raise ValueError(msg)
+
+ def call_cancel_request(self, session_id: str, call_id: str) -> None:
+ """Hook propagating a cancellation to the remote run. The default is
+ a no-op.
+
+ A replay may propagate the cancellation again, so remote
+ cancellations must be idempotent.
+ """
diff --git a/python/flink_agents/runtime/base_subagent.py b/python/flink_agents/runtime/base_subagent.py
new file mode 100644
index 000000000..a19655dcf
--- /dev/null
+++ b/python/flink_agents/runtime/base_subagent.py
@@ -0,0 +1,294 @@
+################################################################################
+# 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.
+################################################################################
+"""The framework-level runtime base shared by every sub-agent execution mode."""
+
+import hashlib
+import json
+import uuid
+from abc import ABC, abstractmethod
+from dataclasses import dataclass, replace
+from typing import Any
+
+from pydantic import PrivateAttr
+
+from flink_agents.api.runner_context import RunnerContext
+from flink_agents.api.subagent import SubagentFuture, SubagentSetup
+from flink_agents.runtime.subagent_handles import PendingSubagentCallRegistry
+from flink_agents.runtime.task_lifecycle_listener import TaskLifecycleListener
+
+
+def _event_attributes(event: Any) -> dict[str, Any]:
+ """Normalize an event's attributes into a plain dict.
+
+ Accepts both a Java ``Event`` reference passed across the bridge and a
+ plain Python mapping, copying Java maps entry by entry.
+ """
+ attributes = event.getAttributes()
+ if attributes is None:
+ return {}
+ if isinstance(attributes, dict):
+ return dict(attributes)
+ try:
+ return {str(k): v for k, v in attributes.entrySet()}
+ except AttributeError:
+ return {str(k): v for k, v in dict(attributes).items()}
+
+
+@dataclass(frozen=True)
+class Namespace:
+ """The caller-side identity of one action task execution.
+
+ Provides the task identity keying the runtime bookkeeping and the
+ namespace digest seeding the deterministic ids of the sub-agent
+ calls the task issues.
+
+ Key, sequence number, action name, and the event's type and
+ attributes are facts of the execution itself, identical for every
+ sub-agent called from it. The subagent name distinguishes the
+ sub-agents called from one action, so it alone keeps their id
+ ranges apart.
+ """
+
+ key: str
+ sequence_number: int
+ action_name: str
+ event_type: str
+ event_attributes: dict[str, Any]
+ event_id: str
+ subagent_name: str = ""
+
+ @staticmethod
+ def from_task(task: Any) -> "Namespace":
+ """Extract the facts from an ``ActionTask`` reference or fake."""
+ return Namespace(
+ key=str(task.getKey()),
+ sequence_number=int(task.getSequenceNumber()),
+ action_name=str(task.getAction().getName()),
+ event_type=str(task.getEvent().getType()),
+ event_attributes=_event_attributes(task.getEvent()),
+ event_id=str(task.getEvent().getId()),
+ )
+
+ @property
+ def task_identity(self) -> str:
+ """A key unique among live task executions and stable across the
+ steps of one task.
+ """
+ return f"{self.key}#{self.sequence_number}#{self.action_name}#{self.event_id}"
+
+ def namespace_digest(self) -> str:
+ """Digest the id-bearing facts into a name-based UUID string.
+
+ The ids are reproducible across a failover replay. The event id
+ stays out of the digest: it keys the runtime bookkeeping only.
+ """
+ fields = {
+ "actionName": self.action_name,
+ "eventAttributes": self.event_attributes,
+ "eventType": self.event_type,
+ "key": self.key,
+ "sequenceNumber": self.sequence_number,
+ "subagentName": self.subagent_name,
+ }
+ payload = json.dumps(
+ fields, sort_keys=True, separators=(",", ":"), default=str
+ ).encode("utf-8")
+ # MD5 with the version/variant bits, as in Java's
+ # UUID.nameUUIDFromBytes (a version 3 UUID).
+ digest = bytearray(hashlib.md5(payload).digest())
+ digest[6] = (digest[6] & 0x0F) | 0x30
+ digest[8] = (digest[8] & 0x3F) | 0x80
+ return str(uuid.UUID(bytes=bytes(digest)))
+
+
+class SubagentIdAllocator:
+ """Deterministic ``(session_id, call_id)`` source for one task execution.
+
+ The namespace digest fixes the counting range, so a failover replay
+ of the same task hands out the same ids in the same call order.
+ """
+
+ def __init__(self, namespace: Namespace) -> None:
+ """Create an allocator over one task's namespace."""
+ self._namespace = namespace
+ self._session_ordinal = 0
+ self._per_session_call_ordinals: dict[str, int] = {}
+
+ def next_session_id(self) -> str:
+ """Create a session id scoped to this task's namespace."""
+ ordinal = self._session_ordinal
+ self._session_ordinal += 1
+ return f"{self._namespace.namespace_digest()}-{ordinal}"
+
+ def next_call_id(self, session_id: str) -> str:
+ """Create a call id by appending the per-session ordinal."""
+ ordinal = self._per_session_call_ordinals.get(session_id, 0) + 1
+ self._per_session_call_ordinals[session_id] = ordinal
+ return f"{session_id}-{ordinal}"
+
+
+class BaseSubagentSetup(SubagentSetup, TaskLifecycleListener, ABC):
+ """Runtime base for sub-agent setups, holding the per-task id allocators
+ and pending-call registries keyed to the currently executing action task.
+ How an invocation is issued stays an execution mode owned by the concrete
+ subclass.
+ """
+
+ _per_task_allocators: dict[str, SubagentIdAllocator] = PrivateAttr(
+ default_factory=dict
+ )
+ _per_task_registries: dict[str, PendingSubagentCallRegistry] = PrivateAttr(
+ default_factory=dict
+ )
+ _current_namespace: Namespace | None = PrivateAttr(default=None)
+ _subagent_name: str | None = PrivateAttr(default=None)
+
+ # --------------------------------------------------------------------------------
+ # Task lifecycle hooks (keyword-invoked by the runtime bridge)
+ # --------------------------------------------------------------------------------
+
+ def on_action_prepared(self, task: Any) -> None:
+ """Record the task whose execution is currently issuing calls."""
+ namespace = Namespace.from_task(task)
+ self._current_namespace = replace(
+ namespace, subagent_name=self._subagent_name or ""
+ )
+
+ def on_action_transferred(self, from_task: Any, to_task: Any) -> None:
+ """Move the finishing task's bookkeeping onto the generated task."""
+ from_identity = Namespace.from_task(from_task).task_identity
+ to_identity = Namespace.from_task(to_task).task_identity
+ allocator = self._per_task_allocators.pop(from_identity, None)
+ if allocator is not None:
+ self._per_task_allocators[to_identity] = allocator
+ registry = self._per_task_registries.pop(from_identity, None)
+ if registry is not None:
+ registry.set_action_name(
+ Namespace.from_task(to_task).action_name
+ )
+ self._per_task_registries[to_identity] = registry
+
+ def on_action_finishing(self, task: Any) -> None:
+ """Drop the task's bookkeeping and enforce resolved handles.
+
+ The replay-reuse path reaches the same finalization through
+ ``on_action_reused``, keeping the prepared/terminal pairing intact on
+ both paths. A failed invocation intentionally skips this cleanup: the
+ failure fails the run and the task is replayed on the restarted
+ operator, so stale entries cannot outlive the run.
+ """
+ self._current_namespace = None
+ identity = Namespace.from_task(task).task_identity
+ self._per_task_allocators.pop(identity, None)
+ registry = self._per_task_registries.pop(identity, None)
+ if registry is not None:
+ registry.check_empty()
+
+ def on_action_reused(self, task: Any) -> None:
+ """Reuse is a terminal outcome like finishing, so share finalization."""
+ self.on_action_finishing(task)
+
+ # --------------------------------------------------------------------------------
+ # Identity injected by the framework
+ # --------------------------------------------------------------------------------
+
+ def set_subagent_name(self, subagent_name: str) -> None:
+ """Record the resource name the framework injects as the subagent name."""
+ self._subagent_name = subagent_name
+
+ @property
+ def subagent_name(self) -> str | None:
+ """The injected subagent name, or None outside the framework."""
+ return self._subagent_name
+
+ # --------------------------------------------------------------------------------
+ # Submit dispatch: complete missing ids, then delegate to the mode
+ # --------------------------------------------------------------------------------
+
+ async def submit(
+ self,
+ ctx: RunnerContext,
+ prompt: Any,
+ session_id: str | None = None,
+ call_id: str | None = None,
+ ) -> SubagentFuture:
+ """Issue an invocation, assigning the missing ids deterministically.
+
+ The ids are assigned when this call is awaited rather than when it is
+ made, so a replay awaiting the invocations in the same order hands
+ out the same ids.
+ """
+ if session_id is None or call_id is None:
+ allocator = self._current_allocator()
+ if session_id is None:
+ session_id = allocator.next_session_id()
+ if call_id is None:
+ call_id = allocator.next_call_id(session_id)
+ return await self.submit_with_identity(ctx, prompt, session_id, call_id)
+
+ @abstractmethod
+ async def submit_with_identity(
+ self,
+ ctx: RunnerContext,
+ prompt: Any,
+ session_id: str,
+ call_id: str,
+ ) -> SubagentFuture:
+ """Issue one invocation under the fully assigned identity.
+
+ The execution-mode hook implementing how the invocation is issued.
+ """
+
+ # --------------------------------------------------------------------------------
+ # Per-task bookkeeping
+ # --------------------------------------------------------------------------------
+
+ def pending_call_registry(self) -> PendingSubagentCallRegistry | None:
+ """The registry of the currently executing task.
+
+ Handles record themselves there on creation. Returns None outside a
+ prepared task, so calls issued without a task context skip tracking.
+ """
+ return self._current_task_registry()
+
+ def _current_task_registry(self) -> PendingSubagentCallRegistry | None:
+ if self._current_namespace is None:
+ return None
+ identity = self._current_namespace.task_identity
+ registry = self._per_task_registries.get(identity)
+ if registry is None:
+ registry = PendingSubagentCallRegistry(
+ self._current_namespace.action_name
+ )
+ self._per_task_registries[identity] = registry
+ return registry
+
+ def _current_allocator(self) -> SubagentIdAllocator:
+ """The allocator of the executing task, scoped to one action
+ execution so ordinals restart for the next action. Replays hand
+ out the same ids.
+ """
+ if self._current_namespace is None:
+ msg = "No prepared action task to assign sub-agent ids from."
+ raise RuntimeError(msg)
+ namespace = self._current_namespace
+ allocator = self._per_task_allocators.get(namespace.task_identity)
+ if allocator is None:
+ allocator = SubagentIdAllocator(namespace)
+ self._per_task_allocators[namespace.task_identity] = allocator
+ return allocator
diff --git a/python/flink_agents/runtime/deferred_subagent.py b/python/flink_agents/runtime/deferred_subagent.py
new file mode 100644
index 000000000..68024b604
--- /dev/null
+++ b/python/flink_agents/runtime/deferred_subagent.py
@@ -0,0 +1,180 @@
+################################################################################
+# 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.
+################################################################################
+"""The framework-level deferred execution mode for sub-agent setups."""
+
+from abc import ABC, abstractmethod
+from concurrent.futures import CancelledError
+from typing import Any, Callable
+
+from flink_agents.api.runner_context import RunnerContext
+from flink_agents.api.subagent import (
+ SubagentFuture,
+ SubagentFutures,
+ SubagentResult,
+)
+from flink_agents.runtime.base_subagent import BaseSubagentSetup
+from flink_agents.runtime.subagent_handles import (
+ PendingSubagentCallRegistry,
+ SubagentFutureGroup,
+)
+
+#: The (durable id, call, reconcile) triple returned by
+#: :meth:`DeferredSubagentSetup.prepare`.
+PreparedTriple = tuple[Any, Any, Any]
+
+
+class DeferredSubagentFuture(SubagentFuture):
+ """Deferred handle to one sub-agent invocation."""
+
+ def __init__(
+ self,
+ session_id: str,
+ call_id: str,
+ ctx: RunnerContext,
+ prepared_factory: Callable[[], PreparedTriple],
+ registry: PendingSubagentCallRegistry | None = None,
+ ) -> None:
+ """Initialize with the identity and the factory preparing the call."""
+ super().__init__(session_id, call_id)
+ self._ctx = ctx
+ self._prepared_factory = prepared_factory
+ self._registry = registry
+ self._prepared: PreparedTriple | None = None
+ self._done = False
+ self._cancelled = False
+ self._value: SubagentResult | None = None
+ if registry is not None:
+ registry.track_pending_subagent_call(self.identity)
+
+ def done(self) -> bool:
+ """Whether the invocation has been resolved or cancelled."""
+ return self._done or self._cancelled
+
+ def cancel(self) -> None:
+ """Cancel before the request is prepared: the request is discarded.
+
+ Resolving a cancelled handle raises :class:`CancelledError`. An
+ already resolved handle ignores the cancellation request.
+ """
+ if self._done:
+ return
+ self._cancelled = True
+ if self._registry is not None:
+ self._registry.untrack_pending_subagent_call(self.identity)
+
+ def prepare(self) -> PreparedTriple:
+ """Prepare the request if it has not been prepared yet.
+
+ Mailbox-confined: must run on the mailbox thread.
+ """
+ if self._cancelled:
+ msg = f"Sub-agent call cancelled: {self.identity}"
+ raise CancelledError(msg)
+ if self._prepared is None:
+ self._prepared = self._prepared_factory()
+ return self._prepared
+
+ def execute(self) -> Any:
+ """Run the prepared request through durable execution and record
+ the outcome; awaitable, releasing the mailbox while waiting.
+
+ A system-level failure escaping durable execution propagates and fails
+ the action instead of being folded into an error result.
+ """
+ durable_id, call, reconcile = self.prepare()
+ value = yield from self._ctx.durable_execute_async(
+ call,
+ reconciler=reconcile,
+ durable_id=durable_id,
+ ).__await__()
+ self._resolve(value)
+
+ def combine(self, *others: SubagentFuture) -> SubagentFutures:
+ """Group this handle with others for a batched resolve."""
+ return SubagentFutureGroup((self, *others))
+
+ def __await__(self) -> Any:
+ """Resolve the invocation, releasing the mailbox while waiting."""
+ if self._cancelled:
+ msg = f"Sub-agent call cancelled: {self.identity}"
+ raise CancelledError(msg)
+ if not self._done:
+ yield from self.execute()
+ return self._value
+
+ def _resolve(self, value: SubagentResult) -> None:
+ self._value = value
+ self._done = True
+ if self._registry is not None:
+ self._registry.untrack_pending_subagent_call(self.identity)
+
+
+class DeferredSubagentSetup(BaseSubagentSetup, ABC):
+ """The framework-level deferred execution mode for sub-agent setups.
+
+ ``submit`` registers the invocation and returns a deferred handle
+ without sending anything; the actual request is issued lazily when
+ the handle is first resolved, and runs through one durable async
+ callable keyed by a failover-reproducible id, so the invocation
+ participates in the task's durable execution.
+ """
+
+ async def submit_with_identity(
+ self,
+ ctx: RunnerContext,
+ prompt: Any,
+ session_id: str,
+ call_id: str,
+ ) -> SubagentFuture:
+ """Register the invocation under the given identity and return its handle."""
+ return DeferredSubagentFuture(
+ session_id,
+ call_id,
+ ctx,
+ prepared_factory=lambda: self.prepare(ctx, prompt, session_id, call_id),
+ registry=self.pending_call_registry(),
+ )
+
+ @abstractmethod
+ def prepare(
+ self,
+ ctx: RunnerContext,
+ prompt: Any,
+ session_id: str,
+ call_id: str,
+ ) -> PreparedTriple:
+ """Prepare one invocation and return its ``(id, call, reconcile)``
+ triple; ids are supplied.
+
+ The durable id MUST be derived solely from the
+ ``(session_id, call_id)`` pair so it is reproducible after
+ failover.
+
+ Called exactly once per invocation, when the deferred handle is
+ first resolved, on the mailbox thread; implementations may
+ therefore perform the mailbox-confined part of issuing the request
+ here, leaving only the off-mailbox part in the returned call.
+
+ The returned call folds its own comprehensible failures into the
+ :class:`SubagentResult` it returns; an exception escaping the call
+ is a system-level failure that propagates and fails the action.
+
+ Skipping ``reconcile`` has a cost: a crash between the call landing
+ and its result being persisted re-invokes the sub-agent on replay,
+ possibly duplicating external side effects.
+ """
diff --git a/python/flink_agents/runtime/durable_execution.py b/python/flink_agents/runtime/durable_execution.py
index a85d704bc..a862d02e8 100644
--- a/python/flink_agents/runtime/durable_execution.py
+++ b/python/flink_agents/runtime/durable_execution.py
@@ -15,12 +15,40 @@
# See the License for the specific language governing permissions and
# limitations under the License.
#################################################################################
+import functools
import hashlib
import inspect
from typing import Any, Callable
import cloudpickle
+_DURABLE_ID_ATTR = "__flink_agents_durable_id__"
+
+
+def with_durable_id(func: Callable, durable_id: str) -> Callable:
+ """Wrap ``func`` so durable execution keys it by ``durable_id``.
+
+ Callers that own a stable identity for a durable call attach it here
+ instead of relying on the module/qualname of the callable, which is
+ shared by every call issued from the same implementation.
+ """
+
+ @functools.wraps(func)
+ def wrapped(*args: Any, **kwargs: Any) -> Any:
+ return func(*args, **kwargs)
+
+ setattr(wrapped, _DURABLE_ID_ATTR, durable_id)
+ return wrapped
+
+
+def get_durable_id(func: Callable) -> str | None:
+ """Return the explicit durable id attached by :func:`with_durable_id`.
+
+ Returns ``None`` when the callable carries no explicit id, in which case
+ callers fall back to deriving the identity from the callable itself.
+ """
+ return getattr(func, _DURABLE_ID_ATTR, None)
+
def durable_identity_for_call(
func: Callable,
@@ -33,7 +61,14 @@ def durable_identity_for_call(
def _compute_function_id(func: Callable) -> str:
- """Compute a stable function identifier from a callable."""
+ """Compute a stable function identifier from a callable.
+
+ An explicit id attached by :func:`with_durable_id` wins over the derived
+ module/qualname.
+ """
+ explicit_id = get_durable_id(func)
+ if explicit_id is not None:
+ return explicit_id
module_obj = inspect.getmodule(func)
module = (
module_obj.__name__
diff --git a/python/flink_agents/runtime/flink_runner_context.py b/python/flink_agents/runtime/flink_runner_context.py
index cf44c10be..4e9fefa88 100644
--- a/python/flink_agents/runtime/flink_runner_context.py
+++ b/python/flink_agents/runtime/flink_runner_context.py
@@ -51,6 +51,7 @@ from flink_agents.runtime.durable_execution import (
_compute_function_id,
_validate_reconciler_callable,
durable_identity_for_call,
+ with_durable_id,
)
from flink_agents.runtime.flink_memory_object import FlinkMemoryObject
from flink_agents.runtime.flink_metric_group import FlinkMetricGroup
@@ -65,6 +66,7 @@ from flink_agents.runtime.resource_cache import (
_failure_of,
_first_or_logged,
)
+from flink_agents.runtime.task_lifecycle_listener import TaskLifecycleListener
logger = logging.getLogger(__name__)
@@ -485,6 +487,10 @@ class FlinkRunnerContext(RunnerContext, ExecutionReporter):
self.__resource_cache.set_java_resource_adapter(j_resource_adapter)
self.__config = self.__agent_plan.config
self.executor = executor
+ # Task lifecycle listeners the operator's callbacks fan out to,
+ # registered via add_task_lifecycle_listener() (aligned with the Java
+ # operator's taskLifecycleListeners).
+ self.__task_lifecycle_listeners: list = []
def set_long_term_memory(self, ltm: InternalBaseLongTermMemory) -> None:
"""Set long term memory instance to this context.
@@ -530,6 +536,75 @@ class FlinkRunnerContext(RunnerContext, ExecutionReporter):
resource.set_metric_group(metric_group or self.action_metric_group)
return resource
+ def eager_materialize(self, resource_type: str) -> Dict[str, Resource]:
+ """Materialize every Python-owned resource of ``resource_type``.
+
+ The Python-side counterpart of the Java ``ResourceCache.eagerMaterialize``:
+ resources declared by Python providers are built, cached and closed here,
+ so the Java side asks for them instead of building its own. Returns them
+ keyed by resource name.
+ """
+ from flink_agents.plan.resource_provider import is_python_owned
+
+ type_ = ResourceType(resource_type)
+ materialized = {}
+ providers = self.__agent_plan.resource_providers.get(type_, {})
+ for name, provider in providers.items():
+ if not is_python_owned(provider):
+ # Java-owned resources are materialized by the Java resource cache.
+ continue
+ materialized[name] = self.__resource_cache.get_resource(name, type_)
+ return materialized
+
+ def add_task_lifecycle_listener(self, listener: Any) -> None:
+ """Register a task lifecycle listener the operator's callbacks fan out to."""
+ self.__task_lifecycle_listeners.append(listener)
+
+ def notify_record_start(self, key: Any) -> None:
+ """Fan out the operator's onRecordStart to the task lifecycle listeners."""
+ for listener in self.__task_lifecycle_listeners:
+ listener.on_record_start(key)
+
+ def notify_action_prepared(self, task: Any) -> None:
+ """Fan out the operator's onActionPrepared to the task lifecycle listeners."""
+ for listener in self.__task_lifecycle_listeners:
+ listener.on_action_prepared(task)
+
+ def notify_action_started(self, task: Any) -> None:
+ """Fan out the operator's onActionStarted to the task lifecycle listeners."""
+ for listener in self.__task_lifecycle_listeners:
+ listener.on_action_started(task)
+
+ def notify_action_transferred(self, from_task: Any, to_task: Any) -> None:
+ """Fan out the operator's onActionTransferred to the listeners."""
+ for listener in self.__task_lifecycle_listeners:
+ listener.on_action_transferred(from_task, to_task)
+
+ def notify_action_finishing(self, task: Any) -> None:
+ """Fan out the operator's onActionFinishing to the task lifecycle listeners."""
+ for listener in self.__task_lifecycle_listeners:
+ listener.on_action_finishing(task)
+
+ def notify_action_finished(self, task: Any) -> None:
+ """Fan out the operator's onActionFinished to the listeners."""
+ for listener in self.__task_lifecycle_listeners:
+ listener.on_action_finished(task)
+
+ def notify_action_reused(self, task: Any) -> None:
+ """Fan out the operator's onActionReused to the task lifecycle listeners."""
+ for listener in self.__task_lifecycle_listeners:
+ listener.on_action_reused(task)
+
+ def notify_action_failed(self, task: Any, error: Any) -> None:
+ """Fan out the operator's onActionFailed to the task lifecycle listeners."""
+ for listener in self.__task_lifecycle_listeners:
+ listener.on_action_failed(task, error)
+
+ def notify_record_finished(self, key: Any) -> None:
+ """Fan out the operator's onRecordFinished to the task lifecycle listeners."""
+ for listener in self.__task_lifecycle_listeners:
+ listener.on_record_finished(key)
+
@property
@override
def action_config(self) -> Dict[str, Any]:
@@ -1113,6 +1188,7 @@ class FlinkRunnerContext(RunnerContext, ExecutionReporter):
func: Callable[[Any], Any],
*args: Any,
reconciler: Callable[[], Any] | None = None,
+ durable_id: str | None = None,
**kwargs: Any,
) -> Any:
"""Synchronously execute the provided function with durable execution support.
@@ -1126,6 +1202,8 @@ class FlinkRunnerContext(RunnerContext, ExecutionReporter):
the operator until completion.
"""
validated_reconciler = _validate_reconciler_callable(reconciler)
+ if durable_id is not None:
+ func = with_durable_id(func, durable_id)
if validated_reconciler is not None:
plan = self._plan_reconciler_execution(
@@ -1153,6 +1231,7 @@ class FlinkRunnerContext(RunnerContext, ExecutionReporter):
func: Callable[[Any], Any],
*args: Any,
reconciler: Callable[[], Any] | None = None,
+ durable_id: str | None = None,
**kwargs: Any,
) -> AsyncExecutionResult:
"""Asynchronously execute the provided function with durable execution support.
@@ -1167,6 +1246,8 @@ class FlinkRunnerContext(RunnerContext, ExecutionReporter):
recorded and cannot be recovered.
"""
validated_reconciler = _validate_reconciler_callable(reconciler)
+ if durable_id is not None:
+ func = with_durable_id(func, durable_id)
if validated_reconciler is not None:
return _ReconcilerDurableAsyncExecutionResult(
@@ -1295,6 +1376,72 @@ def close_flink_runner_context(
ctx.close()
+def eager_materialize(
+ ctx: FlinkRunnerContext, resource_type: str
+) -> Dict[str, Resource]:
+ """Java entry: materialize the Python-owned resources of ``resource_type``."""
+ return ctx.eager_materialize(resource_type)
+
+
+def add_task_lifecycle_listener(ctx: FlinkRunnerContext, listener: Any) -> bool:
+ """Java entry: register a Python object as a task lifecycle listener.
+
+ Returns whether it observes the lifecycle, so the Java side knows whether the
+ Python runtime has anything to be notified about.
+ """
+ if not isinstance(listener, TaskLifecycleListener):
+ return False
+ ctx.add_task_lifecycle_listener(listener)
+ return True
+
+
+def notify_record_start(ctx: FlinkRunnerContext, key: Any) -> None:
+ """Java entry: forward onRecordStart to the Python task lifecycle listeners."""
+ ctx.notify_record_start(key)
+
+
+def notify_action_prepared(ctx: FlinkRunnerContext, task: Any) -> None:
+ """Java entry: forward onActionPrepared to the Python task lifecycle listeners."""
+ ctx.notify_action_prepared(task)
+
+
+def notify_action_started(ctx: FlinkRunnerContext, task: Any) -> None:
+ """Java entry: forward onActionStarted to the Python task lifecycle listeners."""
+ ctx.notify_action_started(task)
+
+
+def notify_action_transferred(
+ ctx: FlinkRunnerContext, from_task: Any, to_task: Any
+) -> None:
+ """Java entry: forward onActionTransferred to the Python listeners."""
+ ctx.notify_action_transferred(from_task, to_task)
+
+
+def notify_action_finishing(ctx: FlinkRunnerContext, task: Any) -> None:
+ """Java entry: forward onActionFinishing to the Python task lifecycle listeners."""
+ ctx.notify_action_finishing(task)
+
+
+def notify_action_finished(ctx: FlinkRunnerContext, task: Any) -> None:
+ """Java entry: forward onActionFinished to the Python listeners."""
+ ctx.notify_action_finished(task)
+
+
+def notify_action_reused(ctx: FlinkRunnerContext, task: Any) -> None:
+ """Java entry: forward onActionReused to the Python task lifecycle listeners."""
+ ctx.notify_action_reused(task)
+
+
+def notify_action_failed(ctx: FlinkRunnerContext, task: Any, error: Any) -> None:
+ """Java entry: forward onActionFailed to the Python task lifecycle listeners."""
+ ctx.notify_action_failed(task, error)
+
+
+def notify_record_finished(ctx: FlinkRunnerContext, key: Any) -> None:
+ """Java entry: forward onRecordFinished to the Python task lifecycle listeners."""
+ ctx.notify_record_finished(key)
+
+
_ASYNC_POOL_ID = itertools.count(1)
"""Process-unique pool ids keeping multiple async executors distinguishable."""
diff --git a/python/flink_agents/runtime/resource_cache.py b/python/flink_agents/runtime/resource_cache.py
index 9c96636b4..2e94b0cae 100644
--- a/python/flink_agents/runtime/resource_cache.py
+++ b/python/flink_agents/runtime/resource_cache.py
@@ -123,6 +123,13 @@ class ResourceCache:
)
if isinstance(resource, FunctionTool) and isinstance(resource.func, JavaFunction):
resource.set_java_resource_adapter(self._j_resource_adapter)
+ # Local import avoids pulling sub-agent machinery for non-sub-agent usage.
+ from flink_agents.runtime.base_subagent import BaseSubagentSetup
+
+ if isinstance(resource, BaseSubagentSetup):
+ # The framework owns the setup's identity: inject the resource name
+ # as its sub-agent name, mirroring the Java ResourceCache.
+ resource.set_subagent_name(name)
resource.open()
self._cache.setdefault(type, {})[name] = resource
return resource
diff --git a/python/flink_agents/runtime/subagent_handles.py b/python/flink_agents/runtime/subagent_handles.py
new file mode 100644
index 000000000..07d4f91e9
--- /dev/null
+++ b/python/flink_agents/runtime/subagent_handles.py
@@ -0,0 +1,144 @@
+################################################################################
+# 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.
+################################################################################
+"""Framework-level handle utilities for sub-agent setups.
+
+These companions are execution-mode agnostic. The execution modes build
+on them, so the framework never needs to know how a request is issued.
+"""
+
+from typing import Any
+
+from flink_agents.api.subagent import (
+ SubagentFuture,
+ SubagentFutures,
+ SubagentResult,
+)
+
+
+class PendingSubagentCallRegistry:
+ """The per-action-execution set of sub-agent handles submitted but
+ not yet resolved.
+ """
+
+ def __init__(self, action_name: str) -> None:
+ """Initialize an empty registry for the given action."""
+ self._action_name = action_name
+ self._pending_calls: list[str] = []
+
+ def set_action_name(self, action_name: str) -> None:
+ """Adopt the continuation's action when the execution moves onto
+ another task.
+ """
+ self._action_name = action_name
+
+ def track_pending_subagent_call(self, call_identity: str) -> None:
+ """Record a pending handle. Duplicate identities collapse to one."""
+ if call_identity not in self._pending_calls:
+ self._pending_calls.append(call_identity)
+
+ def untrack_pending_subagent_call(self, call_identity: str) -> None:
+ """Drop a resolved handle and do nothing when the identity is unknown."""
+ if call_identity in self._pending_calls:
+ self._pending_calls.remove(call_identity)
+
+ def is_empty(self) -> bool:
+ """Whether no handle is pending."""
+ return not self._pending_calls
+
+ def check_empty(self) -> None:
+ """Fail when the finished action left a handle unresolved."""
+ if self._pending_calls:
+ msg = (
+ f"Action {self._action_name} finished without resolving the "
+ f"sub-agent calls it submitted: {self._pending_calls}. "
+ f"Resolve every handle returned by submit(), individually "
+ f"or through SubagentFutures."
+ )
+ raise RuntimeError(msg)
+
+
+class CompletedSubagentFuture(SubagentFuture):
+ """A handle for an invocation that has already produced ``value``."""
+
+ def __init__(self, session_id: str, call_id: str, value: SubagentResult) -> None:
+ """Initialize with the identity and the produced value."""
+ super().__init__(session_id, call_id)
+ self._value = value
+
+ def done(self) -> bool:
+ """The invocation has already reached its terminal state."""
+ return True
+
+ def combine(self, *others: SubagentFuture) -> SubagentFutures:
+ """Group this handle with others to be resolved together."""
+ return SubagentFutureGroup((self, *others))
+
+ def __await__(self) -> Any:
+ """Resolve immediately with the produced value."""
+ return self._value
+ yield # pragma: no cover - makes this a generator function
+
+
+class SubagentFutureGroup(SubagentFutures):
+ """The :class:`SubagentFutures` returned by ``combine``: several
+ handles held together.
+ """
+
+ def __init__(self, futures: tuple) -> None:
+ """Initialize with the handles to resolve together."""
+ self._futures = tuple(futures)
+
+ def done(self) -> bool:
+ """Whether every handle in the group has been resolved."""
+ return all(future.done() for future in self._futures)
+
+ def cancel(self) -> None:
+ """Propagate the cancellation request to every handle in the group."""
+ for future in self._futures:
+ future.cancel()
+
+ def combine(self, *others: SubagentFuture) -> SubagentFutures:
+ """Add more handles to the group."""
+ return SubagentFutureGroup((*self._futures, *others))
+
+ def __await__(self) -> Any:
+ """Wait for every handle in submission order.
+
+ Pending deferred handles are prepared before any execution
+ starts, then executed one by one.
+ """
+ # Late import: deferred handles build on this module.
+ from flink_agents.runtime.deferred_subagent import DeferredSubagentFuture
+
+ pending = [
+ future
+ for future in self._futures
+ if isinstance(future, DeferredSubagentFuture) and not future.done()
+ ]
+ for future in pending:
+ future.prepare()
+ # TODO(#926): execute the prepared calls as one batch once durable
+ # execution supports batched submission; until then the prepared
+ # calls are executed one by one.
+ for future in pending:
+ yield from future.execute()
+ outcomes = []
+ for future in self._futures:
+ outcome = yield from future.__await__()
+ outcomes.append(outcome)
+ return outcomes
diff --git a/python/flink_agents/runtime/task_lifecycle_listener.py b/python/flink_agents/runtime/task_lifecycle_listener.py
new file mode 100644
index 000000000..5f7c1ae1f
--- /dev/null
+++ b/python/flink_agents/runtime/task_lifecycle_listener.py
@@ -0,0 +1,96 @@
+################################################################################
+# 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.
+################################################################################
+"""Python mirror of the Java ``TaskLifecycleListener``."""
+
+from typing import Any
+
+
+class TaskLifecycleListener:
+ """Observes the per-record and per-action lifecycle of the action operator.
+
+ Mirrors the Java ``TaskLifecycleListener``: the operator drives Python
+ actions on the JVM and forwards each lifecycle callback here over pemja.
+ Every callback defaults to a no-op, so an implementation overrides only
+ the ones it cares about.
+
+ Pairing semantics: ``on_action_prepared`` pairs with exactly one terminal
+ callback -- ``on_action_finishing`` followed by ``on_action_finished`` on
+ normal completion, ``on_action_reused`` when a replay skips an
+ already-completed action, ``on_action_failed`` on invocation failure, or
+ ``on_action_transferred`` when a non-finished task hands its context over
+ to the task it generated. ``on_action_started`` fires at most once per
+ action execution; the gate is checkpointed with the task, so a failover
+ replay re-emits ``on_record_start`` but not ``on_action_started``.
+
+ Exception contract: the framework allows a listener to inspect state and
+ raise when necessary; a listener that only observes should avoid raising.
+ """
+
+ def on_record_start(self, key: Any) -> None:
+ """The first task of an input record is about to be prepared.
+
+ Also re-emitted when an in-flight record resumes after a failover,
+ so listeners observe a paired bracket for the replayed round.
+ """
+
+ def on_action_prepared(self, task: Any) -> None:
+ """A task's context is wired up and it is ready to run.
+
+ Fires on every preparation, including re-preparation of a suspended
+ or resumed task.
+ """
+
+ def on_action_started(self, task: Any) -> None:
+ """An action execution is about to run for the first time."""
+
+ def on_action_transferred(self, from_task: Any, to_task: Any) -> None:
+ """A non-finished task handed its context to the task it generated."""
+
+ def on_action_finishing(self, task: Any) -> None:
+ """A task completed but its result is not persisted yet.
+
+ Fires immediately before the result is persisted.
+ """
+
+ def on_action_finished(self, task: Any) -> None:
+ """A task's invocation finished normally and its result was persisted.
+
+ Marks the end of the normal completion path; a later replay of the
+ same action skips the invocation. Not emitted when the action fails.
+ """
+
+ def on_action_reused(self, task: Any) -> None:
+ """A replayed already-completed action skipped its invocation.
+
+ This is the sole terminal callback on the reuse path.
+ """
+
+ def on_action_failed(self, task: Any, error: Any) -> None:
+ """An action invocation failed.
+
+ Purely observational: perceive the failure for logging, metrics, or
+ bookkeeping cleanup, but never compensate or decide on rethrowing.
+ """
+
+ def on_record_finished(self, key: Any) -> None:
+ """Every task spawned by an input record has completed.
+
+ Implementations must make their per-record cleanup idempotent: after
+ a failover replay the notification may not arrive again for records
+ that completed before the snapshot.
+ """
diff --git a/python/flink_agents/runtime/tests/test_async_subagent.py b/python/flink_agents/runtime/tests/test_async_subagent.py
new file mode 100644
index 000000000..090e7b152
--- /dev/null
+++ b/python/flink_agents/runtime/tests/test_async_subagent.py
@@ -0,0 +1,571 @@
+################################################################################
+# 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.
+################################################################################
+"""Tests for the async sub-agent base in pub/sub mode."""
+
+from concurrent.futures import CancelledError
+from typing import Any, NamedTuple
+
+import pytest
+from pydantic import PrivateAttr
+
+from flink_agents.api.subagent import SubagentResult
+from flink_agents.runtime.async_subagent import (
+ BaseAsyncSubagentSetup,
+ RunStatus,
+)
+from flink_agents.runtime.tests.test_base_subagent import _FakeTask, _run
+
+
+class _DurableExecuteCall(NamedTuple):
+ """One recorded durable execution invocation."""
+
+ func: Any
+ args: tuple
+ reconciler: Any
+ durable_id: str | None
+
+
+class _RecordingContext:
+ """Fake context recording durable execution and running it inline."""
+
+ def __init__(self) -> None:
+ self.durable_execute_calls: list[_DurableExecuteCall] = []
+ self.async_durable_calls = 0
+
+ def durable_execute(
+ self,
+ func: Any,
+ *args: Any,
+ reconciler: Any = None,
+ durable_id: str | None = None,
+ **kwargs: Any,
+ ) -> Any:
+ self.durable_execute_calls.append(
+ _DurableExecuteCall(func, args, reconciler, durable_id)
+ )
+ return func(*args)
+
+ def durable_execute_async(
+ self,
+ func: Any,
+ *args: Any,
+ reconciler: Any = None,
+ durable_id: str | None = None,
+ **kwargs: Any,
+ ) -> Any:
+ self.async_durable_calls += 1
+ self.durable_execute_calls.append(
+ _DurableExecuteCall(func, args, reconciler, durable_id)
+ )
+ return _ImmediateAwaitable(func(*args))
+
+
+class _ImmediateAwaitable:
+ """Awaitable resolving without yielding, mirroring a cached durable result."""
+
+ def __init__(self, value: Any) -> None:
+ self._value = value
+
+ def __await__(self) -> Any:
+ return self._value
+ yield # pragma: no cover - makes this a generator function
+
+
+class _MockAsyncSetup(BaseAsyncSubagentSetup):
+ """Example integration: an in-memory asynchronous agent service.
+
+ Demonstrates that an integration only supplies the transport primitives
+ plus the optional cancel hook; counters let tests assert how many times
+ each endpoint was hit.
+ """
+
+ _runs: dict = PrivateAttr(default_factory=dict)
+ _queries_until_complete: int = PrivateAttr(default=2)
+ _fail_on_post: bool = PrivateAttr(default=False)
+ _post_count: int = PrivateAttr(default=0)
+ _status_query_count: int = PrivateAttr(default=0)
+ _fetch_count: int = PrivateAttr(default=0)
+ _cancel_count: int = PrivateAttr(default=0)
+
+ def __init__(
+ self, queries_until_complete: int = 2, fail_on_post: bool = False
+ ) -> None:
+ super().__init__()
+ self._queries_until_complete = queries_until_complete
+ self._fail_on_post = fail_on_post
+ # 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.
+ self.status_poll_interval_millis = 0
+
+ def call_submit_request(self, session_id: str, call_id: str, prompt: Any) -> None:
+ self._post_count += 1
+ if self._fail_on_post:
+ msg = "post failed"
+ raise RuntimeError(msg)
+ self._runs[f"{session_id}#{call_id}"] = {
+ "result": f"done:{prompt}",
+ "error": None,
+ "queries_remaining": self._queries_until_complete,
+ }
+
+ def call_query_status(self, session_id: str, call_id: str) -> RunStatus:
+ self._status_query_count += 1
+ run = self._runs.get(f"{session_id}#{call_id}")
+ if run is None:
+ return RunStatus.not_started()
+ if run["queries_remaining"] > 0:
+ run["queries_remaining"] -= 1
+ return RunStatus.running()
+ if run["error"] is None:
+ return RunStatus.completed()
+ return RunStatus.failed(run["error"])
+
+ def call_fetch_result(self, session_id: str, call_id: str) -> SubagentResult:
+ self._fetch_count += 1
+ run = self._runs.get(f"{session_id}#{call_id}")
+ if run is None:
+ return SubagentResult.error("no run on record")
+ if run["error"] is None:
+ return SubagentResult.ok(run["result"])
+ return SubagentResult.error(run["error"])
+
+ def call_cancel_request(self, session_id: str, call_id: str) -> None:
+ self._cancel_count += 1
+
+ def seed_run(
+ self,
+ session_id: str,
+ call_id: str,
+ result: Any,
+ error: str | None,
+ queries_until_complete: int,
+ ) -> None:
+ """Inject a run that already exists remotely, exercising reconciler reuse."""
+ self._runs[f"{session_id}#{call_id}"] = {
+ "result": result,
+ "error": error,
+ "queries_remaining": queries_until_complete,
+ }
+
+ def forget_run(self, session_id: str, call_id: str) -> None:
+ """Drop the remote record of a run, simulating a POST that never landed."""
+ self._runs.pop(f"{session_id}#{call_id}", None)
+
+ def post_count(self) -> int:
+ """Number of times the POST endpoint has been hit."""
+ return self._post_count
+
+ def status_query_count(self) -> int:
+ """Number of times the status endpoint has been probed."""
+ return self._status_query_count
+
+ def fetch_count(self) -> int:
+ """Number of times the result endpoint has been fetched."""
+ return self._fetch_count
+
+ def cancel_count(self) -> int:
+ """Number of times the cancel hook has been invoked."""
+ return self._cancel_count
+
+
+# ------------------------------------------------------------------------------------------
+# Construction: the status_poll_interval_millis YAML argument
+# ------------------------------------------------------------------------------------------
+
+
+class _PlainAsyncSetup(BaseAsyncSubagentSetup):
+ """Concrete base without transport overrides or a custom constructor, so
+ construction kwargs flow through the pydantic validation of the fields.
+ """
+
+ def call_submit_request(self, session_id: str, call_id: str, prompt: Any) -> None:
+ raise NotImplementedError
+
+ def call_query_status(self, session_id: str, call_id: str) -> RunStatus:
+ raise NotImplementedError
+
+ def call_fetch_result(self, session_id: str, call_id: str) -> SubagentResult:
+ raise NotImplementedError
+
+
+def test_status_poll_interval_is_set_from_the_yaml_argument() -> None:
+ setup = _PlainAsyncSetup(status_poll_interval_millis=123)
+
+ assert setup.status_poll_interval_millis == 123
+
+
+def test_status_poll_interval_defaults_to_500() -> None:
+ assert _PlainAsyncSetup().status_poll_interval_millis == 500
+
+
+# ------------------------------------------------------------------------------------------
+# The pub: one durable POST, issued immediately
+# ------------------------------------------------------------------------------------------
+
+
+def test_submit_posts_immediately_under_the_call_identity() -> None:
+ setup = _MockAsyncSetup()
+ ctx = _RecordingContext()
+
+ future = _run(setup.submit(ctx, "ping", "sid-1", "call-1"))
+
+ assert setup.post_count() == 1
+ assert setup.status_query_count() == 0
+ assert setup.fetch_count() == 0
+ assert len(ctx.durable_execute_calls) == 1
+ # The pub POST ran through async durable execution.
+ assert ctx.async_durable_calls == 1
+ assert ctx.durable_execute_calls[0].durable_id == "sid-1#call-1"
+ assert ctx.durable_execute_calls[0].reconciler is not None
+ assert future.session_id == "sid-1"
+ assert future.call_id == "call-1"
+
+
+def test_post_failure_fails_the_submit() -> None:
+ setup = _MockAsyncSetup(fail_on_post=True)
+ ctx = _RecordingContext()
+
+ with pytest.raises(RuntimeError, match="post failed"):
+ _run(setup.submit(ctx, "ping", "sid-1", "call-1"))
+ assert setup.status_query_count() == 0
+ assert setup.fetch_count() == 0
+
+
+def test_short_forms_fail_without_a_prepared_task() -> None:
+ setup = _MockAsyncSetup()
+ ctx = _RecordingContext()
+
+ with pytest.raises(RuntimeError, match="No prepared action task"):
+ _run(setup.submit(ctx, "ping"))
+ with pytest.raises(RuntimeError, match="No prepared action task"):
+ _run(setup.submit(ctx, "ping", "sid-1"))
+ assert setup.post_count() == 0
+
+
+def test_short_forms_assign_through_the_prepared_task() -> None:
+ """Short forms assign ids from the executing task and POST under them."""
+ setup = _MockAsyncSetup()
+ ctx = _RecordingContext()
+ setup.on_action_prepared(_FakeTask())
+
+ handle = _run(setup.submit(ctx, "ping"))
+
+ assert setup.post_count() == 1
+ assert handle.call_id == f"{handle.session_id}-1"
+ posted = ctx.durable_execute_calls[0]
+ assert posted.durable_id == f"{handle.session_id}#{handle.call_id}"
+
+
+# ------------------------------------------------------------------------------------------
+# The crash-window reconciler of the POST
+# ------------------------------------------------------------------------------------------
+
+
+def _recorded_reconciler(setup: _MockAsyncSetup, ctx: _RecordingContext) -> Any:
+ _run(setup.submit(ctx, "ping", "sid-1", "call-1"))
+ return ctx.durable_execute_calls[0].reconciler
+
+
+def test_reconciler_reposts_when_the_run_is_not_on_record() -> None:
+ setup = _MockAsyncSetup(queries_until_complete=1)
+ ctx = _RecordingContext()
+ reconciler = _recorded_reconciler(setup, ctx)
+ # The remote has no record of the run: the POST never landed.
+ setup.forget_run("sid-1", "call-1")
+
+ reconciler()
+
+ # Probe reported NOT_STARTED, so the missing POST was issued exactly once.
+ assert setup.post_count() == 2
+ assert setup.status_query_count() == 1
+
+
+def test_reconciler_does_not_repost_a_running_run() -> None:
+ setup = _MockAsyncSetup(queries_until_complete=1)
+ ctx = _RecordingContext()
+ reconciler = _recorded_reconciler(setup, ctx)
+ setup.seed_run("sid-1", "call-1", "done:ping", None, 1)
+
+ reconciler()
+
+ assert setup.post_count() == 1
+ assert setup.status_query_count() == 1
+
+
+def test_reconciler_does_not_repost_a_terminal_run() -> None:
+ setup = _MockAsyncSetup(queries_until_complete=0)
+ ctx = _RecordingContext()
+ reconciler = _recorded_reconciler(setup, ctx)
+ setup.seed_run("sid-1", "call-1", "done:ping", None, 0)
+
+ reconciler()
+
+ assert setup.post_count() == 1
+ assert setup.status_query_count() == 1
+
+
+def test_reconciler_treats_a_failed_run_as_landed() -> None:
+ setup = _MockAsyncSetup(queries_until_complete=0)
+ ctx = _RecordingContext()
+ reconciler = _recorded_reconciler(setup, ctx)
+ setup.seed_run("sid-1", "call-1", None, "run exploded", 0)
+
+ reconciler()
+
+ assert setup.post_count() == 1
+ assert setup.status_query_count() == 1
+
+
+# ------------------------------------------------------------------------------------------
+# The sub: status probes and the await composition
+# ------------------------------------------------------------------------------------------
+
+
+def test_done_probes_the_status_directly_without_durable_calls() -> None:
+ setup = _MockAsyncSetup(queries_until_complete=1)
+ ctx = _RecordingContext()
+
+ future = _run(setup.submit(ctx, "ping", "sid-1", "call-1"))
+
+ assert future.done() is False # first probe: RUNNING
+ assert future.done() is True # second probe: COMPLETED
+ assert setup.status_query_count() == 2
+ assert len(ctx.durable_execute_calls) == 1 # only the pub POST
+
+
+def test_await_waits_durably_then_fetches() -> None:
+ setup = _MockAsyncSetup(queries_until_complete=2)
+ ctx = _RecordingContext()
+
+ future = _run(setup.submit(ctx, "ping", "sid-1", "call-1"))
+ result = _run(future)
+
+ assert result.success is True
+ assert result.result == "done:ping"
+ assert ctx.durable_execute_calls[1].durable_id == "sid-1#call-1#await"
+ # Two RUNNING probes, the terminal one, then the separate fetch.
+ assert setup.status_query_count() == 3
+ assert setup.fetch_count() == 1
+
+
+def test_await_surfaces_a_failed_run_without_fetching() -> None:
+ setup = _MockAsyncSetup(queries_until_complete=0)
+ ctx = _RecordingContext()
+
+ future = _run(setup.submit(ctx, "ping", "sid-1", "call-1"))
+ # The remote run fails before the handle resolves.
+ setup.seed_run("sid-1", "call-1", None, "run exploded", 0)
+ result = _run(future)
+
+ assert result.success is False
+ assert "run exploded" in result.error_message
+ assert setup.fetch_count() == 0
+
+
+def test_resolve_twice_runs_one_durable_await() -> None:
+ setup = _MockAsyncSetup(queries_until_complete=0)
+ ctx = _RecordingContext()
+
+ future = _run(setup.submit(ctx, "ping", "sid-1", "call-1"))
+ first = _run(future)
+ second = _run(future)
+
+ assert first is second
+ # Only the pub POST and one await composition.
+ assert len(ctx.durable_execute_calls) == 2
+
+
+def test_resolve_without_probing_goes_straight_to_the_await() -> None:
+ setup = _MockAsyncSetup(queries_until_complete=2)
+ ctx = _RecordingContext()
+
+ future = _run(setup.submit(ctx, "ping", "sid-1", "call-1"))
+ result = _run(future)
+
+ assert result.result == "done:ping"
+ # No done()-style probes: the await composition did them all.
+ assert setup.status_query_count() == 3
+ assert setup.fetch_count() == 1
+
+
+# ------------------------------------------------------------------------------------------
+# Failover replay: fresh probes may take a different path to the same result
+# ------------------------------------------------------------------------------------------
+
+
+def test_replay_after_the_run_completed_takes_fewer_probes() -> None:
+ # Original execution: the run completes only after two RUNNING probes.
+ original = _MockAsyncSetup(queries_until_complete=2)
+ original.seed_run("sid-1", "call-1", "done:ping", None, 2)
+ before = original._await_until_terminal("sid-1", "call-1")
+ assert original.status_query_count() == 3
+
+ # Replay: the run has already reached a terminal state, so the same await
+ # takes a shorter path — fewer probes — to the same result.
+ replay = _MockAsyncSetup(queries_until_complete=2)
+ replay.seed_run("sid-1", "call-1", "done:ping", None, 0)
+ after = replay._await_until_terminal("sid-1", "call-1")
+
+ assert after.success is True
+ assert after.result == before.result
+ assert replay.status_query_count() == 1
+
+
+# ------------------------------------------------------------------------------------------
+# Cancellation: the hook's return governs the cancelled resolve
+# ------------------------------------------------------------------------------------------
+
+
+def test_cancel_then_resolve_raises_cancelled_error_by_default() -> None:
+ setup = _MockAsyncSetup()
+ ctx = _RecordingContext()
+
+ future = _run(setup.submit(ctx, "ping", "sid-1", "call-1"))
+ future.cancel()
+
+ assert setup.cancel_count() == 1
+ assert future.done() is True
+ with pytest.raises(CancelledError):
+ _run(future)
+ # The pub landed, but the cancelled resolve never awaited nor fetched.
+ assert setup.post_count() == 1
+ assert setup.status_query_count() == 0
+ assert setup.fetch_count() == 0
+ assert len(ctx.durable_execute_calls) == 1
+
+
+def test_repeated_cancel_is_a_local_no_op() -> None:
+ setup = _MockAsyncSetup()
+ ctx = _RecordingContext()
+
+ future = _run(setup.submit(ctx, "ping", "sid-1", "call-1"))
+ future.cancel()
+ future.cancel()
+
+ # A repeated cancel on the same handle does not propagate again; a
+ # failover replay creates a fresh handle, which may.
+ assert setup.cancel_count() == 1
+
+
+def test_cancel_after_the_resolve_is_ignored() -> None:
+ setup = _MockAsyncSetup(queries_until_complete=0)
+ ctx = _RecordingContext()
+
+ future = _run(setup.submit(ctx, "ping", "sid-1", "call-1"))
+ assert _run(future).result == "done:ping"
+
+ future.cancel()
+
+ assert setup.cancel_count() == 0
+ assert _run(future).result == "done:ping"
+
+
+def test_cancel_propagates_through_the_group() -> None:
+ setup = _MockAsyncSetup()
+ ctx = _RecordingContext()
+
+ first = _run(setup.submit(ctx, "a", "sid-1", "call-1"))
+ second = _run(setup.submit(ctx, "b", "sid-1", "call-2"))
+ first.combine(second).cancel()
+
+ assert setup.cancel_count() == 2
+ with pytest.raises(CancelledError):
+ _run(first)
+ with pytest.raises(CancelledError):
+ _run(second)
+
+
+def test_combine_resolves_every_handle_of_the_batch() -> None:
+ setup = _MockAsyncSetup(queries_until_complete=0)
+ ctx = _RecordingContext()
+
+ first = _run(setup.submit(ctx, "a", "sid-1", "call-1"))
+ second = _run(setup.submit(ctx, "b", "sid-1", "call-2"))
+
+ outcomes = _run(first.combine(second))
+
+ assert [outcome.result for outcome in outcomes] == ["done:a", "done:b"]
+
+
+def test_the_await_form_waits_through_the_async_durable_composition() -> None:
+ setup = _MockAsyncSetup(queries_until_complete=1)
+ ctx = _RecordingContext()
+
+ future = _run(setup.submit(ctx, "ping", "sid-1", "call-1"))
+ result = _run(future)
+
+ assert result.result == "done:ping"
+ # One async durable call for the pub POST, one for the await.
+ assert ctx.async_durable_calls == 2
+ assert ctx.durable_execute_calls[1].durable_id == "sid-1#call-1#await"
+ assert future.done() is True
+
+
+def test_the_await_form_of_a_cancelled_handle_raises() -> None:
+ setup = _MockAsyncSetup()
+ ctx = _RecordingContext()
+
+ future = _run(setup.submit(ctx, "ping", "sid-1", "call-1"))
+ future.cancel()
+
+ with pytest.raises(CancelledError):
+ _run(future)
+ # Only the pub POST ran through async durable execution.
+ assert ctx.async_durable_calls == 1
+
+
+# ------------------------------------------------------------------------------------------
+# Pending-call registry: the task must resolve every handle it submits
+# ------------------------------------------------------------------------------------------
+
+
+def test_lifecycle_dropped_handles_fail_the_finished_task() -> None:
+ """An async handle left unresolved fails the task on finish, matching Java."""
+ setup = _MockAsyncSetup()
+ ctx = _RecordingContext()
+ setup.on_action_prepared(_FakeTask())
+
+ _run(setup.submit(ctx, "ping"))
+
+ with pytest.raises(RuntimeError, match="finished without resolving"):
+ setup.on_action_finishing(_FakeTask())
+
+
+def test_lifecycle_resolved_handles_let_the_task_finish() -> None:
+ """Awaiting every submitted handle lets the task finish cleanly."""
+ setup = _MockAsyncSetup(queries_until_complete=0)
+ ctx = _RecordingContext()
+ setup.on_action_prepared(_FakeTask())
+
+ handle = _run(setup.submit(ctx, "ping"))
+ _run(handle)
+
+ setup.on_action_finishing(_FakeTask())
+
+
+def test_lifecycle_cancelled_handles_let_the_task_finish() -> None:
+ """Cancelling a submitted handle unregisters it so the task finishes."""
+ setup = _MockAsyncSetup()
+ ctx = _RecordingContext()
+ setup.on_action_prepared(_FakeTask())
+
+ handle = _run(setup.submit(ctx, "ping"))
+ handle.cancel()
+
+ setup.on_action_finishing(_FakeTask())
diff --git a/python/flink_agents/runtime/tests/test_base_subagent.py b/python/flink_agents/runtime/tests/test_base_subagent.py
new file mode 100644
index 000000000..eb96a8c79
--- /dev/null
+++ b/python/flink_agents/runtime/tests/test_base_subagent.py
@@ -0,0 +1,305 @@
+################################################################################
+# 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.
+################################################################################
+"""Tests for the framework base of sub-agent setups.
+
+The Python parity of Java's ``BaseSubagentSetupTest``: lifecycle-driven
+id assignment, replay determinism, continuity across the steps of a
+suspended task, dropped-handle enforcement, and resource-name isolation.
+"""
+
+from typing import Any
+
+import pytest
+
+from flink_agents.api.resource import ResourceType
+from flink_agents.api.subagent import SubagentFuture, SubagentResult
+from flink_agents.runtime.base_subagent import BaseSubagentSetup
+from flink_agents.runtime.resource_cache import ResourceCache
+from flink_agents.runtime.subagent_handles import CompletedSubagentFuture
+
+
+class _FakeAction:
+ """Duck-typed ``Action`` exposing the name getter the base reads."""
+
+ def __init__(self, name: str) -> None:
+ self._name = name
+
+ def getName(self) -> str:
+ return self._name
+
+
+class _FakeEvent:
+ """Duck-typed ``Event`` exposing the getters the base reads."""
+
+ def __init__(
+ self,
+ event_type: str = "TestEvent",
+ attributes: dict[str, Any] | None = None,
+ event_id: str = "event-1",
+ ) -> None:
+ self._type = event_type
+ self._attributes = attributes or {}
+ self._id = event_id
+
+ def getType(self) -> str:
+ return self._type
+
+ def getAttributes(self) -> dict[str, Any]:
+ return self._attributes
+
+ def getId(self) -> str:
+ return self._id
+
+
+class _FakeTask:
+ """Duck-typed ``ActionTask`` carrying the caller-side facts."""
+
+ def __init__(
+ self,
+ key: str = "k",
+ sequence_number: int = 1,
+ action_name: str = "act",
+ event: _FakeEvent | None = None,
+ ) -> None:
+ self._key = key
+ self._sequence_number = sequence_number
+ self._action = _FakeAction(action_name)
+ self._event = event or _FakeEvent()
+
+ def getKey(self) -> str:
+ return self._key
+
+ def getSequenceNumber(self) -> int:
+ return self._sequence_number
+
+ def getAction(self) -> _FakeAction:
+ return self._action
+
+ def getEvent(self) -> _FakeEvent:
+ return self._event
+
+
+class _RecordingBaseSetup(BaseSubagentSetup):
+ """Base subclass completing handles without any transport."""
+
+ async def submit_with_identity(
+ self,
+ ctx: Any,
+ prompt: Any,
+ session_id: str,
+ call_id: str,
+ ) -> SubagentFuture:
+ """Complete the invocation immediately under the assigned identity."""
+ return CompletedSubagentFuture(session_id, call_id, SubagentResult.ok([prompt]))
+
+
+def _run(awaitable: Any) -> Any:
+ """Drive an awaitable the way the runtime drives an action coroutine."""
+ iterator = awaitable.__await__()
+ try:
+ while True:
+ next(iterator)
+ except StopIteration as stop:
+ return stop.value
+
+
+def test_short_forms_assign_through_the_prepared_task() -> None:
+ """The short forms allocate deterministically from the executing task."""
+ setup = _RecordingBaseSetup()
+ setup.on_action_prepared(_FakeTask())
+
+ first = _run(setup.submit(None, "p"))
+ second = _run(setup.submit(None, "p"))
+ under_session = _run(setup.submit(None, "p", "given-session"))
+ third_call = _run(setup.submit(None, "p", first.session_id))
+
+ assert first.session_id.endswith("-0")
+ assert first.call_id == f"{first.session_id}-1"
+ assert second.session_id.endswith("-1")
+ assert second.call_id == f"{second.session_id}-1"
+ assert third_call.call_id == f"{first.session_id}-2"
+ assert under_session.session_id == "given-session"
+ assert under_session.call_id == "given-session-1"
+
+
+def test_replay_assigns_the_same_ids() -> None:
+ """A failover replay of the same task facts hands out the same ids."""
+ first = _RecordingBaseSetup()
+ first.on_action_prepared(_FakeTask(key="k", sequence_number=7, action_name="act"))
+ original = _run(first.submit(None, "p"))
+
+ replay = _RecordingBaseSetup()
+ replay.on_action_prepared(_FakeTask(key="k", sequence_number=7, action_name="act"))
+ replayed = _run(replay.submit(None, "p"))
+
+ assert replayed.session_id == original.session_id
+ assert replayed.call_id == original.call_id
+
+
+def test_allocation_continues_across_task_steps() -> None:
+ """Each step of a suspended task re-prepares with the same facts, and
+ the allocator persists, so the session ordinal continues instead of
+ restarting.
+ """
+ setup = _RecordingBaseSetup()
+ setup.on_action_prepared(_FakeTask())
+ first = _run(setup.submit(None, "p"))
+
+ setup.on_action_prepared(_FakeTask())
+ second = _run(setup.submit(None, "p"))
+
+ assert first.session_id.endswith("-0")
+ assert second.session_id.endswith("-1")
+ assert second.session_id != first.session_id
+
+
+def test_transfer_moves_bookkeeping_onto_a_different_continuation() -> None:
+ """The generated task may carry a different identity than the finishing
+ task, so the allocator and the pending-call registry are re-keyed onto
+ it instead of assumed equal.
+ """
+ setup = _RecordingBaseSetup()
+ from_task = _FakeTask(event=_FakeEvent(event_id="event-from"))
+ to_task = _FakeTask(action_name="act-next", event=_FakeEvent(event_id="event-to"))
+
+ setup.on_action_prepared(from_task)
+ first = _run(setup.submit(None, "p"))
+ setup.pending_call_registry().track_pending_subagent_call("sid#call-1")
+ setup.on_action_transferred(from_task, to_task)
+
+ setup.on_action_prepared(to_task)
+ second = _run(setup.submit(None, "p"))
+
+ # The allocator moved with the execution: the session ordinal continues
+ # instead of restarting under the continuation's own facts.
+ assert first.session_id.endswith("-0")
+ assert second.session_id.endswith("-1")
+
+ # The pending-call registry moved as well, and adopted the continuation's
+ # action: finishing the continuation still reports the handle tracked
+ # under the finishing task.
+ with pytest.raises(RuntimeError, match=r"act-next.*sid#call-1"):
+ setup.on_action_finishing(to_task)
+
+
+def test_finished_task_drops_its_bookkeeping() -> None:
+ """After the task finishes, short forms have no task to assign from."""
+ setup = _RecordingBaseSetup()
+ setup.on_action_prepared(_FakeTask())
+ _run(setup.submit(None, "p"))
+
+ setup.on_action_finishing(_FakeTask())
+
+ with pytest.raises(RuntimeError, match="No prepared action task"):
+ _run(setup.submit(None, "p"))
+
+
+def test_new_action_restarts_call_ordinal_for_a_reused_session_id() -> None:
+ """Ids assigned without an explicit id are only valid within one action
+ execution: a new task starts the per-session call ordinal at 1 again, so
+ reusing a session id across actions reproduces the same call ids.
+ """
+ setup = _RecordingBaseSetup()
+ setup.on_action_prepared(_FakeTask(sequence_number=1))
+ first = _run(setup.submit(None, "p", "shared-session"))
+ second = _run(setup.submit(None, "p", "shared-session"))
+ assert first.call_id == "shared-session-1"
+ assert second.call_id == "shared-session-2"
+
+ # The next action prepares a different task; its fresh allocator hands
+ # out the identical ids under the reused session id.
+ setup.on_action_finishing(_FakeTask(sequence_number=1))
+ setup.on_action_prepared(_FakeTask(sequence_number=2))
+ reused = _run(setup.submit(None, "p", "shared-session"))
+ assert reused.call_id == "shared-session-1"
+
+
+def test_subagent_name_isolates_namespaces() -> None:
+ """Setups sharing one caller's counting range assign disjoint ids."""
+ left = _RecordingBaseSetup()
+ left.set_subagent_name("scope.left")
+ left.on_action_prepared(_FakeTask())
+
+ right = _RecordingBaseSetup()
+ right.set_subagent_name("scope.right")
+ right.on_action_prepared(_FakeTask())
+
+ left_handle = _run(left.submit(None, "p"))
+ right_handle = _run(right.submit(None, "p"))
+
+ assert left_handle.session_id != right_handle.session_id
+
+
+def test_interleaved_executions_of_one_action_keep_apart() -> None:
+ """Tasks of one action triggered by different events within one record
+ interleave: the earlier one may still be suspended when the later one
+ finishes, and their bookkeeping must not mix.
+ """
+ setup = _RecordingBaseSetup()
+ first_task = _FakeTask(event=_FakeEvent(event_id="e1"))
+ second_task = _FakeTask(event=_FakeEvent(event_id="e2"))
+
+ setup.on_action_prepared(first_task)
+ first = _run(setup.submit(None, "p"))
+ left = setup.pending_call_registry()
+ left.track_pending_subagent_call(first.identity)
+
+ setup.on_action_prepared(second_task)
+ _run(setup.submit(None, "p"))
+ setup.on_action_finishing(second_task)
+
+ # The finished execution dropped only its own bookkeeping.
+ setup.on_action_prepared(first_task)
+ resumed = _run(setup.submit(None, "p"))
+ assert resumed.session_id.endswith("-1")
+ assert setup.pending_call_registry() is left
+
+ left.untrack_pending_subagent_call(first.identity)
+ setup.on_action_finishing(first_task)
+
+
+def test_explicit_identity_passes_through_untouched() -> None:
+ """A fully supplied identity skips the allocator entirely."""
+ setup = _RecordingBaseSetup()
+
+ handle = _run(setup.submit(None, "p", "sid-x", "call-y"))
+
+ assert handle.session_id == "sid-x"
+ assert handle.call_id == "call-y"
+
+
+class _FakeProvider:
+ """Resource provider returning a fixed resource instance."""
+
+ def __init__(self, resource: Any) -> None:
+ self._resource = resource
+
+ def provide(self, resource_context: Any, config: Any) -> Any:
+ """Return the pre-built resource."""
+ return self._resource
+
+
+def test_resource_cache_injects_the_subagent_name() -> None:
+ """Materializing a sub-agent setup injects the resource name, like Java."""
+ setup = _RecordingBaseSetup()
+ cache = ResourceCache({ResourceType.AGENT: {"reviewer": _FakeProvider(setup)}})
+
+ resolved = cache.get_resource("reviewer", ResourceType.AGENT)
+
+ assert resolved is setup
+ assert setup.subagent_name == "reviewer"
diff --git a/python/flink_agents/runtime/tests/test_deferred_subagent.py b/python/flink_agents/runtime/tests/test_deferred_subagent.py
new file mode 100644
index 000000000..243eff69a
--- /dev/null
+++ b/python/flink_agents/runtime/tests/test_deferred_subagent.py
@@ -0,0 +1,452 @@
+################################################################################
+# 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.
+################################################################################
+"""Tests for the deferred sub-agent futures and call routing."""
+
+from concurrent.futures import CancelledError
+from typing import Any, Callable, NamedTuple
+
+import pytest
+
+from flink_agents.api.subagent import SubagentResult
+from flink_agents.runtime.deferred_subagent import (
+ DeferredSubagentFuture,
+ DeferredSubagentSetup,
+)
+from flink_agents.runtime.subagent_handles import (
+ CompletedSubagentFuture,
+ PendingSubagentCallRegistry,
+)
+from flink_agents.runtime.tests.test_base_subagent import _FakeTask, _run
+
+
+class _DurableExecuteCall(NamedTuple):
+ """One recorded ``durable_execute`` invocation."""
+
+ func: Any
+ args: tuple
+ reconciler: Any
+ durable_id: str | None
+
+
+class _RecordingContext:
+ """Fake context recording durable execution."""
+
+ def __init__(self) -> None:
+ self.durable_execute_calls: list[_DurableExecuteCall] = []
+
+ def durable_execute(
+ self,
+ func: Any,
+ *args: Any,
+ reconciler: Any = None,
+ durable_id: str | None = None,
+ **kwargs: Any,
+ ) -> Any:
+ self.durable_execute_calls.append(
+ _DurableExecuteCall(func, args, reconciler, durable_id)
+ )
+ return func(*args)
+
+
+class _AwaitingContext(_RecordingContext):
+ """Recording context whose ``durable_execute_async`` is awaitable.
+
+ Records how many requests had been issued when the first wait started:
+ a batched wait prepares every pending deferred handle up front, so the
+ whole batch is issued before any execution starts.
+ """
+
+ def __init__(self) -> None:
+ super().__init__()
+ self.issued_before_first_wait: int | None = None
+ self.issued_count = 0
+
+ def durable_execute_async(
+ self,
+ func: Any,
+ *args: Any,
+ reconciler: Any = None,
+ durable_id: str | None = None,
+ **kwargs: Any,
+ ) -> Any:
+ if self.issued_before_first_wait is None:
+ self.issued_before_first_wait = self.issued_count
+ self.durable_execute_calls.append(
+ _DurableExecuteCall(func, args, reconciler, durable_id)
+ )
+ return _ImmediateAwaitable(func(*args))
+
+
+class _ImmediateAwaitable:
+ """Awaitable resolving without yielding, mirroring a cached durable result."""
+
+ def __init__(self, value: Any) -> None:
+ self._value = value
+
+ def __await__(self) -> Any:
+ return self._value
+ yield # pragma: no cover - makes this a generator function
+
+
+def _echo_callable(prompt: Any) -> Callable[[], SubagentResult]:
+ """Callable echoing the prompt as a successful result."""
+
+ def call() -> SubagentResult:
+ return SubagentResult.ok([prompt])
+
+ return call
+
+
+def _raising_callable(exc: Exception) -> Callable[[], SubagentResult]:
+ """Callable raising a system-level failure instead of returning a result."""
+
+ def call() -> SubagentResult:
+ raise exc
+
+ return call
+
+
+_ISSUED = 0
+_REGISTRY: PendingSubagentCallRegistry | None = None
+
+
+def _reset_echoing_state() -> None:
+ global _ISSUED, _REGISTRY
+ _ISSUED = 0
+ _REGISTRY = None
+
+
+class _MockDeferredSetup(DeferredSubagentSetup):
+ """Setup issuing deferred futures like the runtime bases do.
+
+ Counts how many times a request has been issued in module-level state
+ (one setup instance is shared by every resolving task, mirroring the
+ runtime); an optional per-task registry records deferred handles until
+ they resolve.
+ """
+
+ def prepare(
+ self,
+ ctx: Any,
+ prompt: Any,
+ session_id: str,
+ call_id: str,
+ ) -> tuple:
+ """Count the issue and prepare the echoing triple."""
+ global _ISSUED
+ _ISSUED += 1
+ if isinstance(ctx, _AwaitingContext):
+ ctx.issued_count += 1
+ return (f"{session_id}#{call_id}", _echo_callable(prompt), None)
+
+ def pending_call_registry(self) -> PendingSubagentCallRegistry | None:
+ """Opt into tracking when a registry has been assigned."""
+ return _REGISTRY
+
+
+class _LifecycleDeferredSetup(DeferredSubagentSetup):
+ """Deferred setup tracking handles through the base's per-task registry."""
+
+ def prepare(
+ self,
+ ctx: Any,
+ prompt: Any,
+ session_id: str,
+ call_id: str,
+ ) -> tuple:
+ """Prepare the echoing triple keyed by the assigned identity."""
+ return (f"{session_id}#{call_id}", _echo_callable(prompt), None)
+
+
+def test_prepare_returns_the_prepared_triple() -> None:
+ """``prepare`` supplies the durable id, the call, and the reconciler."""
+ setup = _MockDeferredSetup()
+ ctx = _RecordingContext()
+
+ durable_id, call, reconcile = setup.prepare(ctx, "ping", "sid-1", "call-1")
+
+ assert durable_id == "sid-1#call-1"
+ assert reconcile is None
+ assert call().result == ["ping"]
+
+
+def test_submit_with_explicit_ids_routes_durably() -> None:
+ """``submit`` with explicit ids routes through durable execution."""
+ setup = _MockDeferredSetup()
+ ctx = _AwaitingContext()
+
+ future = _run(setup.submit(ctx, "hello", "explicit-sid", "call-1"))
+ result = _run(future)
+
+ assert len(ctx.durable_execute_calls) == 1
+ assert ctx.durable_execute_calls[0].durable_id == "explicit-sid#call-1"
+ assert result.success is True
+ assert result.result == ["hello"]
+
+
+def test_short_forms_without_a_task_fail_to_assign() -> None:
+ """Short forms assign through the executing task; without one they fail."""
+ setup = _MockDeferredSetup()
+ ctx = _RecordingContext()
+
+ with pytest.raises(RuntimeError, match="No prepared action task"):
+ _run(setup.submit(ctx, "ping"))
+ with pytest.raises(RuntimeError, match="No prepared action task"):
+ _run(setup.submit(ctx, "ping", "sid-1"))
+
+
+def test_callable_id_is_derived_from_the_explicit_identity() -> None:
+ """The framework keys the durable call by the caller-supplied ids."""
+ setup = _MockDeferredSetup()
+ ctx = _AwaitingContext()
+
+ future = _run(setup.submit(ctx, "ping", "sid-1", "call-1"))
+ _run(future)
+
+ assert ctx.durable_execute_calls[0].durable_id == "sid-1#call-1"
+
+
+def test_submit_returns_handle_carrying_the_explicit_identity() -> None:
+ """``submit`` exposes the caller-supplied identity on the handle."""
+ setup = _MockDeferredSetup()
+ ctx = _AwaitingContext()
+
+ future = _run(setup.submit(ctx, "ping", "sid-1", "call-1"))
+
+ assert future.session_id == "sid-1"
+ assert future.call_id == "call-1"
+ assert future.done() is False
+ assert _run(future).result == ["ping"]
+ assert future.done() is True
+
+
+def test_submit_resolves_once_and_keys_by_the_call_identity() -> None:
+ """Resolving twice runs one durable call, keyed by ``session_id#call_id``."""
+ setup = _MockDeferredSetup()
+ ctx = _AwaitingContext()
+
+ future = _run(setup.submit(ctx, "ping", "sid-1", "call-1"))
+ first = _run(future)
+ second = _run(future)
+
+ assert first is second
+ assert len(ctx.durable_execute_calls) == 1
+ assert ctx.durable_execute_calls[0].durable_id == "sid-1#call-1"
+
+
+def test_submit_defers_the_request_until_resolve() -> None:
+ """``submit`` never issues the request up front; resolve does."""
+ global _REGISTRY
+ _reset_echoing_state()
+ ctx = _AwaitingContext()
+ setup = _MockDeferredSetup()
+ registry = PendingSubagentCallRegistry("my_action")
+ _REGISTRY = registry
+
+ future = _run(setup.submit(ctx, "ping", "sid-1", "call-1"))
+
+ assert _ISSUED == 0
+ assert registry.is_empty() is False
+
+ _run(future)
+
+ assert _ISSUED == 1
+ assert registry.is_empty() is True
+
+
+def test_deferred_handles_are_resolved_in_submission_order() -> None:
+ """``combine`` resolves every handle in submission order; the group
+ prepares the whole batch before any execution starts.
+ """
+ global _REGISTRY
+ _reset_echoing_state()
+ ctx = _AwaitingContext()
+ setup = _MockDeferredSetup()
+ registry = PendingSubagentCallRegistry("my_action")
+ _REGISTRY = registry
+
+ first = _run(setup.submit(ctx, "a", "sid-1", "call-1"))
+ second = _run(setup.submit(ctx, "b", "sid-1", "call-2"))
+ third = _run(setup.submit(ctx, "c", "sid-1", "call-3"))
+ assert _ISSUED == 0
+
+ outcomes = _run(first.combine(second, third))
+
+ assert [outcome.result for outcome in outcomes] == [["a"], ["b"], ["c"]]
+ assert first.done()
+ assert second.done()
+ assert third.done()
+ assert registry.is_empty() is True
+ # The group prepared the whole batch before the first execution
+ # started.
+ assert ctx.issued_before_first_wait == 3
+ assert _ISSUED == 3
+
+
+def test_batching_an_already_resolved_handle_joins_the_batch() -> None:
+ """A resolved handle contributes its value to a mixed batch."""
+ ctx = _AwaitingContext()
+ setup = _MockDeferredSetup()
+
+ resolved = CompletedSubagentFuture("s", "c", SubagentResult.ok("x"))
+ pending = _run(setup.submit(ctx, "pending", "sid-1", "call-1"))
+
+ outcomes = _run(resolved.combine(pending))
+
+ # Only the pending handle prepared a request; the resolved one kept its
+ # value.
+ assert ctx.issued_before_first_wait == 1
+ assert [outcome.result for outcome in outcomes] == ["x", ["pending"]]
+ assert pending.done()
+ assert resolved.done()
+
+
+def test_registry_check_empty_fails_on_dropped_handles() -> None:
+ """``check_empty`` names every handle the action dropped unresolved."""
+ registry = PendingSubagentCallRegistry("my_action")
+ registry.track_pending_subagent_call("sid-1#call-1")
+
+ with pytest.raises(RuntimeError, match="sid-1#call-1"):
+ registry.check_empty()
+
+ # The state is left intact, so the caller can inspect the dropped handles.
+ assert registry.is_empty() is False
+
+
+def test_system_level_failure_propagates_out_of_resolve() -> None:
+ """An exception escaping the callable propagates instead of folding.
+
+ The integration folds its own comprehensible failures into the
+ SubagentResult; a raised exception is system-level.
+ """
+ ctx = _AwaitingContext()
+ boom = RuntimeError("durable execution crashed")
+
+ future = DeferredSubagentFuture(
+ "sid-1",
+ "call-1",
+ ctx,
+ prepared_factory=lambda: ("sid-1#call-1", _raising_callable(boom), None),
+ )
+
+ with pytest.raises(RuntimeError, match="durable execution crashed"):
+ _run(future)
+ assert future.done() is False
+
+
+def test_lifecycle_dropped_handles_fail_the_finished_task() -> None:
+ """A short-form handle left unresolved fails the task on finish."""
+ setup = _LifecycleDeferredSetup()
+ setup.on_action_prepared(_FakeTask())
+ _run(setup.submit(_AwaitingContext(), "p"))
+
+ with pytest.raises(RuntimeError, match="finished without resolving"):
+ setup.on_action_finishing(_FakeTask())
+
+
+def test_lifecycle_an_unawaited_submit_registers_nothing() -> None:
+ """Dropping the submit awaitable issues nothing, so the finish check
+ finds no handle to report and the mistake stays invisible to it.
+ """
+ setup = _LifecycleDeferredSetup()
+ setup.on_action_prepared(_FakeTask())
+ submission = setup.submit(_AwaitingContext(), "p")
+
+ setup.on_action_finishing(_FakeTask())
+
+ # Close the dropped coroutine so it does not warn while other tests run.
+ submission.close()
+
+
+def test_lifecycle_resolved_handles_let_the_task_finish() -> None:
+ """Resolving every short-form handle lets the task finish cleanly."""
+ setup = _LifecycleDeferredSetup()
+ setup.on_action_prepared(_FakeTask())
+ handle = _run(setup.submit(_AwaitingContext(), "p"))
+
+ _run(handle)
+ setup.on_action_finishing(_FakeTask())
+
+
+def test_deferred_future_prepares_through_the_factory_once() -> None:
+ """Handles prepare through the supplied factory exactly once."""
+ ctx = _AwaitingContext()
+ calls = 0
+
+ def factory() -> tuple:
+ nonlocal calls
+ calls += 1
+ return ("sid-1#call-1", _echo_callable("ping"), None)
+
+ future: DeferredSubagentFuture = DeferredSubagentFuture(
+ "sid-1", "call-1", ctx, prepared_factory=factory
+ )
+
+ assert _run(future).result == ["ping"]
+ assert calls == 1
+
+
+def test_cancel_before_resolve_discards_the_request() -> None:
+ """A cancelled deferred handle never issues the request."""
+ global _REGISTRY
+ _reset_echoing_state()
+ ctx = _RecordingContext()
+ setup = _MockDeferredSetup()
+ registry = PendingSubagentCallRegistry("my_action")
+ _REGISTRY = registry
+
+ future = _run(setup.submit(ctx, "ping", "sid-1", "call-1"))
+ future.cancel()
+
+ assert _ISSUED == 0
+ # The cancelled handle unregisters, so tracking setups see nothing left.
+ assert registry.is_empty() is True
+ assert future.done() is True
+ with pytest.raises(CancelledError):
+ future.prepare()
+ with pytest.raises(CancelledError):
+ _run(future)
+ assert ctx.durable_execute_calls == []
+
+
+def test_cancel_propagates_through_the_group() -> None:
+ """A group cancel reaches every handle; resolving the batch fails."""
+ _reset_echoing_state()
+ ctx = _AwaitingContext()
+ setup = _MockDeferredSetup()
+
+ first = _run(setup.submit(ctx, "a", "sid-1", "call-1"))
+ second = _run(setup.submit(ctx, "b", "sid-1", "call-2"))
+ first.combine(second).cancel()
+
+ with pytest.raises(CancelledError):
+ _run(first.combine(second))
+ assert _ISSUED == 0
+
+
+def test_cancel_of_a_resolved_handle_is_ignored() -> None:
+ """An already resolved handle keeps its value after a cancel request."""
+ ctx = _AwaitingContext()
+ setup = _MockDeferredSetup()
+
+ future = _run(setup.submit(ctx, "ping", "sid-1", "call-1"))
+ assert _run(future).result == ["ping"]
+
+ future.cancel()
+
+ assert _run(future).result == ["ping"]
diff --git a/python/flink_agents/runtime/tests/test_durable_execution.py b/python/flink_agents/runtime/tests/test_durable_execution.py
index f652b0ee8..e3d6f9001 100644
--- a/python/flink_agents/runtime/tests/test_durable_execution.py
+++ b/python/flink_agents/runtime/tests/test_durable_execution.py
@@ -24,6 +24,8 @@ from flink_agents.runtime.durable_execution import (
_compute_args_digest,
_compute_function_id,
_validate_reconciler_callable,
+ get_durable_id,
+ with_durable_id,
)
@@ -278,3 +280,22 @@ def test_cloudpickle_none_exception_message() -> None:
assert isinstance(deserialized, RuntimeError)
# str() of an exception with None message is "None"
assert str(deserialized) == "None"
+
+
+def test_with_durable_id_overrides_derived_function_id() -> None:
+ """An explicit durable id wins over the derived module/qualname id."""
+ wrapped = with_durable_id(sample_function, "session-1#call-1")
+
+ assert get_durable_id(wrapped) == "session-1#call-1"
+ assert _compute_function_id(wrapped) == "session-1#call-1"
+ # The wrapper stays invocable and does not mutate the original callable.
+ assert wrapped(1, 2) == 3
+ assert get_durable_id(sample_function) is None
+
+
+def test_get_durable_id_returns_none_for_plain_callables() -> None:
+ """Callables without an explicit id report None and keep derived ids."""
+ assert get_durable_id(sample_function) is None
+ derived = _compute_function_id(sample_function)
+ assert derived == _compute_function_id(sample_function)
+ assert "sample_function" in derived
diff --git a/python/flink_agents/runtime/tests/test_eager_materialize.py b/python/flink_agents/runtime/tests/test_eager_materialize.py
new file mode 100644
index 000000000..ee840acba
--- /dev/null
+++ b/python/flink_agents/runtime/tests/test_eager_materialize.py
@@ -0,0 +1,129 @@
+################################################################################
+# 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.
+################################################################################
+"""Tests for materializing the resources the Python runtime owns.
+
+The Java resource cache cannot build a resource declared by a Python provider,
+so it asks the Python runtime to materialize its own resources and keeps a
+handle to each. These tests exercise that Python-side entry with stub providers,
+without a live interpreter.
+"""
+
+from typing import Any
+
+from flink_agents.api.resource import ResourceType
+from flink_agents.plan.resource_provider import (
+ JavaResourceProvider,
+ PythonResourceProvider,
+ PythonSerializableResourceProvider,
+)
+from flink_agents.runtime.flink_runner_context import FlinkRunnerContext
+
+
+class _StubResourceCache:
+ """Resource cache recording every resolution and returning a marker."""
+
+ def __init__(self) -> None:
+ self.resolved: list = []
+
+ def get_resource(self, name: str, type: ResourceType) -> Any:
+ self.resolved.append((name, type))
+ return f"resource:{name}"
+
+
+class _StubAgentPlan:
+ """Agent plan exposing only the resource providers."""
+
+ def __init__(self, resource_providers: dict) -> None:
+ self.resource_providers = resource_providers
+
+
+def _context(resource_providers: dict) -> tuple[FlinkRunnerContext, _StubResourceCache]:
+ """Build a FlinkRunnerContext over the given providers.
+
+ Bypasses ``__init__`` (which needs a Java runner context) and injects the
+ plan and cache the materialization reads.
+ """
+ ctx = FlinkRunnerContext.__new__(FlinkRunnerContext)
+ cache = _StubResourceCache()
+ ctx._FlinkRunnerContext__agent_plan = _StubAgentPlan(resource_providers)
+ ctx._FlinkRunnerContext__resource_cache = cache
+ return ctx, cache
+
+
+def _python_provider(name: str) -> PythonSerializableResourceProvider:
+ return PythonSerializableResourceProvider.model_construct(
+ name=name, type=ResourceType.CHAT_MODEL
+ )
+
+
+def _python_descriptor_provider(name: str) -> PythonResourceProvider:
+ return PythonResourceProvider.model_construct(
+ name=name, type=ResourceType.CHAT_MODEL
+ )
+
+
+def _java_provider(name: str) -> JavaResourceProvider:
+ return JavaResourceProvider.model_construct(name=name, type=ResourceType.CHAT_MODEL)
+
+
+def test_python_owned_resources_are_materialized_and_keyed_by_name() -> None:
+ """Both Python provider kinds are materialized through the resource cache."""
+ ctx, cache = _context(
+ {
+ ResourceType.CHAT_MODEL: {
+ "declared": _python_provider("declared"),
+ "from_yaml": _python_descriptor_provider("from_yaml"),
+ }
+ }
+ )
+
+ materialized = ctx.eager_materialize(ResourceType.CHAT_MODEL.value)
+
+ assert materialized == {
+ "declared": "resource:declared",
+ "from_yaml": "resource:from_yaml",
+ }
+ assert cache.resolved == [
+ ("declared", ResourceType.CHAT_MODEL),
+ ("from_yaml", ResourceType.CHAT_MODEL),
+ ]
+
+
+def test_java_owned_resources_are_left_to_the_java_cache() -> None:
+ """A Java-owned resource is not built a second time in the Python runtime."""
+ ctx, cache = _context(
+ {
+ ResourceType.CHAT_MODEL: {
+ "python": _python_provider("python"),
+ "java": _java_provider("java"),
+ }
+ }
+ )
+
+ materialized = ctx.eager_materialize(ResourceType.CHAT_MODEL.value)
+
+ assert materialized == {"python": "resource:python"}
+ assert cache.resolved == [("python", ResourceType.CHAT_MODEL)]
+
+
+def test_a_type_without_providers_materializes_nothing() -> None:
+ """The type the operator asks for may not exist in the plan at all."""
+ ctx, cache = _context({})
+
+ assert ctx.eager_materialize(ResourceType.CHAT_MODEL.value) == {}
+ assert cache.resolved == []
diff --git a/python/flink_agents/runtime/tests/test_task_lifecycle_bridge.py b/python/flink_agents/runtime/tests/test_task_lifecycle_bridge.py
new file mode 100644
index 000000000..2b82bbd85
--- /dev/null
+++ b/python/flink_agents/runtime/tests/test_task_lifecycle_bridge.py
@@ -0,0 +1,177 @@
+################################################################################
+# 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.
+################################################################################
+"""Tests for the Python side of the task lifecycle bridge.
+
+The Java operator forwards its per-record and per-action lifecycle to Python
+over pemja, and ``FlinkRunnerContext`` fans each callback out to the
+registered ``TaskLifecycleListener``s. These tests exercise that registration
+and fan-out with a recording listener, without a live interpreter.
+"""
+
+from typing import Any
+
+from flink_agents.runtime.flink_runner_context import FlinkRunnerContext
+from flink_agents.runtime.task_lifecycle_listener import TaskLifecycleListener
+
+
+class _RecordingListener(TaskLifecycleListener):
+ """Lifecycle listener recording every callback it receives."""
+
+ def __init__(self) -> None:
+ self.calls: list = []
+
+ def on_record_start(self, key: Any) -> None:
+ self.calls.append(("record_start", key))
+
+ def on_action_prepared(self, task: Any) -> None:
+ self.calls.append(("action_prepared", task))
+
+ def on_action_started(self, task: Any) -> None:
+ self.calls.append(("action_started", task))
+
+ def on_action_transferred(self, from_task: Any, to_task: Any) -> None:
+ self.calls.append(("action_transferred", from_task, to_task))
+
+ def on_action_finishing(self, task: Any) -> None:
+ self.calls.append(("action_finishing", task))
+
+ def on_action_finished(self, task: Any) -> None:
+ self.calls.append(("action_finished", task))
+
+ def on_action_reused(self, task: Any) -> None:
+ self.calls.append(("action_reused", task))
+
+ def on_action_failed(self, task: Any, error: Any) -> None:
+ self.calls.append(("action_failed", task, error))
+
+ def on_record_finished(self, key: Any) -> None:
+ self.calls.append(("record_finished", key))
+
+
+def _context() -> FlinkRunnerContext:
+ """Build a FlinkRunnerContext with an empty listener registry.
+
+ Bypasses ``__init__`` (which needs a Java runner context) and starts from
+ an empty registry, as the operator does before registering listeners.
+ """
+ ctx = FlinkRunnerContext.__new__(FlinkRunnerContext)
+ ctx._FlinkRunnerContext__task_lifecycle_listeners = []
+ return ctx
+
+
+def test_fan_out_forwards_every_callback_in_order() -> None:
+ """Each callback reaches the registered listener, in order."""
+ ctx = _context()
+ listener = _RecordingListener()
+ ctx.add_task_lifecycle_listener(listener)
+
+ ctx.notify_record_start("k")
+ ctx.notify_action_prepared("t")
+ ctx.notify_action_started("t")
+ ctx.notify_action_transferred("t", "t2")
+ ctx.notify_action_finishing("t2")
+ ctx.notify_action_finished("t2")
+ ctx.notify_record_finished("k")
+
+ assert listener.calls == [
+ ("record_start", "k"),
+ ("action_prepared", "t"),
+ ("action_started", "t"),
+ ("action_transferred", "t", "t2"),
+ ("action_finishing", "t2"),
+ ("action_finished", "t2"),
+ ("record_finished", "k"),
+ ]
+
+
+def test_fan_out_forwards_reuse_and_failure_terminals() -> None:
+ """The reuse and failure paths reach the listener as their own callbacks."""
+ ctx = _context()
+ listener = _RecordingListener()
+ ctx.add_task_lifecycle_listener(listener)
+
+ ctx.notify_action_reused("t")
+ ctx.notify_action_failed("t", "boom")
+
+ assert listener.calls == [
+ ("action_reused", "t"),
+ ("action_failed", "t", "boom"),
+ ]
+
+
+def test_fan_out_reaches_every_registered_listener() -> None:
+ """A callback is delivered to all registered listeners."""
+ ctx = _context()
+ first, second = _RecordingListener(), _RecordingListener()
+ ctx.add_task_lifecycle_listener(first)
+ ctx.add_task_lifecycle_listener(second)
+
+ ctx.notify_action_prepared("t")
+
+ assert first.calls == [("action_prepared", "t")]
+ assert second.calls == [("action_prepared", "t")]
+
+
+def test_module_entries_delegate_to_the_context() -> None:
+ """The Java-invoked module functions delegate to the context fan-out."""
+ from flink_agents.runtime import flink_runner_context as frc
+
+ ctx = _context()
+ listener = _RecordingListener()
+ ctx.add_task_lifecycle_listener(listener)
+
+ frc.notify_record_start(ctx, "k")
+ frc.notify_action_prepared(ctx, "t")
+ frc.notify_action_transferred(ctx, "t", "t2")
+ frc.notify_action_finishing(ctx, "t2")
+ frc.notify_action_finished(ctx, "t2")
+ frc.notify_record_finished(ctx, "k")
+
+ assert [call[0] for call in listener.calls] == [
+ "record_start",
+ "action_prepared",
+ "action_transferred",
+ "action_finishing",
+ "action_finished",
+ "record_finished",
+ ]
+
+
+def test_registration_entry_registers_an_observing_listener() -> None:
+ """The Java side registers a Python listener through the module entry."""
+ from flink_agents.runtime import flink_runner_context as frc
+
+ ctx = _context()
+ listener = _RecordingListener()
+
+ assert frc.add_task_lifecycle_listener(ctx, listener) is True
+
+ ctx.notify_record_start("k")
+ assert listener.calls == [("record_start", "k")]
+
+
+def test_registration_entry_rejects_an_object_that_ignores_the_lifecycle() -> None:
+ """Registering a non-listener reports that there is nothing to notify."""
+ from flink_agents.runtime import flink_runner_context as frc
+
+ ctx = _context()
+
+ assert frc.add_task_lifecycle_listener(ctx, object()) is False
+
+ # Nothing was registered, so a fan-out cannot fail on a missing callback.
+ ctx.notify_record_start("k")
diff --git a/runtime/src/main/java/org/apache/flink/agents/runtime/ResourceCache.java b/runtime/src/main/java/org/apache/flink/agents/runtime/ResourceCache.java
index cd84da8c3..620c92803 100644
--- a/runtime/src/main/java/org/apache/flink/agents/runtime/ResourceCache.java
+++ b/runtime/src/main/java/org/apache/flink/agents/runtime/ResourceCache.java
@@ -24,13 +24,19 @@ import org.apache.flink.agents.api.resource.python.PythonResourceAdapter;
import org.apache.flink.agents.plan.resourceprovider.PythonResourceProvider;
import org.apache.flink.agents.plan.resourceprovider.ResourceProvider;
import org.apache.flink.agents.plan.tools.FunctionTool;
+import org.apache.flink.agents.runtime.python.utils.PythonActionExecutor;
import org.apache.flink.agents.runtime.resource.ResourceContextImpl;
+import org.apache.flink.agents.runtime.subagent.BaseSubagentSetup;
import org.apache.flink.util.ExceptionUtils;
+import java.util.ArrayList;
import java.util.HashMap;
+import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
+import static org.apache.flink.util.Preconditions.checkState;
+
/**
* Lazily resolves and caches Resource instances from ResourceProviders.
*
@@ -47,6 +53,7 @@ public class ResourceCache implements AutoCloseable {
private final Map<ResourceType, Map<String, ResourceProvider>> resourceProviders;
private final Map<ResourceType, Map<String, Resource>> cache = new ConcurrentHashMap<>();
private volatile PythonResourceAdapter pythonResourceAdapter;
+ private volatile PythonActionExecutor pythonActionExecutor;
private final ResourceContextImpl resourceContext;
/**
@@ -86,6 +93,15 @@ public class ResourceCache implements AutoCloseable {
this.pythonResourceAdapter = adapter;
}
+ /**
+ * Wires the executor that reaches the Python runtime, so the cache can ask that runtime to
+ * materialize the resources it owns. The runtime bridge calls this while the operator opens,
+ * before any resource is resolved.
+ */
+ public void setPythonActionExecutor(PythonActionExecutor pythonActionExecutor) {
+ this.pythonActionExecutor = pythonActionExecutor;
+ }
+
public ResourceContextImpl getResourceContext() {
return resourceContext;
}
@@ -137,6 +153,12 @@ public class ResourceCache implements AutoCloseable {
Resource resource = provider.provide(resourceContext);
+ if (resource instanceof BaseSubagentSetup) {
+ // The framework owns the setup's identity: inject the resource name as its
+ // subagent name.
+ ((BaseSubagentSetup) resource).setSubagentName(name);
+ }
+
if (pythonResourceAdapter != null && resource instanceof FunctionTool) {
((FunctionTool) resource).setPythonResourceAdapter(pythonResourceAdapter);
}
@@ -157,6 +179,52 @@ public class ResourceCache implements AutoCloseable {
cache.computeIfAbsent(type, k -> new ConcurrentHashMap<>()).put(name, resource);
}
+ /**
+ * Eagerly materializes every resource of the given type, wherever it lives. Java-owned
+ * resources are resolved through their provider exactly like a first {@link #getResource}
+ * access, while Python-owned resources are materialized in the Python runtime and represented
+ * by a handle. Every instance is returned and cached, so a later lookup by name resolves to the
+ * same instance. Providers are resolved in no particular order, and resource construction must
+ * not depend on it.
+ *
+ * @param type the resource type to materialize.
+ * @return the materialized resources, empty when the type has none.
+ * @throws IllegalStateException if the type has Python-owned resources while the Python runtime
+ * is unavailable, which leaves them unreachable for the whole job.
+ */
+ public synchronized List<Resource> eagerMaterialize(ResourceType type) throws Exception {
+ Map<String, ResourceProvider> providers = resourceProviders.get(type);
+ List<Resource> materialized = new ArrayList<>();
+ if (providers == null) {
+ return materialized;
+ }
+ boolean hasPythonOwned = false;
+ for (Map.Entry<String, ResourceProvider> entry : providers.entrySet()) {
+ ResourceProvider provider = entry.getValue();
+ if (ResourceProvider.isPythonOwned(provider)) {
+ hasPythonOwned = true;
+ continue;
+ }
+ materialized.add(getResource(entry.getKey(), type));
+ }
+ if (!hasPythonOwned) {
+ return materialized;
+ }
+ checkState(
+ pythonActionExecutor != null,
+ "Resources of type %s are declared in Python but no Python runtime was"
+ + " initialized for this plan, so they cannot be materialized.",
+ type);
+ // The Python runtime owns these resources: it built and opened them, so the handles are
+ // cached as they are instead of being opened again here.
+ for (Map.Entry<String, Resource> handle :
+ pythonActionExecutor.eagerMaterialize(type).entrySet()) {
+ put(handle.getKey(), type, handle.getValue());
+ materialized.add(handle.getValue());
+ }
+ return materialized;
+ }
+
@Override
public void close() throws Exception {
// Close every cached resource, then the resource context, even when an earlier close
diff --git a/runtime/src/main/java/org/apache/flink/agents/runtime/context/RunnerContextImpl.java b/runtime/src/main/java/org/apache/flink/agents/runtime/context/RunnerContextImpl.java
index a40e56b03..49d56e37b 100644
--- a/runtime/src/main/java/org/apache/flink/agents/runtime/context/RunnerContextImpl.java
+++ b/runtime/src/main/java/org/apache/flink/agents/runtime/context/RunnerContextImpl.java
@@ -33,13 +33,13 @@ import org.apache.flink.agents.api.resource.Resource;
import org.apache.flink.agents.api.resource.ResourceType;
import org.apache.flink.agents.api.trace.ExecutionLifecycleEvents;
import org.apache.flink.agents.api.trace.ExecutionReporter;
-import org.apache.flink.agents.api.trace.ExecutionTraceContext;
import org.apache.flink.agents.plan.AgentPlan;
import org.apache.flink.agents.plan.actions.Action;
import org.apache.flink.agents.plan.utils.JsonUtils;
import org.apache.flink.agents.runtime.ResourceCache;
import org.apache.flink.agents.runtime.actionstate.ActionState;
import org.apache.flink.agents.runtime.actionstate.CallResult;
+import org.apache.flink.agents.runtime.lifecycle.ComponentExecutionListener;
import org.apache.flink.agents.runtime.memory.CachedMemoryStore;
import org.apache.flink.agents.runtime.memory.InteranlBaseLongTermMemory;
import org.apache.flink.agents.runtime.memory.MemoryEventBuilder;
@@ -47,8 +47,6 @@ import org.apache.flink.agents.runtime.memory.MemoryEventSettings;
import org.apache.flink.agents.runtime.memory.MemoryObjectImpl;
import org.apache.flink.agents.runtime.memory.MemoryValueObservation;
import org.apache.flink.agents.runtime.metrics.FlinkAgentsMetricGroupImpl;
-import org.apache.flink.agents.runtime.trace.ExecutionEventSink;
-import org.apache.flink.agents.runtime.trace.ReportedExecutionKey;
import org.apache.flink.util.Preconditions;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
@@ -122,7 +120,8 @@ public class RunnerContextImpl implements RunnerContext, ExecutionReporter {
private static final Logger LOG = LoggerFactory.getLogger(RunnerContextImpl.class);
- protected final List<Event> pendingEvents = new ArrayList<>();
+ protected List<Event> pendingEvents = new ArrayList<>();
+
protected final FlinkAgentsMetricGroupImpl agentMetricGroup;
protected final Runnable mailboxThreadChecker;
protected final AgentPlan agentPlan;
@@ -150,9 +149,8 @@ public class RunnerContextImpl implements RunnerContext, ExecutionReporter {
/** Whether the fixed job-level configuration enables any LTM observation. */
private final boolean ltmObservationConfigured;
- @Nullable protected ExecutionTraceContext actionTraceContext;
- @Nullable protected ExecutionEventSink executionEventSink;
- @Nullable private Map<ReportedExecutionKey, ExecutionTraceContext> activeReportedExecutions;
+ /** Component execution listeners of the current action execution, fanned out best-effort. */
+ @Nullable protected List<ComponentExecutionListener> componentExecutionListeners;
/** Context for fine-grained durable execution, may be null if not enabled. */
@Nullable protected DurableExecutionContext durableExecutionContext;
@@ -182,60 +180,24 @@ public class RunnerContextImpl implements RunnerContext, ExecutionReporter {
public void switchActionContext(
String actionName,
MemoryContext memoryContext,
- String contextKey,
- String observationId,
- boolean observationSuppressed) {
- switchActionContext(
- actionName,
- memoryContext,
- contextKey,
- observationId,
- observationSuppressed,
- null,
- null);
- }
-
- public void switchActionContext(
- String actionName,
- MemoryContext memoryContext,
- String contextKey,
- @Nullable ExecutionTraceContext actionTraceContext,
- @Nullable Map<ReportedExecutionKey, ExecutionTraceContext> activeReportedExecutions) {
- switchActionContext(
- actionName,
- memoryContext,
- contextKey,
- null,
- false,
- actionTraceContext,
- activeReportedExecutions);
- }
-
- public void switchActionContext(
- String actionName,
- MemoryContext memoryContext,
+ List<Event> pendingEvents,
String contextKey,
@Nullable String observationId,
boolean observationSuppressed,
- @Nullable ExecutionTraceContext actionTraceContext,
- @Nullable Map<ReportedExecutionKey, ExecutionTraceContext> activeReportedExecutions) {
+ @Nullable List<ComponentExecutionListener> componentExecutionListeners) {
this.actionName = actionName;
this.memoryContext = memoryContext;
+ this.pendingEvents = pendingEvents;
this.contextKey = contextKey;
this.observationId = observationId;
this.observationSuppressed = observationSuppressed;
this.ltmObservationEnabled = !observationSuppressed && ltmObservationConfigured;
- this.actionTraceContext = actionTraceContext;
- this.activeReportedExecutions = activeReportedExecutions;
+ this.componentExecutionListeners = componentExecutionListeners;
if (ltm != null) {
ltm.switchContext(contextKey, observationId, observationSuppressed);
}
}
- public void setExecutionEventSink(@Nullable ExecutionEventSink executionEventSink) {
- this.executionEventSink = executionEventSink;
- }
-
public MemoryContext getMemoryContext() {
return memoryContext;
}
@@ -350,6 +312,15 @@ public class RunnerContextImpl implements RunnerContext, ExecutionReporter {
this.pendingEvents.isEmpty(), "There are pending events remaining in the context.");
}
+ public List<Event> getPendingEvents() {
+ return this.pendingEvents;
+ }
+
+ @Nullable
+ public List<ComponentExecutionListener> getComponentExecutionListeners() {
+ return this.componentExecutionListeners;
+ }
+
public List<MemoryUpdate> getSensoryMemoryUpdates() {
mailboxThreadChecker.run();
return List.copyOf(memoryContext.getSensoryMemoryUpdates());
@@ -404,42 +375,27 @@ public class RunnerContextImpl implements RunnerContext, ExecutionReporter {
ExecutionLifecycleEvents.executionFailed(error, problemCategory));
}
+ /**
+ * Fans the report out to the current action execution's component listeners best-effort: a
+ * listener that throws is logged and skipped, so reporting never fails the caller.
+ */
protected void reportChildExecution(
String entityType, String entityName, Map<String, Object> entityMetadata, Event event) {
mailboxThreadChecker.run();
- if (actionTraceContext == null
- || executionEventSink == null
- || activeReportedExecutions == null) {
+ if (componentExecutionListeners == null) {
return;
}
-
- ReportedExecutionKey key = new ReportedExecutionKey(entityType, entityName, entityMetadata);
- ExecutionTraceContext reportTraceContext;
- if (ExecutionLifecycleEvents.EXECUTION_STARTED_EVENT_TYPE.equals(event.getType())) {
- reportTraceContext =
- actionTraceContext.childExecution(
- entityType, entityName, key.getEntityMetadata());
- ExecutionTraceContext previous = activeReportedExecutions.put(key, reportTraceContext);
- if (previous != null) {
- LOG.debug(
- "Execution start report for {}:{} replaced an active report with the same metadata.",
- entityType,
- entityName);
- }
- } else {
- reportTraceContext = activeReportedExecutions.remove(key);
- if (reportTraceContext == null) {
- LOG.debug(
- "Execution terminal report for {}:{} has no matching start report; emitting it with a new execution id.",
- entityType,
- entityName);
- reportTraceContext =
- actionTraceContext.childExecution(
- entityType, entityName, key.getEntityMetadata());
+ for (ComponentExecutionListener listener : componentExecutionListeners) {
+ try {
+ listener.onComponentExecution(entityType, entityName, entityMetadata, event);
+ } catch (Exception | LinkageError e) {
+ LOG.warn(
+ "Component execution listener {} failed on a report for action '{}' ({})",
+ listener.getClass().getSimpleName(),
+ actionName,
+ e.getClass().getSimpleName());
}
}
-
- executionEventSink.emit(event, reportTraceContext);
}
@Override
diff --git a/runtime/src/main/java/org/apache/flink/agents/runtime/lifecycle/ComponentExecutionListener.java b/runtime/src/main/java/org/apache/flink/agents/runtime/lifecycle/ComponentExecutionListener.java
new file mode 100644
index 000000000..e7099fb66
--- /dev/null
+++ b/runtime/src/main/java/org/apache/flink/agents/runtime/lifecycle/ComponentExecutionListener.java
@@ -0,0 +1,59 @@
+/*
+ * 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.lifecycle;
+
+import org.apache.flink.agents.api.Event;
+import org.apache.flink.agents.api.trace.ExecutionLifecycleEvents;
+
+import java.util.Map;
+
+/**
+ * Observes component executions reported from within an action, at LLM, parser, and tool
+ * granularity.
+ *
+ * <p>A component reports its lifecycle as a status event rather than as one callback per outcome,
+ * so a listener that only cares about a subset matches on the event type and ignores the rest.
+ *
+ * <p>Invariants a listener may rely on, and must not break:
+ *
+ * <ul>
+ * <li>The callback runs on the mailbox thread, so a listener needs no synchronization of its own.
+ * <li>An exception thrown by a listener is logged and swallowed, so reporting never fails the
+ * reporting component and never starves the remaining listeners.
+ * <li>The event carries the lifecycle status only; the reporting component is identified by the
+ * entity triple, which repeats on every report of the same execution.
+ * <li>The event instance is shared with every other listener, so a listener must treat it as
+ * read-only.
+ * </ul>
+ */
+@FunctionalInterface
+public interface ComponentExecutionListener {
+
+ /**
+ * A component execution reported a lifecycle event.
+ *
+ * @param entityType the component entity type, one of {@code
+ * org.apache.flink.agents.api.trace.ExecutionReporter.EntityTypes}.
+ * @param entityName the component entity name.
+ * @param entityMetadata the entity metadata reported with the execution.
+ * @param event the lifecycle event, one of those produced by {@link ExecutionLifecycleEvents}.
+ */
+ void onComponentExecution(
+ String entityType, String entityName, Map<String, Object> entityMetadata, Event event);
+}
diff --git a/runtime/src/main/java/org/apache/flink/agents/runtime/lifecycle/PythonTaskLifecycleListener.java b/runtime/src/main/java/org/apache/flink/agents/runtime/lifecycle/PythonTaskLifecycleListener.java
new file mode 100644
index 000000000..4f95fa622
--- /dev/null
+++ b/runtime/src/main/java/org/apache/flink/agents/runtime/lifecycle/PythonTaskLifecycleListener.java
@@ -0,0 +1,82 @@
+/*
+ * 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.lifecycle;
+
+import org.apache.flink.agents.runtime.operator.ActionTask;
+import org.apache.flink.agents.runtime.python.utils.PythonActionExecutor;
+
+/**
+ * Forwards the operator's per-record and per-action lifecycle to the Python runtime over the pemja
+ * bridge, where the Python side dispatches each callback to its registered listeners. The whole
+ * {@link TaskLifecycleListener} contract is forwarded, so a Python listener observes the same
+ * lifecycle as any Java listener.
+ */
+public final class PythonTaskLifecycleListener implements TaskLifecycleListener {
+
+ private final PythonActionExecutor pythonActionExecutor;
+
+ public PythonTaskLifecycleListener(PythonActionExecutor pythonActionExecutor) {
+ this.pythonActionExecutor = pythonActionExecutor;
+ }
+
+ @Override
+ public void onRecordStart(Object key) {
+ pythonActionExecutor.notifyRecordStart(key);
+ }
+
+ @Override
+ public void onActionPrepared(ActionTask task) {
+ pythonActionExecutor.notifyActionPrepared(task);
+ }
+
+ @Override
+ public void onActionStarted(ActionTask task) {
+ pythonActionExecutor.notifyActionStarted(task);
+ }
+
+ @Override
+ public void onActionTransferred(ActionTask from, ActionTask to) {
+ pythonActionExecutor.notifyActionTransferred(from, to);
+ }
+
+ @Override
+ public void onActionFinishing(ActionTask task) {
+ pythonActionExecutor.notifyActionFinishing(task);
+ }
+
+ @Override
+ public void onActionFinished(ActionTask task) {
+ pythonActionExecutor.notifyActionFinished(task);
+ }
+
+ @Override
+ public void onActionReused(ActionTask task) {
+ pythonActionExecutor.notifyActionReused(task);
+ }
+
+ @Override
+ public void onActionFailed(ActionTask task, Throwable error) {
+ pythonActionExecutor.notifyActionFailed(task, error);
+ }
+
+ @Override
+ public void onRecordFinished(Object key) {
+ pythonActionExecutor.notifyRecordFinished(key);
+ }
+}
diff --git a/runtime/src/main/java/org/apache/flink/agents/runtime/lifecycle/TaskLifecycleListener.java b/runtime/src/main/java/org/apache/flink/agents/runtime/lifecycle/TaskLifecycleListener.java
new file mode 100644
index 000000000..9a3e339fa
--- /dev/null
+++ b/runtime/src/main/java/org/apache/flink/agents/runtime/lifecycle/TaskLifecycleListener.java
@@ -0,0 +1,130 @@
+/*
+ * 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.lifecycle;
+
+import org.apache.flink.agents.runtime.operator.ActionTask;
+
+/**
+ * Observes the per-record and per-action lifecycle of {@code ActionExecutionOperator}.
+ *
+ * <p>All callbacks run on the mailbox thread.
+ *
+ * <p>Event pairing semantics:
+ *
+ * <ul>
+ * <li>{@code onRecordStart}/{@code onRecordFinished} bracket the processing of one input record
+ * for a key; a record that triggers no actions emits neither callback.
+ * <li>{@code onActionPrepared} fires on every preparation of an action task, including
+ * re-preparation of a suspended or resumed task, and pairs with exactly one of the terminal
+ * callbacks: {@code onActionFinishing} followed by {@code onActionFinished} on normal
+ * completion, {@code onActionReused} when a replay skips an already-completed action, {@code
+ * onActionFailed} when the invocation fails, or {@code onActionTransferred} when a
+ * non-finished task hands its contexts over to the task it generated.
+ * <li>{@code onActionStarted} fires at most once per action execution, before the first real
+ * invocation and never on a re-preparation; the gate is checkpointed with the task, so a
+ * failover replay re-emits {@code onRecordStart} for the resumed round but not {@code
+ * onActionStarted}.
+ * </ul>
+ *
+ * <p>Exception contract: the framework allows a listener to inspect state and throw when necessary;
+ * a listener that only observes should avoid throwing.
+ */
+public interface TaskLifecycleListener {
+
+ /**
+ * The first action task of the input record of {@code key} has just been created and is about
+ * to be prepared. Also re-emitted when the task chain of a record that was in flight at
+ * snapshot time resumes after a failover, so listeners observe a paired bracket for the
+ * replayed round.
+ *
+ * @param key the Flink key of the input record starting processing.
+ */
+ default void onRecordStart(Object key) {}
+
+ /**
+ * An action task's runner context has been wired up and the task is ready to run. Fires on
+ * every preparation, including re-preparation of a suspended or resumed task.
+ *
+ * @param task the prepared action task.
+ */
+ default void onActionPrepared(ActionTask task) {}
+
+ /**
+ * An action execution is about to run for the first time. Fires at most once per action
+ * execution: re-preparations of a suspended or resumed task do not emit it again.
+ *
+ * @param task the action task whose first invocation is imminent.
+ */
+ default void onActionStarted(ActionTask task) {}
+
+ /**
+ * A non-finished task handed its per-task contexts to the task it generated. Listeners that
+ * keep per-task bookkeeping must move their entries from {@code from} to {@code to} so the
+ * continuation keeps its state.
+ *
+ * @param from the finishing task whose contexts were transferred.
+ * @param to the generated task that inherited the contexts.
+ */
+ default void onActionTransferred(ActionTask from, ActionTask to) {}
+
+ /**
+ * A task's invocation completed and its contexts record has been removed; fires immediately
+ * before its result (including its completed state) is persisted. Not emitted on the
+ * replay-reuse path.
+ *
+ * @param task the completing action task.
+ */
+ default void onActionFinishing(ActionTask task) {}
+
+ /**
+ * A task's invocation finished normally and its result has been persisted, so a later replay of
+ * the same action skips the invocation. Marks the end of the normal completion path. Not
+ * emitted when the invocation fails.
+ *
+ * @param task the finished action task.
+ */
+ default void onActionFinished(ActionTask task) {}
+
+ /**
+ * A replayed already-completed action had its persisted result applied and its invocation
+ * skipped. This is the sole terminal callback on the reuse path.
+ *
+ * @param task the reused action task.
+ */
+ default void onActionReused(ActionTask task) {}
+
+ /**
+ * An action invocation failed. Purely observational: listeners perceive the failure for
+ * logging, metrics, or bookkeeping cleanup, but must not compensate or decide on rethrowing.
+ *
+ * @param task the failed action task.
+ * @param error the failure thrown by the invocation.
+ */
+ default void onActionFailed(ActionTask task, Throwable error) {}
+
+ /**
+ * Every task spawned by the input record of {@code key} has completed and the record is fully
+ * processed. Implementations must make their per-record cleanup idempotent: after a failover
+ * replay the notification may not be delivered again for records that completed before the
+ * snapshot.
+ *
+ * @param key the Flink key of the finished input record.
+ */
+ default void onRecordFinished(Object key) {}
+}
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 f82d8047a..a42f216f1 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
@@ -21,8 +21,8 @@ import org.apache.flink.agents.api.Event;
import org.apache.flink.agents.api.OutputEvent;
import org.apache.flink.agents.api.agents.AgentExecutionOptions;
import org.apache.flink.agents.api.event.AgentRunBeginEvent;
-import org.apache.flink.agents.api.trace.ExecutionLifecycleEvents;
-import org.apache.flink.agents.api.trace.ExecutionReporter;
+import org.apache.flink.agents.api.resource.Resource;
+import org.apache.flink.agents.api.resource.ResourceType;
import org.apache.flink.agents.api.trace.ExecutionTraceContext;
import org.apache.flink.agents.plan.AgentPlan;
import org.apache.flink.agents.plan.JavaFunction;
@@ -32,6 +32,9 @@ import org.apache.flink.agents.runtime.ResourceCache;
import org.apache.flink.agents.runtime.actionstate.ActionState;
import org.apache.flink.agents.runtime.actionstate.ActionStateStore;
import org.apache.flink.agents.runtime.eventlog.EventLogWriter;
+import org.apache.flink.agents.runtime.lifecycle.ComponentExecutionListener;
+import org.apache.flink.agents.runtime.lifecycle.PythonTaskLifecycleListener;
+import org.apache.flink.agents.runtime.lifecycle.TaskLifecycleListener;
import org.apache.flink.agents.runtime.memory.Mem0LongTermMemory;
import org.apache.flink.agents.runtime.memory.MemoryEventBuilder;
import org.apache.flink.agents.runtime.memory.MemoryObjectImpl;
@@ -39,7 +42,10 @@ import org.apache.flink.agents.runtime.memory.MemoryUpdateReplayer;
import org.apache.flink.agents.runtime.metrics.BuiltInMetrics;
import org.apache.flink.agents.runtime.metrics.FlinkAgentsMetricGroupImpl;
import org.apache.flink.agents.runtime.python.operator.PythonActionTask;
+import org.apache.flink.agents.runtime.python.resource.PythonRuntimeResource;
import org.apache.flink.agents.runtime.python.utils.PythonActionExecutor;
+import org.apache.flink.agents.runtime.trace.EventLogComponentExecutionListener;
+import org.apache.flink.agents.runtime.trace.EventLogTaskLifecycleListener;
import org.apache.flink.agents.runtime.trace.ExecutionEventLogger;
import org.apache.flink.agents.runtime.utils.EventUtil;
import org.apache.flink.annotation.VisibleForTesting;
@@ -141,6 +147,13 @@ public class ActionExecutionOperator<IN, OUT> extends AbstractStreamOperator<OUT
private final boolean pythonKeyIsPickled;
private final boolean agentRunBeginEventEnabled;
+ // Broadcast targets for the per-record/per-action lifecycle events.
+ private transient List<TaskLifecycleListener> taskLifecycleListeners = new ArrayList<>();
+
+ // Broadcast targets for component execution reports, injected per action execution.
+ private transient List<ComponentExecutionListener> componentExecutionListeners =
+ new ArrayList<>();
+
public ActionExecutionOperator(
AgentPlan agentPlan,
Boolean inputIsJava,
@@ -217,6 +230,16 @@ public class ActionExecutionOperator<IN, OUT> extends AbstractStreamOperator<OUT
// runner context created by ActionTaskContextManager.
ltm = pythonBridge.getLongTermMemory();
+ if (taskLifecycleListeners == null) {
+ taskLifecycleListeners = new ArrayList<>();
+ }
+ if (componentExecutionListeners == null) {
+ componentExecutionListeners = new ArrayList<>();
+ }
+
+ registerEventLogListeners();
+ registerSubagentSetups();
+
// init context manager for runner context creation and memory contexts
contextManager =
new ActionTaskContextManager(
@@ -305,8 +328,14 @@ public class ActionExecutionOperator<IN, OUT> extends AbstractStreamOperator<OUT
output.collect(eventRouter.getReusedStreamRecord().replace(outputData));
}
} else {
+ boolean freshRecordRound = false;
if (isInputEvent) {
// If the event is an InputEvent, we mark that the key is currently being processed.
+ if (!stateManager.hasMoreActionTasks()) {
+ // No tasks in flight for this key: this input record starts a fresh record
+ // processing round.
+ freshRecordRound = true;
+ }
stateManager.addProcessingKey(key);
stateManager.initOrIncSequenceNumber();
tryEmitAgentRunBeginEvent(key, contextKey, event, traceContext);
@@ -317,7 +346,16 @@ public class ActionExecutionOperator<IN, OUT> extends AbstractStreamOperator<OUT
if (triggerActions != null && !triggerActions.isEmpty()) {
for (Action triggerAction : triggerActions) {
stateManager.addActionTask(
- createActionTask(key, triggerAction, event, traceContext));
+ createActionTask(
+ key,
+ triggerAction,
+ event,
+ stateManager.getSequenceNumber(),
+ traceContext));
+ if (freshRecordRound) {
+ notifyRecordStart(key);
+ freshRecordRound = false;
+ }
}
}
}
@@ -423,7 +461,8 @@ public class ActionExecutionOperator<IN, OUT> extends AbstractStreamOperator<OUT
stateManager.getShortTermMemState(),
pythonBridge.getPythonRunnerContext(),
ltm,
- executionEventLogger);
+ this::createComponentListeners);
+ notifyActionPrepared(actionTask);
long sequenceNumber = stateManager.getSequenceNumber();
boolean isFinished;
@@ -449,6 +488,7 @@ public class ActionExecutionOperator<IN, OUT> extends AbstractStreamOperator<OUT
actionTask.getRunnerContext().getSensoryMemory(),
actionState.getSensoryMemoryUpdates());
notifyActionReused(actionTask);
+ contextManager.removeContexts(actionTask);
} else {
// Initialize ActionState if not exists, or use existing one for recovery
if (actionState == null) {
@@ -459,8 +499,8 @@ public class ActionExecutionOperator<IN, OUT> extends AbstractStreamOperator<OUT
key, sequenceNumber, actionTask.action, actionTask.event);
}
- notifyActionStarted(actionTask);
try {
+ notifyActionStarted(actionTask);
// Set up durable execution context for fine-grained recovery
durableExecManager.setupDurableExecutionContext(
actionTask, actionState, sequenceNumber);
@@ -483,11 +523,16 @@ public class ActionExecutionOperator<IN, OUT> extends AbstractStreamOperator<OUT
throw new AssertionError("Unreachable after rethrowing action failure");
}
- // Drop task-local contexts after each step; continuations transfer them back.
- contextManager.removeMemoryContext(actionTask);
+ // We remove the contexts record from the map after the task is processed. It
+ // will be recreated by transferContexts below if the action task has a generated
+ // action task, meaning it is not finished.
+ contextManager.removeContexts(actionTask);
durableExecManager.removeDurableContext(actionTask);
- contextManager.removeContinuationContext(actionTask);
- contextManager.removePythonAwaitableRef(actionTask);
+ if (actionTaskResult.isFinished()) {
+ // Notify before persisting the result, so listeners observe the task
+ // before its completion becomes durable.
+ notifyActionFinishing(actionTask);
+ }
durableExecManager.maybePersistTaskResult(
key,
sequenceNumber,
@@ -502,25 +547,15 @@ public class ActionExecutionOperator<IN, OUT> extends AbstractStreamOperator<OUT
notifyActionFinished(actionTask);
}
} catch (Throwable t) {
- try {
- notifyActionFailed(actionTask, t);
- } finally {
- contextManager.completeActionExecution(actionTask);
- }
+ notifyActionFailed(actionTask, t);
ExceptionUtils.rethrowException(t);
// Unreachable; required for Java definite-assignment analysis.
return;
}
}
- try {
- for (Event actionOutputEvent : outputEvents) {
- processEvent(key, contextKey, actionOutputEvent, actionTask.getTraceContext());
- }
- } finally {
- if (isFinished) {
- contextManager.completeActionExecution(actionTask);
- }
+ for (Event actionOutputEvent : outputEvents) {
+ processEvent(key, contextKey, actionOutputEvent, actionTask.getTraceContext());
}
boolean currentInputEventFinished = false;
@@ -542,12 +577,14 @@ public class ActionExecutionOperator<IN, OUT> extends AbstractStreamOperator<OUT
// If the action task is not finished, we keep the contexts in memory for the
// next generated ActionTask to be invoked.
contextManager.transferContexts(actionTask, generatedActionTask, durableExecManager);
+ notifyActionTransferred(actionTask, generatedActionTask);
stateManager.addActionTask(generatedActionTask);
}
// 3. Process the next InputEvent or next action task
if (currentInputEventFinished) {
+ notifyRecordFinished(key);
// Clean up sensory memory when a single run finished.
actionTask.getRunnerContext().clearSensoryMemory();
durableExecManager.updateLastCompletedSequenceNumber(sequenceNumber);
@@ -701,40 +738,141 @@ public class ActionExecutionOperator<IN, OUT> extends AbstractStreamOperator<OUT
if (actionTask.hasExecutionStartedEventEmitted()) {
return;
}
- notifyExecutionLifecycleEvent(
- actionTask.getTraceContext(), ExecutionLifecycleEvents.executionStarted());
+ for (TaskLifecycleListener listener : taskLifecycleListeners) {
+ listener.onActionStarted(actionTask);
+ }
actionTask.markExecutionStartedEventEmitted();
}
- private void notifyActionFinished(ActionTask actionTask) {
- notifyExecutionLifecycleEvent(
- actionTask.getTraceContext(), ExecutionLifecycleEvents.executionFinished());
+ private void registerEventLogListeners() {
+ addTaskLifecycleListener(new EventLogTaskLifecycleListener(executionEventLogger));
+ }
+
+ /**
+ * Builds the component execution listeners of one action execution: the per-execution event log
+ * adapter first, followed by the globally registered listeners.
+ */
+ private List<ComponentExecutionListener> createComponentListeners(ActionTask actionTask) {
+ List<ComponentExecutionListener> listeners = new ArrayList<>();
+ listeners.add(
+ new EventLogComponentExecutionListener(
+ actionTask.getTraceContext(), executionEventLogger));
+ listeners.addAll(componentExecutionListeners);
+ return listeners;
+ }
+
+ /**
+ * Materializes every sub-agent setup, in either language, and registers the ones that observe
+ * the task lifecycle. A Java setup joins this operator's listeners directly; a Python setup
+ * lives in the Python runtime, so it joins the Python runtime's listeners and this operator
+ * notifies them through a single bridge listener.
+ *
+ * <p>Runs while the operator opens, after the Python bridge is up, because the Python runtime
+ * materializes the setups it owns.
+ */
+ private void registerSubagentSetups() throws Exception {
+ boolean pythonSetupRegistered = false;
+ for (Resource setup : resourceCache.eagerMaterialize(ResourceType.AGENT)) {
+ if (setup instanceof PythonRuntimeResource) {
+ pythonSetupRegistered |=
+ pythonBridge
+ .getPythonActionExecutor()
+ .addTaskLifecycleListener(
+ ((PythonRuntimeResource) setup).getPythonResource());
+ } else if (setup instanceof TaskLifecycleListener) {
+ addTaskLifecycleListener((TaskLifecycleListener) setup);
+ }
+ }
+ if (pythonSetupRegistered) {
+ addTaskLifecycleListener(
+ new PythonTaskLifecycleListener(pythonBridge.getPythonActionExecutor()));
+ }
+ }
+
+ /**
+ * Registers a listener to be notified of per-record/per-action lifecycle events. The
+ * registration itself is not part of the operator state, so it must happen before records are
+ * processed.
+ */
+ public void addTaskLifecycleListener(TaskLifecycleListener listener) {
+ taskLifecycleListeners.add(listener);
+ }
+
+ /**
+ * Registers a listener to be notified of component execution reports of every action execution.
+ * The registration itself is not part of the operator state, so it must happen before records
+ * are processed.
+ */
+ public void addComponentExecutionListener(ComponentExecutionListener listener) {
+ componentExecutionListeners.add(listener);
+ }
+
+ private void notifyRecordStart(Object key) {
+ for (TaskLifecycleListener listener : taskLifecycleListeners) {
+ listener.onRecordStart(key);
+ }
+ }
+
+ private void notifyActionPrepared(ActionTask task) {
+ for (TaskLifecycleListener listener : taskLifecycleListeners) {
+ listener.onActionPrepared(task);
+ }
+ }
+
+ private void notifyActionTransferred(ActionTask from, ActionTask to) {
+ for (TaskLifecycleListener listener : taskLifecycleListeners) {
+ listener.onActionTransferred(from, to);
+ }
+ }
+
+ private void notifyActionFinishing(ActionTask task) {
+ for (TaskLifecycleListener listener : taskLifecycleListeners) {
+ listener.onActionFinishing(task);
+ }
+ }
+
+ private void notifyActionFinished(ActionTask task) {
+ for (TaskLifecycleListener listener : taskLifecycleListeners) {
+ listener.onActionFinished(task);
+ }
}
- private void notifyActionReused(ActionTask actionTask) {
- notifyExecutionLifecycleEvent(
- actionTask.getTraceContext(), ExecutionLifecycleEvents.executionReused());
+ private void notifyActionReused(ActionTask task) {
+ for (TaskLifecycleListener listener : taskLifecycleListeners) {
+ listener.onActionReused(task);
+ }
}
- private void notifyActionFailed(ActionTask actionTask, Throwable error) {
- notifyExecutionLifecycleEvent(
- actionTask.getTraceContext(),
- ExecutionLifecycleEvents.executionFailed(
- error, ExecutionReporter.ProblemCategories.ACTION_EXECUTION_FAILED));
+ private void notifyActionFailed(ActionTask task, Throwable error) {
+ for (TaskLifecycleListener listener : taskLifecycleListeners) {
+ try {
+ listener.onActionFailed(task, error);
+ } catch (Throwable listenerError) {
+ if (listenerError != error) {
+ error.addSuppressed(listenerError);
+ }
+ }
+ }
}
- private void notifyExecutionLifecycleEvent(ExecutionTraceContext traceContext, Event event) {
- executionEventLogger.emit(event, traceContext);
+ private void notifyRecordFinished(Object key) {
+ for (TaskLifecycleListener listener : taskLifecycleListeners) {
+ listener.onRecordFinished(key);
+ }
}
private ActionTask createActionTask(
- Object key, Action action, Event event, ExecutionTraceContext sourceTraceContext) {
+ Object key,
+ Action action,
+ Event event,
+ long sequenceNumber,
+ ExecutionTraceContext sourceTraceContext) {
ExecutionTraceContext actionTraceContext =
ExecutionTraceContext.forAction(sourceTraceContext, action.getName());
if (action.getExec() instanceof JavaFunction) {
- return new JavaActionTask(key, event, action, actionTraceContext);
+ return new JavaActionTask(key, event, action, sequenceNumber, actionTraceContext);
} else if (action.getExec() instanceof PythonFunction) {
- return new PythonActionTask(key, event, action, actionTraceContext);
+ return new PythonActionTask(key, event, action, sequenceNumber, actionTraceContext);
} else {
throw new IllegalStateException(
"Unsupported action type: " + action.getExec().getClass());
@@ -781,6 +919,9 @@ public class ActionExecutionOperator<IN, OUT> extends AbstractStreamOperator<OUT
}
eventRouter.getKeySegmentQueue().addKeyToLastSegment(key);
String contextKey = resolveContextKey(key);
+ // Align with the task-level replay: re-emit the record start for the resumed
+ // round so listeners observe a paired start/finished bracket as well.
+ notifyRecordStart(key);
mailboxExecutor.submit(
() -> tryProcessActionTaskForKey(key, contextKey), "process action task");
}
diff --git a/runtime/src/main/java/org/apache/flink/agents/runtime/operator/ActionTask.java b/runtime/src/main/java/org/apache/flink/agents/runtime/operator/ActionTask.java
index 4a1ee3716..1831a001b 100644
--- a/runtime/src/main/java/org/apache/flink/agents/runtime/operator/ActionTask.java
+++ b/runtime/src/main/java/org/apache/flink/agents/runtime/operator/ActionTask.java
@@ -59,6 +59,13 @@ public abstract class ActionTask implements Serializable {
protected final ExecutionTraceContext traceContext;
private boolean executionStartedEventEmitted;
+ /**
+ * The sequence number of the input record that triggered this task, counted per key by {@link
+ * OperatorStateManager#initOrIncSequenceNumber}. Every task generated while processing one
+ * record inherits the same number, so the number identifies the record rather than the task.
+ */
+ protected final long sequenceNumber;
+
/**
* Since RunnerContextImpl contains references to the Operator and state, it should not be
* serialized and included in the state with ActionTask. Instead, we should check if a valid
@@ -66,40 +73,49 @@ public abstract class ActionTask implements Serializable {
*/
protected transient RunnerContextImpl runnerContext;
- public ActionTask(Object key, Event event, Action action) {
+ public ActionTask(Object key, Event event, Action action, long sequenceNumber) {
this(
key,
event,
action,
+ sequenceNumber,
UUID.randomUUID().toString(),
ExecutionTraceContext.forExecution(
null, null, null, ExecutionReporter.EntityTypes.ACTION, action.getName()));
}
- protected ActionTask(Object key, Event event, Action action, String observationId) {
+ protected ActionTask(
+ Object key, Event event, Action action, long sequenceNumber, String observationId) {
this(
key,
event,
action,
+ sequenceNumber,
observationId,
ExecutionTraceContext.forExecution(
null, null, null, ExecutionReporter.EntityTypes.ACTION, action.getName()));
}
protected ActionTask(
- Object key, Event event, Action action, ExecutionTraceContext traceContext) {
- this(key, event, action, UUID.randomUUID().toString(), traceContext);
+ Object key,
+ Event event,
+ Action action,
+ long sequenceNumber,
+ ExecutionTraceContext traceContext) {
+ this(key, event, action, sequenceNumber, UUID.randomUUID().toString(), traceContext);
}
protected ActionTask(
Object key,
Event event,
Action action,
+ long sequenceNumber,
String observationId,
ExecutionTraceContext traceContext) {
this.key = key;
this.event = event;
this.action = action;
+ this.sequenceNumber = sequenceNumber;
this.observationId = Objects.requireNonNull(observationId, "observationId");
this.traceContext = Objects.requireNonNull(traceContext, "traceContext must not be null");
}
@@ -116,6 +132,18 @@ public abstract class ActionTask implements Serializable {
return key;
}
+ public Event getEvent() {
+ return event;
+ }
+
+ public Action getAction() {
+ return action;
+ }
+
+ public long getSequenceNumber() {
+ return sequenceNumber;
+ }
+
public String getObservationId() {
if (observationId == null) {
// Tasks restored from state written before observation IDs were introduced have no
@@ -125,7 +153,7 @@ public abstract class ActionTask implements Serializable {
return observationId;
}
- ExecutionTraceContext getTraceContext() {
+ public ExecutionTraceContext getTraceContext() {
return traceContext;
}
diff --git a/runtime/src/main/java/org/apache/flink/agents/runtime/operator/ActionTaskContextManager.java b/runtime/src/main/java/org/apache/flink/agents/runtime/operator/ActionTaskContextManager.java
index 64e9e9ed9..fb107a878 100644
--- a/runtime/src/main/java/org/apache/flink/agents/runtime/operator/ActionTaskContextManager.java
+++ b/runtime/src/main/java/org/apache/flink/agents/runtime/operator/ActionTaskContextManager.java
@@ -17,8 +17,8 @@
*/
package org.apache.flink.agents.runtime.operator;
+import org.apache.flink.agents.api.Event;
import org.apache.flink.agents.api.event.MemoryEvent;
-import org.apache.flink.agents.api.trace.ExecutionTraceContext;
import org.apache.flink.agents.plan.AgentPlan;
import org.apache.flink.agents.plan.JavaFunction;
import org.apache.flink.agents.plan.PythonFunction;
@@ -27,20 +27,23 @@ import org.apache.flink.agents.runtime.async.ContinuationActionExecutor;
import org.apache.flink.agents.runtime.async.ContinuationContext;
import org.apache.flink.agents.runtime.context.JavaRunnerContextImpl;
import org.apache.flink.agents.runtime.context.RunnerContextImpl;
+import org.apache.flink.agents.runtime.lifecycle.ComponentExecutionListener;
import org.apache.flink.agents.runtime.memory.CachedMemoryStore;
import org.apache.flink.agents.runtime.memory.InteranlBaseLongTermMemory;
import org.apache.flink.agents.runtime.memory.MemoryObjectImpl;
import org.apache.flink.agents.runtime.metrics.FlinkAgentsMetricGroupImpl;
import org.apache.flink.agents.runtime.python.context.PythonRunnerContextImpl;
-import org.apache.flink.agents.runtime.trace.ExecutionEventSink;
-import org.apache.flink.agents.runtime.trace.ReportedExecutionKey;
import org.apache.flink.api.common.state.MapState;
import org.apache.flink.util.ExceptionUtils;
+import org.apache.flink.util.Preconditions;
import javax.annotation.Nullable;
+import java.util.ArrayList;
import java.util.HashMap;
+import java.util.List;
import java.util.Map;
+import java.util.function.Function;
/**
* Owns the per-{@link ActionTask} runtime context bookkeeping for {@link ActionExecutionOperator}.
@@ -50,47 +53,80 @@ import java.util.Map;
* <ul>
* <li>The shared (Java) {@link RunnerContextImpl} that is reused across action tasks via {@link
* RunnerContextImpl#switchActionContext}.
- * <li>Three per-{@link ActionTask} maps that survive across the boundary between one task and its
- * generated continuation task: memory contexts, continuation contexts (for async Java
- * actions), and Python awaitable references.
- * <li>Active child-execution reports, keyed by Action execution id, that pair start and terminal
- * reports across continuation tasks without entering Flink state.
+ * <li>A single per-{@link ActionTask} contexts record ({@link ActionTaskContexts}) that survives
+ * across the boundary between a finishing action and the action it generates: memory context,
+ * continuation context (for async Java actions), the Python awaitable reference, and the
+ * component execution listeners, created, transferred, and removed as one unit.
* <li>The {@link ContinuationActionExecutor} thread pool used to run async Java continuations.
* </ul>
*
- * <p>Lifecycle: instantiated by the operator's {@code open()} with the configured async-thread
- * count from the agent plan. Has no separate {@code open()} step — fully constructed in the
- * operator's {@code open()}. {@link #close()} closes the shared runner context and the continuation
- * executor.
+ * <p>The manager is fully constructed in the operator's {@code open()} with the configured
+ * async-thread count from the agent plan, so it has no separate open step.
*
- * <p>Note: the Python {@link RunnerContextImpl} is not owned here — it is owned by {@link
- * PythonBridgeManager} and passed in as a parameter to {@link #createOrGetRunnerContext} and {@link
- * #createAndSetRunnerContext}. The durable-execution context map likewise lives on {@link
- * DurableExecutionManager} and is accessed via the manager parameter passed to {@link
- * #transferContexts}.
- *
- * <p>Design constraint: package-private; no manager-to-manager held references. Cross-cutting data
- * flows via method parameters.
+ * <p>No manager-to-manager references are held here, so cross-cutting data flows in as method
+ * parameters. The Python {@link RunnerContextImpl} stays owned by {@link PythonBridgeManager} and
+ * the durable-execution context stays on {@link DurableExecutionManager}, and both are passed in
+ * when a method needs them.
*/
class ActionTaskContextManager implements AutoCloseable {
private RunnerContextImpl runnerContext;
- private final Map<ActionTask, RunnerContextImpl.MemoryContext> actionTaskMemoryContexts;
- private final Map<ActionTask, ContinuationContext> continuationContexts;
- private final Map<ActionTask, String> pythonAwaitableRefs;
- private final Map<String, Map<ReportedExecutionKey, ExecutionTraceContext>>
- activeReportedExecutionsByActionExecutionId;
- private final ContinuationActionExecutor continuationActionExecutor;
+ private final Map<ActionTask, ActionTaskContexts> actionTaskContexts;
+
+ private ContinuationActionExecutor continuationActionExecutor;
ActionTaskContextManager(int numAsyncThreads) {
- this.actionTaskMemoryContexts = new HashMap<>();
- this.continuationContexts = new HashMap<>();
- this.pythonAwaitableRefs = new HashMap<>();
- this.activeReportedExecutionsByActionExecutionId = new HashMap<>();
+ this.actionTaskContexts = new HashMap<>();
this.continuationActionExecutor = new ContinuationActionExecutor(numAsyncThreads);
}
+ /**
+ * Mutable holder for every per-task context except durable execution. The pending output events
+ * live here rather than in the memory context because they are an output buffer, not memory.
+ */
+ private static final class ActionTaskContexts {
+ @Nullable private RunnerContextImpl.MemoryContext memoryContext;
+ @Nullable private ContinuationContext continuationContext;
+ @Nullable private String pythonAwaitableRef;
+ private List<Event> pendingEvents = new ArrayList<>();
+ @Nullable private List<ComponentExecutionListener> componentListeners;
+ }
+
+ private boolean hasContexts(ActionTask actionTask) {
+ return actionTaskContexts.containsKey(actionTask);
+ }
+
+ /**
+ * Explicitly creates the single contexts record for a task. Fails if one already exists so that
+ * creation is always intentional and destroyed contexts can never be silently resurrected by a
+ * stray mutator call.
+ */
+ void createContexts(ActionTask actionTask) {
+ Preconditions.checkState(
+ !actionTaskContexts.containsKey(actionTask),
+ "Contexts already exist for action task");
+ actionTaskContexts.put(actionTask, new ActionTaskContexts());
+ }
+
+ /**
+ * Returns the existing contexts record for a task, failing fast if it was never created or
+ * removed.
+ */
+ private ActionTaskContexts requireContexts(ActionTask actionTask) {
+ return Preconditions.checkNotNull(
+ actionTaskContexts.get(actionTask), "Missing contexts for action task");
+ }
+
+ /**
+ * Removes the whole per-task contexts record as one unit. Fails if there is nothing to remove.
+ */
+ void removeContexts(ActionTask actionTask) {
+ Preconditions.checkState(
+ actionTaskContexts.remove(actionTask) != null,
+ "No contexts to remove for action task");
+ }
+
/**
* Returns a runner context for an action's exec language.
*
@@ -156,9 +192,10 @@ class ActionTaskContextManager implements AutoCloseable {
* <li>Selects a Java or Python runner context based on the action's {@code Exec} type.
* <li>Reuses any existing {@link RunnerContextImpl.MemoryContext} for this task; otherwise
* builds a fresh one backed by the supplied sensory/short-term memory states.
- * <li>Wires the runtime-level execution event sink onto the runner context.
+ * <li>Creates or reuses the per-action-execution component listener list and wires it onto
+ * the runner context.
* <li>Calls {@link RunnerContextImpl#switchActionContext} so the shared context now points at
- * this action's name, memory, key namespace, trace context, and reported-execution state.
+ * this action's name, memory, key namespace, and component listener list.
* <li>For Java contexts, attaches a continuation context (re-used if the task is resuming
* from an async suspend, fresh otherwise).
* <li>For Python contexts, attaches the per-task awaitable reference (or {@code null} if the
@@ -189,7 +226,15 @@ class ActionTaskContextManager implements AutoCloseable {
MapState<String, MemoryObjectImpl.MemoryItem> shortTermMemState,
PythonRunnerContextImpl pythonRunnerContext,
@Nullable InteranlBaseLongTermMemory longTermMemory,
- @Nullable ExecutionEventSink executionEventSink) {
+ @Nullable
+ Function<ActionTask, List<ComponentExecutionListener>>
+ componentListenerFactory) {
+ if (!hasContexts(actionTask)) {
+ // First preparation of a root task materializes its contexts. Re-preparations of a
+ // suspended task, or preparation of a generated successor, already have one (created by
+ // transferContexts), so we never recreate here.
+ createContexts(actionTask);
+ }
RunnerContextImpl context;
if (actionTask.action.getExec() instanceof JavaFunction) {
context =
@@ -217,26 +262,24 @@ class ActionTaskContextManager implements AutoCloseable {
throw new IllegalStateException(
"Unsupported action type: " + actionTask.action.getExec().getClass());
}
- context.setExecutionEventSink(executionEventSink);
- RunnerContextImpl.MemoryContext memoryContext;
- if (actionTaskMemoryContexts.containsKey(actionTask)) {
- memoryContext = actionTaskMemoryContexts.get(actionTask);
- } else {
+ RunnerContextImpl.MemoryContext memoryContext = getMemoryContext(actionTask);
+ if (memoryContext == null) {
memoryContext =
new RunnerContextImpl.MemoryContext(
new CachedMemoryStore(sensoryMemState),
new CachedMemoryStore(shortTermMemState));
+ putMemoryContext(actionTask, memoryContext);
}
context.switchActionContext(
actionTask.action.getName(),
memoryContext,
+ requireContexts(actionTask).pendingEvents,
contextKey,
actionTask.getObservationId(),
MemoryEvent.isMemoryType(actionTask.event.getType()),
- actionTask.getTraceContext(),
- getOrCreateActiveReportedExecutions(actionTask));
+ getOrCreateComponentListeners(actionTask, componentListenerFactory));
if (context instanceof JavaRunnerContextImpl) {
ContinuationContext continuationContext;
@@ -246,6 +289,7 @@ class ActionTaskContextManager implements AutoCloseable {
continuationContext = this.getContinuationContext(actionTask);
} else {
continuationContext = new ContinuationContext();
+ putContinuationContext(actionTask, continuationContext);
}
((JavaRunnerContextImpl) context).setContinuationContext(continuationContext);
}
@@ -260,22 +304,19 @@ class ActionTaskContextManager implements AutoCloseable {
private void putMemoryContext(
ActionTask actionTask, RunnerContextImpl.MemoryContext memoryContext) {
- actionTaskMemoryContexts.put(actionTask, memoryContext);
+ requireContexts(actionTask).memoryContext = memoryContext;
}
@Nullable
- RunnerContextImpl.MemoryContext removeMemoryContext(ActionTask actionTask) {
- return actionTaskMemoryContexts.remove(actionTask);
+ private RunnerContextImpl.MemoryContext getMemoryContext(ActionTask actionTask) {
+ return requireContexts(actionTask).memoryContext;
}
/**
* Transfers per-task contexts from a finishing action task to the action task it generated.
*
* <p>Always transfers the memory context. For Java tasks, transfers the continuation context.
- * For Python tasks, transfers the awaitable reference when present. The durable-execution
- * context map lives on {@link DurableExecutionManager}, so that manager is passed in as a
- * parameter rather than held as a field — this keeps the no-manager-to-manager-references
- * design constraint intact.
+ * For Python tasks, transfers the awaitable reference when present.
*
* @param fromTask the finishing task whose contexts should be transferred.
* @param toTask the newly generated task that will inherit the contexts.
@@ -283,8 +324,18 @@ class ActionTaskContextManager implements AutoCloseable {
*/
void transferContexts(
ActionTask fromTask, ActionTask toTask, DurableExecutionManager durableExecManager) {
+ createContexts(toTask);
putMemoryContext(toTask, fromTask.getRunnerContext().getMemoryContext());
toTask.inheritLifecycleState(fromTask);
+ // Share the finishing task's live buffer, which is sourced from its runner context and
+ // outlives the removed contexts, so events emitted before a suspend survive into the
+ // generated task.
+ requireContexts(toTask).pendingEvents = fromTask.getRunnerContext().getPendingEvents();
+ // Carry over the execution's very listener instances: one that pairs a component's start
+ // report with its terminal report keeps that pairing in itself, so rebuilding them here
+ // would orphan the reports of components that started before the suspend.
+ requireContexts(toTask).componentListeners =
+ fromTask.getRunnerContext().getComponentExecutionListeners();
RunnerContextImpl.DurableExecutionContext durableContext =
fromTask.getRunnerContext().getDurableExecutionContext();
if (durableContext != null) {
@@ -304,51 +355,45 @@ class ActionTaskContextManager implements AutoCloseable {
}
}
- void completeActionExecution(ActionTask actionTask) {
- activeReportedExecutionsByActionExecutionId.remove(
- actionTask.getTraceContext().getExecutionId());
- }
-
- private Map<ReportedExecutionKey, ExecutionTraceContext> getOrCreateActiveReportedExecutions(
- ActionTask actionTask) {
- String executionId = actionTask.getTraceContext().getExecutionId();
- if (executionId == null) {
- throw new IllegalStateException("Action execution id must not be null.");
+ @Nullable
+ private List<ComponentExecutionListener> getOrCreateComponentListeners(
+ ActionTask actionTask,
+ @Nullable
+ Function<ActionTask, List<ComponentExecutionListener>>
+ componentListenerFactory) {
+ if (componentListenerFactory == null) {
+ return null;
}
- return activeReportedExecutionsByActionExecutionId.computeIfAbsent(
- executionId, ignored -> new HashMap<>());
+ ActionTaskContexts contexts = requireContexts(actionTask);
+ if (contexts.componentListeners == null) {
+ contexts.componentListeners = componentListenerFactory.apply(actionTask);
+ }
+ return contexts.componentListeners;
}
@Nullable
ContinuationContext getContinuationContext(ActionTask actionTask) {
- return continuationContexts.get(actionTask);
+ return requireContexts(actionTask).continuationContext;
}
void putContinuationContext(ActionTask actionTask, ContinuationContext context) {
- continuationContexts.put(actionTask, context);
- }
-
- void removeContinuationContext(ActionTask actionTask) {
- continuationContexts.remove(actionTask);
+ requireContexts(actionTask).continuationContext = context;
}
boolean hasContinuationContext(ActionTask actionTask) {
- return continuationContexts.containsKey(actionTask);
+ return getContinuationContext(actionTask) != null;
}
@Nullable
String getPythonAwaitableRef(ActionTask actionTask) {
- return pythonAwaitableRefs.get(actionTask);
+ return requireContexts(actionTask).pythonAwaitableRef;
}
void putPythonAwaitableRef(ActionTask actionTask, String ref) {
- pythonAwaitableRefs.put(actionTask, ref);
- }
-
- void removePythonAwaitableRef(ActionTask actionTask) {
- pythonAwaitableRefs.remove(actionTask);
+ requireContexts(actionTask).pythonAwaitableRef = ref;
}
+ /** Closes the shared runner context and the continuation executor. */
@Override
public void close() throws Exception {
// Close the continuation executor even when the runner context fails to close. The first
diff --git a/runtime/src/main/java/org/apache/flink/agents/runtime/operator/JavaActionTask.java b/runtime/src/main/java/org/apache/flink/agents/runtime/operator/JavaActionTask.java
index 4523d9be0..702d82a8e 100644
--- a/runtime/src/main/java/org/apache/flink/agents/runtime/operator/JavaActionTask.java
+++ b/runtime/src/main/java/org/apache/flink/agents/runtime/operator/JavaActionTask.java
@@ -41,14 +41,18 @@ public class JavaActionTask extends ActionTask {
private boolean executionStarted = false;
- public JavaActionTask(Object key, Event event, Action action) {
- super(key, event, action);
+ public JavaActionTask(Object key, Event event, Action action, long sequenceNumber) {
+ super(key, event, action, sequenceNumber);
checkState(action.getExec() instanceof JavaFunction);
}
public JavaActionTask(
- Object key, Event event, Action action, ExecutionTraceContext traceContext) {
- super(key, event, action, traceContext);
+ Object key,
+ Event event,
+ Action action,
+ long sequenceNumber,
+ ExecutionTraceContext traceContext) {
+ super(key, event, action, sequenceNumber, traceContext);
checkState(action.getExec() instanceof JavaFunction);
}
diff --git a/runtime/src/main/java/org/apache/flink/agents/runtime/operator/PythonBridgeManager.java b/runtime/src/main/java/org/apache/flink/agents/runtime/operator/PythonBridgeManager.java
index ba44a889a..ba5b77c74 100644
--- a/runtime/src/main/java/org/apache/flink/agents/runtime/operator/PythonBridgeManager.java
+++ b/runtime/src/main/java/org/apache/flink/agents/runtime/operator/PythonBridgeManager.java
@@ -21,7 +21,7 @@ import org.apache.flink.agents.api.memory.LongTermMemoryOptions;
import org.apache.flink.agents.plan.AgentPlan;
import org.apache.flink.agents.plan.JavaFunction;
import org.apache.flink.agents.plan.PythonFunction;
-import org.apache.flink.agents.plan.resourceprovider.PythonResourceProvider;
+import org.apache.flink.agents.plan.resourceprovider.ResourceProvider;
import org.apache.flink.agents.runtime.PythonMCPResourceDiscovery;
import org.apache.flink.agents.runtime.ResourceCache;
import org.apache.flink.agents.runtime.env.EmbeddedPythonEnvironment;
@@ -58,7 +58,8 @@ import static org.apache.flink.agents.plan.actions.Utils.supportAsync;
* <ul>
* <li>The {@link PythonEnvironmentManager} that prepares dependencies and the Pemja runtime.
* <li>The {@link PythonInterpreter} obtained from that environment.
- * <li>The {@link PythonActionExecutor} (when the plan contains Python actions or Mem0).
+ * <li>The {@link PythonActionExecutor} (when the plan contains Python actions, Python-owned
+ * resources, or Mem0).
* <li>The {@link PythonRunnerContextImpl} consumed by Python actions.
* <li>The Java/Python resource adapters that bridge resource lookups across languages.
* <li>The Java wrapper around Python Mem0 long-term memory (when configured).
@@ -96,15 +97,17 @@ class PythonBridgeManager implements AutoCloseable {
/**
* Initializes the Python runtime if the agent plan needs it.
*
- * <p>Scans the agent plan for any {@link PythonFunction} action or {@link
- * PythonResourceProvider}. If neither is present, this method is a no-op and {@link
+ * <p>Scans the agent plan for any {@link PythonFunction} action or Python-owned resource
+ * provider. If neither is present and Mem0 is not configured, this method is a no-op and {@link
* #isInitialized()} stays {@code false}. Otherwise it builds the {@link
* PythonEnvironmentManager}, opens an embedded {@link PythonInterpreter}, refreshes the shared
* import state for the current dependency generation, constructs the shared {@link
* PythonRunnerContextImpl}, wires the Java/Python resource adapters, and conditionally
- * initializes the Python action executor and the Python resource adapter (each only when the
- * corresponding component is present in the plan). The generation guard runs immediately after
- * interpreter construction and before any user module import.
+ * initializes the Python resource adapter (when Python-owned resources or Mem0 are present) and
+ * the Python action executor (when Python actions, Python-owned resources, or Mem0 are present,
+ * since the executor is also the bridge that materializes Python-owned resources). The
+ * generation guard runs immediately after interpreter construction and before any user module
+ * import.
*
* @param agentPlan the agent plan describing actions and resources.
* @param resourceCache the resource cache visible to both languages.
@@ -140,11 +143,7 @@ class PythonBridgeManager implements AutoCloseable {
.anyMatch(
resourceProviderMap ->
resourceProviderMap.values().stream()
- .anyMatch(
- resourceProvider ->
- resourceProvider
- instanceof
- PythonResourceProvider));
+ .anyMatch(ResourceProvider::isPythonOwned));
boolean mem0Configured = isMem0Configured(agentPlan);
@@ -189,8 +188,9 @@ class PythonBridgeManager implements AutoCloseable {
if (containPythonResource || mem0Configured) {
initPythonResourceAdapter(agentPlan, resourceCache);
}
- if (containPythonAction || mem0Configured) {
+ if (containPythonAction || containPythonResource || mem0Configured) {
initPythonActionExecutor(agentPlan, jobIdentifier);
+ resourceCache.setPythonActionExecutor(pythonActionExecutor);
}
if (mem0Configured) {
wireLongTermMemory(agentPlan, mailboxThreadChecker);
@@ -289,8 +289,8 @@ class PythonBridgeManager implements AutoCloseable {
}
/**
- * @return the Python action executor, or {@code null} if the agent plan contains no Python
- * actions (or {@link #open} has not yet been called).
+ * @return the Python action executor, or {@code null} if the agent plan contains neither Python
+ * actions nor Python-owned resources (or {@link #open} has not yet been called).
*/
@Nullable
PythonActionExecutor getPythonActionExecutor() {
diff --git a/runtime/src/main/java/org/apache/flink/agents/runtime/python/operator/PythonActionTask.java b/runtime/src/main/java/org/apache/flink/agents/runtime/python/operator/PythonActionTask.java
index 65399ab7c..32088a3d1 100644
--- a/runtime/src/main/java/org/apache/flink/agents/runtime/python/operator/PythonActionTask.java
+++ b/runtime/src/main/java/org/apache/flink/agents/runtime/python/operator/PythonActionTask.java
@@ -35,19 +35,24 @@ import static org.apache.flink.util.Preconditions.checkState;
*/
public class PythonActionTask extends ActionTask {
- public PythonActionTask(Object key, Event event, Action action) {
- super(key, event, action);
+ public PythonActionTask(Object key, Event event, Action action, long sequenceNumber) {
+ super(key, event, action, sequenceNumber);
checkState(action.getExec() instanceof PythonFunction);
}
- protected PythonActionTask(Object key, Event event, Action action, String observationId) {
- super(key, event, action, observationId);
+ protected PythonActionTask(
+ Object key, Event event, Action action, long sequenceNumber, String observationId) {
+ super(key, event, action, sequenceNumber, observationId);
checkState(action.getExec() instanceof PythonFunction);
}
public PythonActionTask(
- Object key, Event event, Action action, ExecutionTraceContext traceContext) {
- super(key, event, action, traceContext);
+ Object key,
+ Event event,
+ Action action,
+ long sequenceNumber,
+ ExecutionTraceContext traceContext) {
+ super(key, event, action, sequenceNumber, traceContext);
checkState(action.getExec() instanceof PythonFunction);
}
@@ -55,9 +60,10 @@ public class PythonActionTask extends ActionTask {
Object key,
Event event,
Action action,
+ long sequenceNumber,
String observationId,
ExecutionTraceContext traceContext) {
- super(key, event, action, observationId, traceContext);
+ super(key, event, action, sequenceNumber, observationId, traceContext);
checkState(action.getExec() instanceof PythonFunction);
}
@@ -81,7 +87,7 @@ public class PythonActionTask extends ActionTask {
((PythonRunnerContextImpl) runnerContext).setPythonAwaitableRef(pythonAwaitableRef);
ActionTask tempGeneratedActionTask =
new PythonGeneratorActionTask(
- key, event, action, getObservationId(), traceContext);
+ key, event, action, sequenceNumber, getObservationId(), traceContext);
tempGeneratedActionTask.setRunnerContext(runnerContext);
return tempGeneratedActionTask.invoke(userCodeClassLoader, executor);
}
diff --git a/runtime/src/main/java/org/apache/flink/agents/runtime/python/operator/PythonGeneratorActionTask.java b/runtime/src/main/java/org/apache/flink/agents/runtime/python/operator/PythonGeneratorActionTask.java
index 3b0f1a1e6..3f03fedf3 100644
--- a/runtime/src/main/java/org/apache/flink/agents/runtime/python/operator/PythonGeneratorActionTask.java
+++ b/runtime/src/main/java/org/apache/flink/agents/runtime/python/operator/PythonGeneratorActionTask.java
@@ -27,22 +27,28 @@ import org.apache.flink.agents.runtime.python.utils.PythonActionExecutor;
/** An {@link ActionTask} wrapper a Python awaitable to represent a code block in Python action. */
public class PythonGeneratorActionTask extends PythonActionTask {
- public PythonGeneratorActionTask(Object key, Event event, Action action, String observationId) {
- super(key, event, action, observationId);
+ public PythonGeneratorActionTask(
+ Object key, Event event, Action action, long sequenceNumber, String observationId) {
+ super(key, event, action, sequenceNumber, observationId);
}
public PythonGeneratorActionTask(
- Object key, Event event, Action action, ExecutionTraceContext traceContext) {
- super(key, event, action, traceContext);
+ Object key,
+ Event event,
+ Action action,
+ long sequenceNumber,
+ ExecutionTraceContext traceContext) {
+ super(key, event, action, sequenceNumber, traceContext);
}
public PythonGeneratorActionTask(
Object key,
Event event,
Action action,
+ long sequenceNumber,
String observationId,
ExecutionTraceContext traceContext) {
- super(key, event, action, observationId, traceContext);
+ super(key, event, action, sequenceNumber, observationId, traceContext);
}
@Override
@@ -63,7 +69,8 @@ public class PythonGeneratorActionTask extends PythonActionTask {
+ "re-executing from beginning.",
action.getName());
PythonActionTask freshTask =
- new PythonActionTask(key, event, action, getObservationId(), traceContext);
+ new PythonActionTask(
+ key, event, action, sequenceNumber, getObservationId(), traceContext);
freshTask.setRunnerContext(runnerContext);
return freshTask.invoke(userCodeClassLoader, executor);
}
diff --git a/runtime/src/main/java/org/apache/flink/agents/runtime/python/resource/PythonRuntimeResource.java b/runtime/src/main/java/org/apache/flink/agents/runtime/python/resource/PythonRuntimeResource.java
new file mode 100644
index 000000000..bbe5295ac
--- /dev/null
+++ b/runtime/src/main/java/org/apache/flink/agents/runtime/python/resource/PythonRuntimeResource.java
@@ -0,0 +1,57 @@
+/*
+ * 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.python.resource;
+
+import org.apache.flink.agents.api.resource.Resource;
+import org.apache.flink.agents.api.resource.ResourceType;
+import pemja.core.object.PyObject;
+
+/**
+ * Java-side handle to a resource that lives in the Python runtime, letting Java code reach a
+ * resource it cannot construct itself.
+ *
+ * <p>The Python runtime owns the resource: it constructs it, keeps it in its own cache and closes
+ * it. This handle is therefore non-owning — {@link #open()} and {@link #close()} deliberately do
+ * nothing, because opening or closing the same Python resource a second time from Java would break
+ * the invariants its owner already established.
+ *
+ * <p>The handle carries no behaviour of its own, because what a Python resource can do is expressed
+ * in Python: a caller that needs more than the resource's type drives the Python object from {@link
+ * #getPythonResource()} over the bridge.
+ */
+public final class PythonRuntimeResource extends Resource {
+
+ private final ResourceType type;
+ private final PyObject pythonResource;
+
+ public PythonRuntimeResource(ResourceType type, PyObject pythonResource) {
+ this.type = type;
+ this.pythonResource = pythonResource;
+ }
+
+ @Override
+ public ResourceType getResourceType() {
+ return type;
+ }
+
+ /** Returns the Python object this handle stands for. */
+ public PyObject getPythonResource() {
+ return pythonResource;
+ }
+}
diff --git a/runtime/src/main/java/org/apache/flink/agents/runtime/python/utils/PythonActionExecutor.java b/runtime/src/main/java/org/apache/flink/agents/runtime/python/utils/PythonActionExecutor.java
index f39958e07..e645b071a 100644
--- a/runtime/src/main/java/org/apache/flink/agents/runtime/python/utils/PythonActionExecutor.java
+++ b/runtime/src/main/java/org/apache/flink/agents/runtime/python/utils/PythonActionExecutor.java
@@ -21,15 +21,22 @@ import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.databind.ObjectMapper;
import org.apache.flink.agents.api.Event;
import org.apache.flink.agents.api.agents.AgentExecutionOptions;
+import org.apache.flink.agents.api.resource.Resource;
+import org.apache.flink.agents.api.resource.ResourceType;
import org.apache.flink.agents.plan.AgentPlan;
import org.apache.flink.agents.plan.PythonFunction;
+import org.apache.flink.agents.runtime.operator.ActionTask;
import org.apache.flink.agents.runtime.python.context.PythonRunnerContextImpl;
+import org.apache.flink.agents.runtime.python.resource.PythonRuntimeResource;
import org.apache.flink.types.Row;
import org.apache.flink.util.ExceptionUtils;
import pemja.core.PythonInterpreter;
import pemja.core.object.PyObject;
import java.io.IOException;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.Map;
import java.util.concurrent.atomic.AtomicLong;
import static org.apache.flink.util.Preconditions.checkState;
@@ -49,6 +56,28 @@ public class PythonActionExecutor implements AutoCloseable {
private static final String CLOSE_FLINK_RUNNER_CONTEXT =
"flink_runner_context.close_flink_runner_context";
+ // =========== PYTHON RESOURCE MATERIALIZATION ===========
+ private static final String EAGER_MATERIALIZE = "flink_runner_context.eager_materialize";
+
+ // =========== TASK LIFECYCLE FORWARDING ===========
+ private static final String ADD_TASK_LIFECYCLE_LISTENER =
+ "flink_runner_context.add_task_lifecycle_listener";
+ private static final String NOTIFY_RECORD_START = "flink_runner_context.notify_record_start";
+ private static final String NOTIFY_ACTION_PREPARED =
+ "flink_runner_context.notify_action_prepared";
+ private static final String NOTIFY_ACTION_STARTED =
+ "flink_runner_context.notify_action_started";
+ private static final String NOTIFY_ACTION_TRANSFERRED =
+ "flink_runner_context.notify_action_transferred";
+ private static final String NOTIFY_ACTION_FINISHING =
+ "flink_runner_context.notify_action_finishing";
+ private static final String NOTIFY_ACTION_FINISHED =
+ "flink_runner_context.notify_action_finished";
+ private static final String NOTIFY_ACTION_REUSED = "flink_runner_context.notify_action_reused";
+ private static final String NOTIFY_ACTION_FAILED = "flink_runner_context.notify_action_failed";
+ private static final String NOTIFY_RECORD_FINISHED =
+ "flink_runner_context.notify_record_finished";
+
// ========== ASYNC THREAD POOL ===========
private static final String CREATE_ASYNC_THREAD_POOL =
"flink_runner_context.create_async_thread_pool";
@@ -97,6 +126,87 @@ public class PythonActionExecutor implements AutoCloseable {
return pythonRunnerContext;
}
+ /**
+ * Materializes every resource of the given type that the Python runtime owns and returns one
+ * handle per resource, keyed by resource name.
+ *
+ * <p>See {@link PythonRuntimeResource} for what the returned handle may and may not do.
+ */
+ @SuppressWarnings("unchecked")
+ public Map<String, Resource> eagerMaterialize(ResourceType type) {
+ Object pythonResources =
+ interpreter.invoke(EAGER_MATERIALIZE, pythonRunnerContext, type.getValue());
+ if (pythonResources == null) {
+ return Collections.emptyMap();
+ }
+ Map<String, Resource> handles = new HashMap<>();
+ ((Map<String, PyObject>) pythonResources)
+ .forEach(
+ (name, pythonResource) ->
+ handles.put(name, new PythonRuntimeResource(type, pythonResource)));
+ return handles;
+ }
+
+ /**
+ * Registers a Python object in the Python runtime's task lifecycle registry. The Python side
+ * fans the operator's callbacks out to that registry when {@link
+ * org.apache.flink.agents.runtime.lifecycle.PythonTaskLifecycleListener} forwards them.
+ *
+ * @return whether the object observes the lifecycle, so the caller can tell whether the Python
+ * runtime has anything to be notified about.
+ */
+ public boolean addTaskLifecycleListener(PyObject pythonListener) {
+ Object registered =
+ interpreter.invoke(
+ ADD_TASK_LIFECYCLE_LISTENER, pythonRunnerContext, pythonListener);
+ return Boolean.TRUE.equals(registered);
+ }
+
+ /** Forwards {@code onRecordStart} to the Python runtime lifecycle listeners. */
+ public void notifyRecordStart(Object key) {
+ interpreter.invoke(NOTIFY_RECORD_START, pythonRunnerContext, key);
+ }
+
+ /** Forwards {@code onActionPrepared} to the Python runtime lifecycle listeners. */
+ public void notifyActionPrepared(ActionTask task) {
+ interpreter.invoke(NOTIFY_ACTION_PREPARED, pythonRunnerContext, task);
+ }
+
+ /** Forwards {@code onActionStarted} to the Python runtime lifecycle listeners. */
+ public void notifyActionStarted(ActionTask task) {
+ interpreter.invoke(NOTIFY_ACTION_STARTED, pythonRunnerContext, task);
+ }
+
+ /** Forwards {@code onActionTransferred} to the Python runtime lifecycle listeners. */
+ public void notifyActionTransferred(ActionTask fromTask, ActionTask toTask) {
+ interpreter.invoke(NOTIFY_ACTION_TRANSFERRED, pythonRunnerContext, fromTask, toTask);
+ }
+
+ /** Forwards {@code onActionFinishing} to the Python runtime lifecycle listeners. */
+ public void notifyActionFinishing(ActionTask task) {
+ interpreter.invoke(NOTIFY_ACTION_FINISHING, pythonRunnerContext, task);
+ }
+
+ /** Forwards {@code onActionFinished} to the Python runtime lifecycle listeners. */
+ public void notifyActionFinished(ActionTask task) {
+ interpreter.invoke(NOTIFY_ACTION_FINISHED, pythonRunnerContext, task);
+ }
+
+ /** Forwards {@code onActionReused} to the Python runtime lifecycle listeners. */
+ public void notifyActionReused(ActionTask task) {
+ interpreter.invoke(NOTIFY_ACTION_REUSED, pythonRunnerContext, task);
+ }
+
+ /** Forwards {@code onActionFailed} to the Python runtime lifecycle listeners. */
+ public void notifyActionFailed(ActionTask task, Throwable error) {
+ interpreter.invoke(NOTIFY_ACTION_FAILED, pythonRunnerContext, task, error);
+ }
+
+ /** Forwards {@code onRecordFinished} to the Python runtime lifecycle listeners. */
+ public void notifyRecordFinished(Object key) {
+ interpreter.invoke(NOTIFY_RECORD_FINISHED, pythonRunnerContext, key);
+ }
+
public void open() throws Exception {
interpreter.exec(PYTHON_IMPORTS);
diff --git a/runtime/src/main/java/org/apache/flink/agents/runtime/subagent/AsyncSubagentFuture.java b/runtime/src/main/java/org/apache/flink/agents/runtime/subagent/AsyncSubagentFuture.java
new file mode 100644
index 000000000..a2d053aea
--- /dev/null
+++ b/runtime/src/main/java/org/apache/flink/agents/runtime/subagent/AsyncSubagentFuture.java
@@ -0,0 +1,125 @@
+/*
+ * 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.subagent;
+
+import org.apache.flink.agents.api.context.DurableCallable;
+import org.apache.flink.agents.api.context.RunnerContext;
+import org.apache.flink.agents.api.subagent.SubagentFuture;
+import org.apache.flink.agents.api.subagent.SubagentFutures;
+import org.apache.flink.agents.api.subagent.SubagentResult;
+import org.apache.flink.agents.runtime.subagent.BaseAsyncSubagentSetup.RunStatus;
+
+import javax.annotation.Nullable;
+
+import java.util.concurrent.CancellationException;
+
+/**
+ * The sub side of an async-job invocation: the run was already started by the durable POST of
+ * {@code submit}, so the handle only subscribes to it.
+ */
+final class AsyncSubagentFuture extends SubagentFuture {
+
+ private final BaseAsyncSubagentSetup setup;
+ private final RunnerContext ctx;
+ @Nullable private final PendingSubagentCallRegistry registry;
+
+ private boolean consumed;
+ private boolean cancelled;
+ @Nullable private SubagentResult value;
+
+ AsyncSubagentFuture(
+ BaseAsyncSubagentSetup setup,
+ RunnerContext ctx,
+ String sessionId,
+ String callId,
+ @Nullable PendingSubagentCallRegistry registry) {
+ super(sessionId, callId);
+ this.setup = setup;
+ this.ctx = ctx;
+ this.registry = registry;
+ if (registry != null) {
+ registry.trackPendingSubagentCall(identity());
+ }
+ }
+
+ /**
+ * Probes the remote status directly. The probe runs outside durable execution, so a failover
+ * replay may probe a different number of times than the original execution. A probe failure
+ * propagates and fails the action.
+ */
+ @Override
+ public boolean isDone() {
+ if (consumed || cancelled) {
+ return true;
+ }
+ RunStatus probe = setup.queryStatus(getSessionId(), getCallId());
+ return probe.getState() == RunStatus.State.COMPLETED
+ || probe.getState() == RunStatus.State.FAILED;
+ }
+
+ /**
+ * Waits for the run through the durable await composition. A cancelled handle fails as a {@link
+ * CancellationException}.
+ */
+ @Override
+ public SubagentResult await() throws Exception {
+ if (cancelled) {
+ throw new CancellationException(
+ "Sub-agent call cancelled: " + getSessionId() + "#" + getCallId());
+ }
+ if (!consumed) {
+ DurableCallable<SubagentResult> awaitCall =
+ setup.awaitResult(ctx, getSessionId(), getCallId());
+ value = ctx.durableExecuteAsync(awaitCall);
+ consumed = true;
+ if (registry != null) {
+ registry.untrackPendingSubagentCall(identity());
+ }
+ }
+ return value;
+ }
+
+ /**
+ * Propagates the cancellation through the setup's {@link
+ * BaseAsyncSubagentSetup#callCancelRequest} hook. The propagation runs synchronously through
+ * the hook and is replayed with the enclosing action, so a failover may propagate the same
+ * cancellation again. A repeated cancel on the same handle and a cancel after the resolve are
+ * local no-ops. A hook failure propagates and fails the action.
+ */
+ @Override
+ public void cancel() {
+ if (consumed || cancelled) {
+ return;
+ }
+ setup.cancelRequest(ctx, getSessionId(), getCallId());
+ cancelled = true;
+ if (registry != null) {
+ registry.untrackPendingSubagentCall(identity());
+ }
+ }
+
+ private String identity() {
+ return getSessionId() + "#" + getCallId();
+ }
+
+ @Override
+ public SubagentFutures combine(SubagentFuture... others) {
+ return new SubagentFutureGroup(this, others);
+ }
+}
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
new file mode 100644
index 000000000..4f0044c21
--- /dev/null
+++ b/runtime/src/main/java/org/apache/flink/agents/runtime/subagent/BaseAsyncSubagentSetup.java
@@ -0,0 +1,287 @@
+/*
+ * 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.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;
+
+import javax.annotation.Nullable;
+
+import java.util.concurrent.Callable;
+
+/**
+ * Production base for sub-agents whose protocol is an asynchronous job, run in durable pub/sub
+ * mode: {@code submit} publishes the run through one durable POST, the returned handle subscribes
+ * to it.
+ */
+public abstract class BaseAsyncSubagentSetup extends BaseSubagentSetup {
+
+ /**
+ * 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
+ * status_poll_interval_millis} argument over it, and subclasses may override it directly.
+ */
+ 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
+ * 500} when absent.
+ */
+ protected BaseAsyncSubagentSetup(
+ ResourceDescriptor descriptor, ResourceContext resourceContext) {
+ Number statusPollInterval = descriptor.getArgument("status_poll_interval_millis");
+ if (statusPollInterval != null) {
+ this.statusPollIntervalMillis = statusPollInterval.longValue();
+ }
+ }
+
+ // ------------------------------------------------------------------------------------------
+ // pub: submit starts the run immediately through one durable POST
+ // ------------------------------------------------------------------------------------------
+
+ /**
+ * Starts the remote run through the durable POST and returns its handle. A POST failure throws
+ * and fails the action.
+ */
+ @Override
+ public final SubagentFuture submit(
+ RunnerContext ctx, Object prompt, String sessionId, String callId) throws Exception {
+ ctx.durableExecuteAsync(submitRequest(ctx, sessionId, callId, prompt));
+ return new AsyncSubagentFuture(this, ctx, sessionId, callId, currentTaskRegistry());
+ }
+
+ // ------------------------------------------------------------------------------------------
+ // Framework wrappers: defaults composing the primitives, overridable
+ // ------------------------------------------------------------------------------------------
+
+ /** The durable POST of one invocation. It is the only wrapper wired to a reconciler. */
+ protected DurableCallable<Void> submitRequest(
+ RunnerContext ctx, String sessionId, String callId, Object prompt) {
+ return new DurableCallable<Void>() {
+ @Override
+ public String getId() {
+ return sessionId + "#" + callId;
+ }
+
+ @Override
+ public Class<Void> getResultClass() {
+ return Void.class;
+ }
+
+ @Override
+ public Void call() throws Exception {
+ callSubmitRequest(sessionId, callId, prompt);
+ return null;
+ }
+
+ @Override
+ public Callable<Void> reconciler() {
+ // Recovery probes first through reconcileSubmitRequest, so a landed POST is never
+ // duplicated.
+ return () -> {
+ reconcileSubmitRequest(sessionId, callId, prompt);
+ return null;
+ };
+ }
+ };
+ }
+
+ /**
+ * The status probe. It is a direct read-only query on the mailbox thread, so durable execution
+ * does not record it and a failover replay probes again.
+ */
+ protected RunStatus queryStatus(String sessionId, String callId) {
+ return callQueryStatus(sessionId, callId);
+ }
+
+ /**
+ * The durable wait of one resolve: poll the status until the run reaches a terminal state, then
+ * fetch the result. Keyed by {@code sessionId#callId#await}. A probe or fetch failure that
+ * escapes the body is a system-level failure: it propagates instead of being folded into an
+ * error result.
+ */
+ protected DurableCallable<SubagentResult> awaitResult(
+ RunnerContext ctx, String sessionId, String callId) {
+ return new DurableCallable<SubagentResult>() {
+ @Override
+ public String getId() {
+ return sessionId + "#" + callId + "#await";
+ }
+
+ @Override
+ public Class<SubagentResult> getResultClass() {
+ return SubagentResult.class;
+ }
+
+ @Override
+ public SubagentResult call() throws Exception {
+ while (true) {
+ RunStatus probe = callQueryStatus(sessionId, callId);
+ switch (probe.getState()) {
+ case COMPLETED:
+ return callFetchResult(sessionId, callId);
+ case FAILED:
+ return SubagentResult.error(probe.getError());
+ default:
+ // NOT_STARTED or RUNNING: keep probing. A NOT_STARTED run after a
+ // durable POST means the remote session expired; the replay then
+ // observes the fresh state instead of the original probe path.
+ Thread.sleep(statusPollIntervalMillis);
+ }
+ }
+ }
+ };
+ }
+
+ /**
+ * The cancellation propagation. The wrapper calls the hook synchronously, so durable execution
+ * does not record the propagation and a failover replay propagates it again.
+ */
+ protected void cancelRequest(RunnerContext ctx, String sessionId, String callId) {
+ callCancelRequest(sessionId, callId);
+ }
+
+ // ------------------------------------------------------------------------------------------
+ // Transport primitives provided by the integration
+ // ------------------------------------------------------------------------------------------
+
+ /** Starts the run remotely. A thrown exception fails the enclosing action. */
+ protected abstract void callSubmitRequest(String sessionId, String callId, Object prompt)
+ throws Exception;
+
+ /**
+ * Read-only probe of the run's current state; must not alter the remote run. The status never
+ * carries the result payload — the result is fetched separately through {@link
+ * #callFetchResult}.
+ *
+ * <p>Implementations must report comprehensible failures (an expired endpoint, expired
+ * credentials, a rejected run) as a FAILED status rather than throwing; a RuntimeException
+ * escaping this probe is treated as a system-level failure, propagates, and triggers a job
+ * failover.
+ */
+ protected abstract RunStatus callQueryStatus(String sessionId, String callId);
+
+ /**
+ * Fetches the result of a run that reached a terminal state; comprehensible failures go into
+ * the {@link SubagentResult}, while an escaping exception is a system-level failure that
+ * propagates. The fetch must be an idempotent read: a failover re-executes it when the crash
+ * hit the fetch in flight.
+ */
+ protected abstract SubagentResult callFetchResult(String sessionId, String callId)
+ throws Exception;
+
+ /**
+ * The crash-window recovery of the POST: probes the status and handles every state explicitly,
+ * so a landed POST is never duplicated. A probe failure propagates and fails the recovery.
+ */
+ protected void reconcileSubmitRequest(String sessionId, String callId, Object prompt)
+ throws Exception {
+ RunStatus probe = callQueryStatus(sessionId, callId);
+ switch (probe.getState()) {
+ case NOT_STARTED:
+ // The service has no record of the run: the POST never landed. Start it.
+ callSubmitRequest(sessionId, callId, prompt);
+ break;
+ case RUNNING:
+ // The POST landed and the run is in flight; the subsequent await keeps
+ // polling it. Nothing to repair.
+ break;
+ case COMPLETED:
+ case FAILED:
+ // The run reached a terminal state while the caller was down; the
+ // subsequent await picks up the outcome — the fetch or the reported
+ // error. Nothing to repair.
+ break;
+ default:
+ // Fail loudly instead of silently skipping an unknown state.
+ throw new IllegalStateException("Unknown run state: " + probe.getState());
+ }
+ }
+
+ /**
+ * Hook propagating a cancellation to the remote run. The default is a no-op. A replay may
+ * propagate the cancellation again, so remote cancellations must be idempotent.
+ */
+ protected void callCancelRequest(String sessionId, String callId) {}
+
+ // ------------------------------------------------------------------------------------------
+ // The state snapshot of a remote run
+ // ------------------------------------------------------------------------------------------
+
+ /**
+ * The state snapshot of a remote run, as reported by the read-only {@link #callQueryStatus}
+ * probe. A state other than {@link State#NOT_STARTED} means the submission landed on the
+ * service, which is the sole basis for {@link #reconcileSubmitRequest} deciding between
+ * re-posting and polling. The snapshot never carries the result payload.
+ */
+ public static final class RunStatus {
+
+ /** Lifecycle of the remote run. */
+ public enum State {
+ NOT_STARTED,
+ RUNNING,
+ COMPLETED,
+ FAILED
+ }
+
+ private final State state;
+ @Nullable private final String error;
+
+ private RunStatus(State state, @Nullable String error) {
+ this.state = state;
+ this.error = error;
+ }
+
+ /** The service has no record of the run: the POST never landed (or the id mismatches). */
+ public static RunStatus notStarted() {
+ return new RunStatus(State.NOT_STARTED, null);
+ }
+
+ /** The run is in progress. */
+ public static RunStatus running() {
+ return new RunStatus(State.RUNNING, null);
+ }
+
+ /** The run finished successfully. */
+ public static RunStatus completed() {
+ return new RunStatus(State.COMPLETED, null);
+ }
+
+ /** The run failed, carrying the error message. */
+ public static RunStatus failed(String error) {
+ return new RunStatus(State.FAILED, error);
+ }
+
+ public State getState() {
+ return state;
+ }
+
+ @Nullable
+ public String getError() {
+ return error;
+ }
+ }
+}
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
new file mode 100644
index 000000000..666898eb6
--- /dev/null
+++ b/runtime/src/main/java/org/apache/flink/agents/runtime/subagent/BaseDeferredSubagentSetup.java
@@ -0,0 +1,66 @@
+/*
+ * 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.subagent;
+
+import org.apache.flink.agents.api.context.DurableCallable;
+import org.apache.flink.agents.api.context.RunnerContext;
+import org.apache.flink.agents.api.subagent.SubagentFuture;
+import org.apache.flink.agents.api.subagent.SubagentResult;
+
+/**
+ * Framework-level deferred execution mode for sub-agent setups: {@code submit} registers the
+ * invocation and returns a deferred handle without sending anything; the actual request is issued
+ * lazily when the handle is first resolved, and runs through one durable async callable keyed by a
+ * failover-reproducible id, so the invocation participates in the task's durable execution.
+ */
+public abstract class BaseDeferredSubagentSetup extends BaseSubagentSetup {
+
+ /** Registers the invocation and returns its deferred handle. */
+ @Override
+ public SubagentFuture submit(RunnerContext ctx, Object prompt, String sessionId, String callId)
+ throws Exception {
+ return new DeferredSubagentFuture(
+ sessionId,
+ callId,
+ ctx,
+ currentTaskRegistry(),
+ () -> prepare(ctx, prompt, sessionId, callId));
+ }
+
+ /**
+ * Prepares one invocation and returns the {@link DurableCallable} running it. Both ids are
+ * already assigned; the durable id MUST be derived solely from the {@code (sessionId, callId)}
+ * pair so it is reproducible after failover.
+ *
+ * <p>Called exactly once per invocation, when the deferred handle is first resolved, on the
+ * mailbox thread. Implementations may therefore perform the mailbox-confined part of issuing
+ * the request here; the returned callable's {@link DurableCallable#call()} carries only the
+ * part that runs off the mailbox thread.
+ *
+ * <p>The callable folds its own comprehensible failures into the returned {@link
+ * SubagentResult}; an exception escaping {@link DurableCallable#call()} is a system-level
+ * failure that propagates and fails the action.
+ *
+ * <p>Skipping the reconciler on the returned callable has a cost: a crash between the call
+ * landing and its result being persisted re-invokes the sub-agent on replay, possibly
+ * duplicating external side effects.
+ */
+ protected abstract DurableCallable<SubagentResult> prepare(
+ 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
new file mode 100644
index 000000000..a3b07c3ce
--- /dev/null
+++ b/runtime/src/main/java/org/apache/flink/agents/runtime/subagent/BaseSubagentSetup.java
@@ -0,0 +1,153 @@
+/*
+ * 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.subagent;
+
+import org.apache.flink.agents.api.context.RunnerContext;
+import org.apache.flink.agents.api.subagent.SubagentFuture;
+import org.apache.flink.agents.api.subagent.SubagentSetup;
+import org.apache.flink.agents.runtime.lifecycle.TaskLifecycleListener;
+import org.apache.flink.agents.runtime.operator.ActionTask;
+
+import javax.annotation.Nullable;
+
+import java.util.HashMap;
+import java.util.Map;
+
+/**
+ * Runtime base for sub-agent setups, holding the per-task id allocators and pending-call registries
+ * keyed to the currently executing action task. How an invocation is issued stays an execution mode
+ * owned by the concrete subclass.
+ */
+public abstract class BaseSubagentSetup extends SubagentSetup implements TaskLifecycleListener {
+
+ private final Map<ActionTask, SubagentIdAllocator> perTaskAllocators = new HashMap<>();
+ private final Map<ActionTask, PendingSubagentCallRegistry> perTaskRegistries = new HashMap<>();
+
+ /** The task whose execution is currently issuing calls. */
+ @Nullable private ActionTask currentTask;
+
+ @Override
+ public void onActionPrepared(ActionTask task) {
+ currentTask = task;
+ }
+
+ @Override
+ public void onActionTransferred(ActionTask from, ActionTask to) {
+ SubagentIdAllocator allocator = perTaskAllocators.remove(from);
+ if (allocator != null) {
+ perTaskAllocators.put(to, allocator);
+ }
+ PendingSubagentCallRegistry registry = perTaskRegistries.remove(from);
+ if (registry != null) {
+ registry.setActionName(to.getAction().getName());
+ perTaskRegistries.put(to, registry);
+ }
+ }
+
+ /**
+ * Finalizes the task's bookkeeping once its outcome is fixed. The replay-reuse path reaches the
+ * same finalization through {@link #onActionReused}, keeping the prepared/terminal pairing
+ * intact on both paths. A failed invocation intentionally skips this cleanup: the failure fails
+ * the run and the task is replayed on the restarted operator, so stale entries cannot outlive
+ * the run.
+ */
+ @Override
+ public void onActionFinishing(ActionTask task) {
+ currentTask = null;
+ perTaskAllocators.remove(task);
+ PendingSubagentCallRegistry registry = perTaskRegistries.remove(task);
+ if (registry != null) {
+ registry.checkEmpty();
+ }
+ }
+
+ /** Reuse is a terminal outcome like finishing, so it shares the finalization hook. */
+ @Override
+ public void onActionReused(ActionTask task) {
+ onActionFinishing(task);
+ }
+
+ /**
+ * The registry of the currently executing task, where handles record themselves on creation.
+ * Returns {@code null} outside a prepared task, so calls issued without a task context skip
+ * tracking.
+ */
+ @Nullable
+ protected final PendingSubagentCallRegistry currentTaskRegistry() {
+ if (currentTask == null) {
+ return null;
+ }
+ ActionTask task = currentTask;
+ return perTaskRegistries.computeIfAbsent(
+ task, t -> new PendingSubagentCallRegistry(t.getAction().getName()));
+ }
+
+ /** The task whose execution is currently issuing calls, or {@code null} outside one. */
+ @Nullable
+ protected final ActionTask currentTask() {
+ return currentTask;
+ }
+
+ /**
+ * Injected by the framework with the setup's resource name when the resource is materialized.
+ */
+ @Nullable private String subagentName;
+
+ public final void setSubagentName(String subagentName) {
+ this.subagentName = subagentName;
+ }
+
+ public final String getSubagentName() {
+ return subagentName;
+ }
+
+ @Override
+ public SubagentFuture submit(RunnerContext ctx, Object prompt, String sessionId)
+ throws Exception {
+ return submit(ctx, prompt, sessionId, currentAllocator().nextCallId(sessionId));
+ }
+
+ @Override
+ public SubagentFuture submit(RunnerContext ctx, Object prompt) throws Exception {
+ SubagentIdAllocator allocator = currentAllocator();
+ String sessionId = allocator.nextSessionId();
+ return submit(ctx, prompt, sessionId, allocator.nextCallId(sessionId));
+ }
+
+ /**
+ * The allocator of the currently executing task, scoped to one action execution so ordinals
+ * restart for the next action. Failover replays hand out the same ids.
+ */
+ protected final SubagentIdAllocator currentAllocator() {
+ if (currentTask == null) {
+ throw new IllegalStateException(
+ "No prepared action task to assign sub-agent ids from.");
+ }
+ ActionTask task = currentTask;
+ return perTaskAllocators.computeIfAbsent(
+ task,
+ t ->
+ new SubagentIdAllocator(
+ t.getKey(),
+ t.getSequenceNumber(),
+ t.getAction().getName(),
+ t.getEvent(),
+ subagentName));
+ }
+}
diff --git a/runtime/src/main/java/org/apache/flink/agents/runtime/subagent/CompletedSubagentFuture.java b/runtime/src/main/java/org/apache/flink/agents/runtime/subagent/CompletedSubagentFuture.java
new file mode 100644
index 000000000..2bd550660
--- /dev/null
+++ b/runtime/src/main/java/org/apache/flink/agents/runtime/subagent/CompletedSubagentFuture.java
@@ -0,0 +1,49 @@
+/*
+ * 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.subagent;
+
+import org.apache.flink.agents.api.subagent.SubagentFuture;
+import org.apache.flink.agents.api.subagent.SubagentFutures;
+import org.apache.flink.agents.api.subagent.SubagentResult;
+
+/** A handle for an invocation that has already produced its value. */
+public final class CompletedSubagentFuture extends SubagentFuture {
+
+ private final SubagentResult value;
+
+ public CompletedSubagentFuture(String sessionId, String callId, SubagentResult value) {
+ super(sessionId, callId);
+ this.value = value;
+ }
+
+ @Override
+ public boolean isDone() {
+ return true;
+ }
+
+ @Override
+ public SubagentResult await() {
+ return value;
+ }
+
+ @Override
+ public SubagentFutures combine(SubagentFuture... others) {
+ return new SubagentFutureGroup(this, others);
+ }
+}
diff --git a/runtime/src/main/java/org/apache/flink/agents/runtime/subagent/DeferredSubagentFuture.java b/runtime/src/main/java/org/apache/flink/agents/runtime/subagent/DeferredSubagentFuture.java
new file mode 100644
index 000000000..913ccdb89
--- /dev/null
+++ b/runtime/src/main/java/org/apache/flink/agents/runtime/subagent/DeferredSubagentFuture.java
@@ -0,0 +1,129 @@
+/*
+ * 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.subagent;
+
+import org.apache.flink.agents.api.context.DurableCallable;
+import org.apache.flink.agents.api.context.RunnerContext;
+import org.apache.flink.agents.api.subagent.SubagentFuture;
+import org.apache.flink.agents.api.subagent.SubagentFutures;
+import org.apache.flink.agents.api.subagent.SubagentResult;
+
+import javax.annotation.Nullable;
+
+import java.util.concurrent.CancellationException;
+import java.util.function.Supplier;
+
+/** Deferred handle to one sub-agent invocation. */
+public final class DeferredSubagentFuture extends SubagentFuture {
+
+ private final RunnerContext ctx;
+ @Nullable private final PendingSubagentCallRegistry registry;
+ private final Supplier<DurableCallable<SubagentResult>> preparedSupplier;
+
+ @Nullable private DurableCallable<SubagentResult> prepared;
+ private boolean done;
+ private boolean cancelled;
+ @Nullable private SubagentResult value;
+
+ public DeferredSubagentFuture(
+ String sessionId,
+ String callId,
+ RunnerContext ctx,
+ @Nullable PendingSubagentCallRegistry registry,
+ Supplier<DurableCallable<SubagentResult>> preparedSupplier) {
+ super(sessionId, callId);
+ this.ctx = ctx;
+ this.registry = registry;
+ this.preparedSupplier = preparedSupplier;
+ if (registry != null) {
+ registry.trackPendingSubagentCall(identity());
+ }
+ }
+
+ /** Prepares the request if it has not been prepared yet; must run on the mailbox thread. */
+ DurableCallable<SubagentResult> prepare() {
+ if (cancelled) {
+ throw new CancellationException("Sub-agent call cancelled: " + identity());
+ }
+ if (prepared == null) {
+ prepared = preparedSupplier.get();
+ }
+ return prepared;
+ }
+
+ /**
+ * Runs the prepared request through durable execution and records the outcome. Mailbox releases
+ * happen inside the durable execution itself.
+ *
+ * <p>A system-level failure escaping durable execution propagates and fails the action instead
+ * of being folded into an error result.
+ */
+ void execute() throws Exception {
+ complete(ctx.durableExecuteAsync(prepare()));
+ }
+
+ /**
+ * Cancels before the request is prepared: the request is discarded and resolving the handle
+ * fails. An already resolved handle ignores the cancellation request.
+ */
+ @Override
+ public void cancel() {
+ if (done) {
+ return;
+ }
+ cancelled = true;
+ if (registry != null) {
+ registry.untrackPendingSubagentCall(identity());
+ }
+ }
+
+ private String identity() {
+ return getSessionId() + "#" + getCallId();
+ }
+
+ /** Records the outcome produced by a batched wait. */
+ private void complete(SubagentResult outcome) {
+ this.value = outcome;
+ this.done = true;
+ if (registry != null) {
+ registry.untrackPendingSubagentCall(identity());
+ }
+ }
+
+ @Override
+ public boolean isDone() {
+ return done || cancelled;
+ }
+
+ @Override
+ public SubagentResult await() throws Exception {
+ if (cancelled) {
+ throw new CancellationException("Sub-agent call cancelled: " + identity());
+ }
+ if (!done) {
+ execute();
+ }
+ return value;
+ }
+
+ @Override
+ public SubagentFutures combine(SubagentFuture... others) {
+ return new SubagentFutureGroup(this, others);
+ }
+}
diff --git a/runtime/src/main/java/org/apache/flink/agents/runtime/subagent/PendingSubagentCallRegistry.java b/runtime/src/main/java/org/apache/flink/agents/runtime/subagent/PendingSubagentCallRegistry.java
new file mode 100644
index 000000000..b48746a00
--- /dev/null
+++ b/runtime/src/main/java/org/apache/flink/agents/runtime/subagent/PendingSubagentCallRegistry.java
@@ -0,0 +1,67 @@
+/*
+ * 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.subagent;
+
+import java.util.LinkedHashSet;
+import java.util.Set;
+
+/** The per-action-execution set of sub-agent handles submitted but not yet resolved. */
+public final class PendingSubagentCallRegistry {
+
+ private final Set<String> pendingCalls = new LinkedHashSet<>();
+
+ /** The action the pending handles belong to, named in the failure message. */
+ private String actionName;
+
+ public PendingSubagentCallRegistry(String actionName) {
+ this.actionName = actionName;
+ }
+
+ /** Adopts the continuation's action when the execution moves onto another task. */
+ public void setActionName(String actionName) {
+ this.actionName = actionName;
+ }
+
+ /** Records a handle. Duplicate identities collapse to one entry. */
+ public void trackPendingSubagentCall(String callIdentity) {
+ pendingCalls.add(callIdentity);
+ }
+
+ /** Drops a resolved handle and does nothing when the identity is unknown. */
+ public void untrackPendingSubagentCall(String callIdentity) {
+ pendingCalls.remove(callIdentity);
+ }
+
+ public boolean isEmpty() {
+ return pendingCalls.isEmpty();
+ }
+
+ /** Fails the action when it left a sub-agent handle unresolved. */
+ public void checkEmpty() {
+ if (!pendingCalls.isEmpty()) {
+ throw new IllegalStateException(
+ "Action "
+ + actionName
+ + " finished without resolving the sub-agent calls it submitted: "
+ + pendingCalls
+ + ". Resolve every handle returned by submit(), individually or through "
+ + "SubagentFutures.");
+ }
+ }
+}
diff --git a/runtime/src/main/java/org/apache/flink/agents/runtime/subagent/SubagentFutureGroup.java b/runtime/src/main/java/org/apache/flink/agents/runtime/subagent/SubagentFutureGroup.java
new file mode 100644
index 000000000..c7312de46
--- /dev/null
+++ b/runtime/src/main/java/org/apache/flink/agents/runtime/subagent/SubagentFutureGroup.java
@@ -0,0 +1,93 @@
+/*
+ * 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.subagent;
+
+import org.apache.flink.agents.api.subagent.SubagentFuture;
+import org.apache.flink.agents.api.subagent.SubagentFutures;
+import org.apache.flink.agents.api.subagent.SubagentResult;
+
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.List;
+
+/** The {@link SubagentFutures} returned by {@code combine}: several handles held together. */
+final class SubagentFutureGroup extends SubagentFutures {
+
+ private final List<SubagentFuture> futures;
+
+ SubagentFutureGroup(SubagentFuture first, SubagentFuture[] others) {
+ this(withFirst(first, others));
+ }
+
+ private static List<SubagentFuture> withFirst(SubagentFuture first, SubagentFuture[] others) {
+ List<SubagentFuture> all = new ArrayList<>(1 + others.length);
+ all.add(first);
+ all.addAll(Arrays.asList(others));
+ return all;
+ }
+
+ private SubagentFutureGroup(List<SubagentFuture> futures) {
+ this.futures = futures;
+ }
+
+ @Override
+ public boolean isDone() {
+ for (SubagentFuture future : futures) {
+ if (!future.isDone()) {
+ return false;
+ }
+ }
+ return true;
+ }
+
+ @Override
+ public List<SubagentResult> awaitAll() throws Exception {
+ for (SubagentFuture future : futures) {
+ if (future instanceof DeferredSubagentFuture && !future.isDone()) {
+ ((DeferredSubagentFuture) future).prepare();
+ }
+ }
+ // TODO(#926): execute the prepared calls as one batch once durable execution supports
+ // batched submission; until then the prepared calls are executed one by one.
+ for (SubagentFuture future : futures) {
+ if (future instanceof DeferredSubagentFuture && !future.isDone()) {
+ ((DeferredSubagentFuture) future).execute();
+ }
+ }
+ List<SubagentResult> outcomes = new ArrayList<>(futures.size());
+ for (SubagentFuture future : futures) {
+ outcomes.add(future.await());
+ }
+ return outcomes;
+ }
+
+ @Override
+ public void cancel() {
+ for (SubagentFuture future : futures) {
+ future.cancel();
+ }
+ }
+
+ @Override
+ public SubagentFutures combine(SubagentFuture... others) {
+ List<SubagentFuture> grown = new ArrayList<>(futures);
+ grown.addAll(Arrays.asList(others));
+ return new SubagentFutureGroup(grown);
+ }
+}
diff --git a/runtime/src/main/java/org/apache/flink/agents/runtime/subagent/SubagentIdAllocator.java b/runtime/src/main/java/org/apache/flink/agents/runtime/subagent/SubagentIdAllocator.java
new file mode 100644
index 000000000..41da94ce6
--- /dev/null
+++ b/runtime/src/main/java/org/apache/flink/agents/runtime/subagent/SubagentIdAllocator.java
@@ -0,0 +1,143 @@
+/*
+ * 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.subagent;
+
+import com.fasterxml.jackson.annotation.JsonIgnore;
+import com.fasterxml.jackson.annotation.JsonInclude;
+import com.fasterxml.jackson.annotation.JsonProperty;
+import com.fasterxml.jackson.core.JsonProcessingException;
+import com.fasterxml.jackson.databind.MapperFeature;
+import com.fasterxml.jackson.databind.ObjectMapper;
+import com.fasterxml.jackson.databind.SerializationFeature;
+import com.fasterxml.jackson.databind.json.JsonMapper;
+import org.apache.flink.agents.api.Event;
+
+import javax.annotation.Nullable;
+
+import java.util.HashMap;
+import java.util.Map;
+import java.util.UUID;
+
+/**
+ * Deterministically assigns sub-agent session and call ids for one action execution. The {@link
+ * Namespace} of caller-side facts fixes the counting range, so a failover replay reproduces the
+ * same id sequence.
+ */
+public final class SubagentIdAllocator {
+
+ private final Namespace namespace;
+
+ private int sessionOrdinal = 0;
+ private final Map<String, Integer> perSessionCallOrdinals = new HashMap<>();
+
+ /** Creates an allocator for one action execution from the execution's caller-side facts. */
+ public SubagentIdAllocator(
+ Object key, long sequenceNumber, String actionName, Event event, String subagentName) {
+ this.namespace = new Namespace(key, sequenceNumber, actionName, event, subagentName);
+ }
+
+ /** Creates a new, ordinal-increasing session id scoped to this task's namespace. */
+ public String nextSessionId() {
+ return namespace.digest() + "-" + (sessionOrdinal++);
+ }
+
+ /**
+ * Creates a new call id by appending the per-session ordinal (starting at 1) to the session id.
+ * Ordinals restart per action execution, so ids assigned here stay valid only within it.
+ */
+ public String nextCallId(String sessionId) {
+ int ordinal = perSessionCallOrdinals.merge(sessionId, 1, Integer::sum);
+ return sessionId + "-" + ordinal;
+ }
+
+ /**
+ * The caller-side identity of one action execution, seeding the deterministic ids of the
+ * sub-agent calls it issues.
+ *
+ * <p>Key, sequence number, action name, and the event's type and attributes are facts of the
+ * execution itself, identical for every sub-agent called from it. The subagent name
+ * distinguishes the sub-agents called from one action, so it alone keeps their id ranges apart.
+ */
+ @JsonInclude(JsonInclude.Include.NON_NULL)
+ public static final class Namespace {
+
+ /**
+ * Sorts map entries and bean properties so the namespace bytes do not depend on map
+ * iteration order, which is not guaranteed across JVMs.
+ */
+ private static final ObjectMapper DIGEST_MAPPER =
+ JsonMapper.builder()
+ .configure(SerializationFeature.ORDER_MAP_ENTRIES_BY_KEYS, true)
+ .configure(MapperFeature.SORT_PROPERTIES_ALPHABETICALLY, true)
+ .build();
+
+ @JsonProperty("key")
+ private final String key;
+
+ @JsonProperty("sequenceNumber")
+ private final long sequenceNumber;
+
+ @JsonProperty("actionName")
+ private final String actionName;
+
+ @JsonProperty("eventType")
+ private final String eventType;
+
+ @JsonProperty("eventAttributes")
+ private final Map<String, Object> eventAttributes;
+
+ @JsonProperty("subagentName")
+ private final String subagentName;
+
+ /**
+ * Computed lazily on the first allocation. Digesting is mailbox-confined, so it needs no
+ * synchronization.
+ */
+ @JsonIgnore @Nullable private String digest;
+
+ public Namespace(
+ Object key,
+ long sequenceNumber,
+ String actionName,
+ Event event,
+ String subagentName) {
+ this.key = key.toString();
+ this.sequenceNumber = sequenceNumber;
+ this.actionName = actionName;
+ this.eventType = event.getType();
+ this.eventAttributes = event.getAttributes();
+ this.subagentName = subagentName;
+ }
+
+ /** Digests the id-bearing facts into a name-based UUID string, stable across replays. */
+ public String digest() {
+ if (digest == null) {
+ try {
+ digest =
+ String.valueOf(
+ UUID.nameUUIDFromBytes(DIGEST_MAPPER.writeValueAsBytes(this)));
+ } catch (JsonProcessingException e) {
+ throw new IllegalStateException(
+ "Failed to digest the sub-agent identity namespace", e);
+ }
+ }
+ return digest;
+ }
+ }
+}
diff --git a/runtime/src/main/java/org/apache/flink/agents/runtime/trace/EventLogComponentExecutionListener.java b/runtime/src/main/java/org/apache/flink/agents/runtime/trace/EventLogComponentExecutionListener.java
new file mode 100644
index 000000000..30f9e18e9
--- /dev/null
+++ b/runtime/src/main/java/org/apache/flink/agents/runtime/trace/EventLogComponentExecutionListener.java
@@ -0,0 +1,86 @@
+/*
+ * 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.trace;
+
+import org.apache.flink.agents.api.Event;
+import org.apache.flink.agents.api.trace.ExecutionLifecycleEvents;
+import org.apache.flink.agents.api.trace.ExecutionTraceContext;
+import org.apache.flink.agents.runtime.lifecycle.ComponentExecutionListener;
+import org.apache.flink.annotation.Internal;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.util.HashMap;
+import java.util.Map;
+
+/**
+ * Per-action-execution adapter that turns component execution reports into event log records under
+ * the action's trace context. Its bookkeeping never leaks across actions because each execution
+ * gets its own instance, and the start/terminal pairing survives continuation task transfers
+ * because the adapter is tied to the action execution rather than the individual task.
+ */
+@Internal
+public final class EventLogComponentExecutionListener implements ComponentExecutionListener {
+
+ private static final Logger LOG =
+ LoggerFactory.getLogger(EventLogComponentExecutionListener.class);
+
+ private final ExecutionTraceContext actionTraceContext;
+ private final ExecutionEventSink executionEventSink;
+ private final Map<ReportedExecutionKey, ExecutionTraceContext> activeReportedExecutions =
+ new HashMap<>();
+
+ public EventLogComponentExecutionListener(
+ ExecutionTraceContext actionTraceContext, ExecutionEventSink executionEventSink) {
+ this.actionTraceContext = actionTraceContext;
+ this.executionEventSink = executionEventSink;
+ }
+
+ @Override
+ public void onComponentExecution(
+ String entityType, String entityName, Map<String, Object> entityMetadata, Event event) {
+ ReportedExecutionKey key = new ReportedExecutionKey(entityType, entityName, entityMetadata);
+ ExecutionTraceContext reportTraceContext;
+ if (ExecutionLifecycleEvents.EXECUTION_STARTED_EVENT_TYPE.equals(event.getType())) {
+ reportTraceContext =
+ actionTraceContext.childExecution(
+ entityType, entityName, key.getEntityMetadata());
+ ExecutionTraceContext previous = activeReportedExecutions.put(key, reportTraceContext);
+ if (previous != null) {
+ LOG.debug(
+ "Execution start report for {}:{} replaced an active report with the same metadata.",
+ entityType,
+ entityName);
+ }
+ } else {
+ reportTraceContext = activeReportedExecutions.remove(key);
+ if (reportTraceContext == null) {
+ LOG.debug(
+ "Execution terminal report for {}:{} has no matching start report; emitting it with a new execution id.",
+ entityType,
+ entityName);
+ reportTraceContext =
+ actionTraceContext.childExecution(
+ entityType, entityName, key.getEntityMetadata());
+ }
+ }
+
+ executionEventSink.emit(event, reportTraceContext);
+ }
+}
diff --git a/runtime/src/main/java/org/apache/flink/agents/runtime/trace/EventLogTaskLifecycleListener.java b/runtime/src/main/java/org/apache/flink/agents/runtime/trace/EventLogTaskLifecycleListener.java
new file mode 100644
index 000000000..376d27f59
--- /dev/null
+++ b/runtime/src/main/java/org/apache/flink/agents/runtime/trace/EventLogTaskLifecycleListener.java
@@ -0,0 +1,64 @@
+/*
+ * 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.trace;
+
+import org.apache.flink.agents.api.trace.ExecutionLifecycleEvents;
+import org.apache.flink.agents.api.trace.ExecutionReporter;
+import org.apache.flink.agents.runtime.lifecycle.TaskLifecycleListener;
+import org.apache.flink.agents.runtime.operator.ActionTask;
+import org.apache.flink.annotation.Internal;
+
+/**
+ * Bridges the operator's action lifecycle callbacks onto the event log, emitting the execution
+ * lifecycle events independently from the business-event router.
+ */
+@Internal
+public final class EventLogTaskLifecycleListener implements TaskLifecycleListener {
+
+ private final ExecutionEventSink executionEventSink;
+
+ public EventLogTaskLifecycleListener(ExecutionEventSink executionEventSink) {
+ this.executionEventSink = executionEventSink;
+ }
+
+ @Override
+ public void onActionStarted(ActionTask task) {
+ executionEventSink.emit(
+ ExecutionLifecycleEvents.executionStarted(), task.getTraceContext());
+ }
+
+ @Override
+ public void onActionReused(ActionTask task) {
+ executionEventSink.emit(ExecutionLifecycleEvents.executionReused(), task.getTraceContext());
+ }
+
+ @Override
+ public void onActionFinished(ActionTask task) {
+ executionEventSink.emit(
+ ExecutionLifecycleEvents.executionFinished(), task.getTraceContext());
+ }
+
+ @Override
+ public void onActionFailed(ActionTask task, Throwable error) {
+ executionEventSink.emit(
+ ExecutionLifecycleEvents.executionFailed(
+ error, ExecutionReporter.ProblemCategories.ACTION_EXECUTION_FAILED),
+ task.getTraceContext());
+ }
+}
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 b51f7e0c8..c39498d25 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
@@ -36,19 +36,25 @@ import org.apache.flink.agents.api.resource.python.PythonResourceAdapter;
import org.apache.flink.agents.api.resource.python.PythonResourceWrapper;
import org.apache.flink.agents.api.skills.SkillSourceSpec;
import org.apache.flink.agents.api.skills.Skills;
+import org.apache.flink.agents.api.subagent.SubagentFuture;
import org.apache.flink.agents.api.vectorstores.Document;
import org.apache.flink.agents.api.vectorstores.VectorStoreQuery;
import org.apache.flink.agents.api.vectorstores.VectorStoreQueryResult;
import org.apache.flink.agents.plan.AgentPlan;
+import org.apache.flink.agents.plan.resourceprovider.JavaSerializableResourceProvider;
+import org.apache.flink.agents.plan.resourceprovider.ResourceProvider;
+import org.apache.flink.agents.runtime.python.utils.PythonActionExecutor;
import org.apache.flink.agents.runtime.resource.ResourceContextImpl;
import org.apache.flink.agents.runtime.skill.AgentSkill;
import org.apache.flink.agents.runtime.skill.SkillManager;
import org.apache.flink.agents.runtime.skill.SkillRepository;
import org.apache.flink.agents.runtime.skill.SkillSourceRegistry;
+import org.apache.flink.agents.runtime.subagent.BaseSubagentSetup;
import org.junit.jupiter.api.Test;
import pemja.core.object.PyObject;
import java.lang.reflect.Field;
+import java.util.Collections;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
@@ -59,6 +65,7 @@ import static org.assertj.core.api.Assertions.assertThatThrownBy;
import static org.assertj.core.api.Assertions.catchThrowable;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
/** Tests for {@link ResourceCache}. */
public class ResourceCacheTest {
@@ -220,6 +227,57 @@ public class ResourceCacheTest {
}
}
+ /** Stands in for a handle to a resource that the Python runtime owns. */
+ public static class TestPythonHandle extends Resource {
+ @Override
+ public ResourceType getResourceType() {
+ return ResourceType.CHAT_MODEL;
+ }
+ }
+
+ @Test
+ public void testEagerMaterializeResolvesEveryJavaOwnedResourceOfTheType() throws Exception {
+ TestAgentWithResources agent = new TestAgentWithResources();
+ AgentPlan agentPlan = new AgentPlan(agent);
+ ResourceCache cache = new ResourceCache(agentPlan.getResourceProviders());
+
+ List<Resource> materialized = cache.eagerMaterialize(ResourceType.TOOL);
+
+ assertThat(materialized).hasSize(2).allMatch(resource -> resource instanceof TestTool);
+ assertThat(materialized).contains(cache.getResource("myTool", ResourceType.TOOL));
+ assertThat(materialized).contains(cache.getResource("anotherTool", ResourceType.TOOL));
+ }
+
+ @Test
+ public void testEagerMaterializeAsksThePythonRuntimeForTheResourcesItOwns() throws Exception {
+ TestAgentWithResources agent = new TestAgentWithResources();
+ AgentPlan agentPlan = new AgentPlan(agent);
+ ResourceCache cache = new ResourceCache(agentPlan.getResourceProviders());
+ TestPythonHandle handle = new TestPythonHandle();
+ // No Python resource adapter is wired, so resolving the Python provider here would fail:
+ // the type materializes only because the Python runtime is asked for its own resources.
+ PythonActionExecutor pythonActionExecutor = mock(PythonActionExecutor.class);
+ when(pythonActionExecutor.eagerMaterialize(ResourceType.CHAT_MODEL))
+ .thenReturn(Collections.singletonMap("pythonChatModel", handle));
+ cache.setPythonActionExecutor(pythonActionExecutor);
+
+ List<Resource> materialized = cache.eagerMaterialize(ResourceType.CHAT_MODEL);
+
+ assertThat(materialized).hasSize(2).contains(handle);
+ assertThat(cache.getResource("pythonChatModel", ResourceType.CHAT_MODEL)).isSameAs(handle);
+ }
+
+ @Test
+ public void testEagerMaterializeFailsWhenNoPythonRuntimeWasInitialized() throws Exception {
+ TestAgentWithResources agent = new TestAgentWithResources();
+ AgentPlan agentPlan = new AgentPlan(agent);
+ ResourceCache cache = new ResourceCache(agentPlan.getResourceProviders());
+
+ assertThatThrownBy(() -> cache.eagerMaterialize(ResourceType.CHAT_MODEL))
+ .isInstanceOf(IllegalStateException.class)
+ .hasMessageContaining("declared in Python but no Python runtime was initialized");
+ }
+
@Test
public void testGetResourceNotFound() throws Exception {
Agent agent = new Agent();
@@ -467,4 +525,32 @@ public class ResourceCacheTest {
}
}
}
+
+ /** Test Java sub-agent setup, registered as an AGENT resource. */
+ public static class TestAgentSetup extends BaseSubagentSetup {
+ @Override
+ public SubagentFuture submit(
+ RunnerContext ctx, Object prompt, String sessionId, String callId) {
+ return null;
+ }
+ }
+
+ @Test
+ public void testMaterializingAnAgentInjectsTheResourceNameAsSubagentName() throws Exception {
+ Map<ResourceType, Map<String, ResourceProvider>> providers = new HashMap<>();
+ Map<String, ResourceProvider> agentProviders = new HashMap<>();
+ agentProviders.put(
+ "reviewer",
+ JavaSerializableResourceProvider.createResourceProvider(
+ "reviewer", ResourceType.AGENT, new TestAgentSetup()));
+ providers.put(ResourceType.AGENT, agentProviders);
+
+ ResourceCache cache = new ResourceCache(providers);
+ List<Resource> materialized = cache.eagerMaterialize(ResourceType.AGENT);
+
+ assertThat(materialized).hasSize(1);
+ assertThat(materialized.get(0)).isInstanceOf(TestAgentSetup.class);
+ // The framework owns the identity: the resource name becomes the sub-agent name.
+ assertThat(((TestAgentSetup) materialized.get(0)).getSubagentName()).isEqualTo("reviewer");
+ }
}
diff --git a/runtime/src/test/java/org/apache/flink/agents/runtime/context/RunnerContextImplExecutionReporterTest.java b/runtime/src/test/java/org/apache/flink/agents/runtime/context/RunnerContextImplExecutionReporterTest.java
index be3090ae5..68857cad0 100644
--- a/runtime/src/test/java/org/apache/flink/agents/runtime/context/RunnerContextImplExecutionReporterTest.java
+++ b/runtime/src/test/java/org/apache/flink/agents/runtime/context/RunnerContextImplExecutionReporterTest.java
@@ -20,118 +20,133 @@ package org.apache.flink.agents.runtime.context;
import org.apache.flink.agents.api.Event;
import org.apache.flink.agents.api.trace.ExecutionLifecycleEvents;
import org.apache.flink.agents.api.trace.ExecutionReporter;
-import org.apache.flink.agents.api.trace.ExecutionTraceContext;
import org.apache.flink.agents.plan.AgentPlan;
+import org.apache.flink.agents.runtime.lifecycle.ComponentExecutionListener;
import org.apache.flink.agents.runtime.python.context.PythonRunnerContextImpl;
-import org.apache.flink.agents.runtime.trace.ReportedExecutionKey;
import org.junit.jupiter.api.Test;
+import javax.annotation.Nullable;
+
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatCode;
-/** Tests for execution reports emitted from {@link RunnerContextImpl}. */
+/** Tests for execution reports fanned out from {@link RunnerContextImpl} to its listeners. */
class RunnerContextImplExecutionReporterTest {
@Test
- void reportedExecutionReusesChildTraceContextBetweenStartAndFinish() throws Exception {
- List<RecordedReport> reports = new ArrayList<>();
+ void reportsFanOutToComponentExecutionListeners() throws Exception {
+ RecordingComponentListener listener = new RecordingComponentListener();
RunnerContextImpl runnerContext =
new RunnerContextImpl(null, () -> {}, emptyAgentPlan(), null, "job");
- ExecutionTraceContext actionTraceContext =
- ExecutionTraceContext.forInputRun("business-key", "agent")
- .childExecution("action", "chat_model_action");
- runnerContext.setExecutionEventSink(
- (event, context) -> reports.add(new RecordedReport(event, context)));
- runnerContext.switchActionContext(
- "chat_model_action", null, "business-key", actionTraceContext, new HashMap<>());
+ switchToChatModelAction(runnerContext, List.of(listener));
runnerContext.reportExecutionStarted(
- ExecutionReporter.EntityTypes.LLM, "model-a", Map.of());
+ ExecutionReporter.EntityTypes.LLM, "model-a", Map.of("temperature", 0.7));
runnerContext.reportExecutionSucceeded(
- ExecutionReporter.EntityTypes.LLM, "model-a", Map.of());
-
- assertThat(reports).hasSize(2);
- RecordedReport started = reports.get(0);
- RecordedReport finished = reports.get(1);
-
- assertThat(started.event.getType())
- .isEqualTo(ExecutionLifecycleEvents.EXECUTION_STARTED_EVENT_TYPE);
- assertThat(finished.event.getType())
- .isEqualTo(ExecutionLifecycleEvents.EXECUTION_FINISHED_EVENT_TYPE);
- assertThat(started.status()).isEqualTo(ExecutionLifecycleEvents.STATUS_STARTED);
- assertThat(finished.status()).isEqualTo(ExecutionLifecycleEvents.STATUS_SUCCESS);
-
- assertThat(started.traceContext().getExecutionId()).isNotBlank();
- assertThat(finished.traceContext().getExecutionId())
- .isEqualTo(started.traceContext().getExecutionId());
- assertThat(started.traceContext().getParentExecutionId())
- .isEqualTo(actionTraceContext.getExecutionId());
- assertThat(started.traceContext().getEntityType())
- .isEqualTo(ExecutionReporter.EntityTypes.LLM);
- assertThat(started.traceContext().getEntityName()).isEqualTo("model-a");
+ ExecutionReporter.EntityTypes.LLM, "model-a", Map.of("temperature", 0.7));
+
+ assertThat(listener.started).hasSize(1);
+ assertThat(listener.started.get(0))
+ .containsExactly(
+ ExecutionReporter.EntityTypes.LLM, "model-a", Map.of("temperature", 0.7));
+ assertThat(listener.succeeded).hasSize(1);
+ assertThat(listener.succeeded.get(0))
+ .containsExactly(
+ ExecutionReporter.EntityTypes.LLM, "model-a", Map.of("temperature", 0.7));
}
@Test
- void reportedExecutionStateFollowsActionContextAcrossSwitches() throws Exception {
- List<RecordedReport> reports = new ArrayList<>();
+ void failedReportResolvesRootCauseTypeAndMessage() throws Exception {
+ RecordingComponentListener listener = new RecordingComponentListener();
RunnerContextImpl runnerContext =
new RunnerContextImpl(null, () -> {}, emptyAgentPlan(), null, "job");
- runnerContext.setExecutionEventSink(
- (event, context) -> reports.add(new RecordedReport(event, context)));
-
- ExecutionTraceContext actionA =
- ExecutionTraceContext.forInputRun("business-key", "agent")
- .childExecution("action", "chat_model_action");
- ExecutionTraceContext actionB =
- ExecutionTraceContext.forInputRun("business-key", "agent")
- .childExecution("action", "tool_call_action");
- Map<ReportedExecutionKey, ExecutionTraceContext> activeReportsA = new HashMap<>();
- Map<ReportedExecutionKey, ExecutionTraceContext> activeReportsB = new HashMap<>();
+ switchToChatModelAction(runnerContext, List.of(listener));
- runnerContext.switchActionContext(
- "chat_model_action", null, "business-key", actionA, activeReportsA);
- runnerContext.reportExecutionStarted(
- ExecutionReporter.EntityTypes.LLM, "model-a", Map.of());
+ runnerContext.reportExecutionFailed(
+ ExecutionReporter.EntityTypes.TOOL,
+ "search",
+ Map.of("toolCallId", "call-1"),
+ new RuntimeException(new IllegalStateException("backend down")),
+ ExecutionReporter.ProblemCategories.TOOL_CALL_FAILED);
- runnerContext.switchActionContext(
- "tool_call_action", null, "business-key", actionB, activeReportsB);
- runnerContext.reportExecutionStarted(
- ExecutionReporter.EntityTypes.TOOL, "search", Map.of("toolCallId", "call-1"));
+ assertThat(listener.failed).hasSize(1);
+ RecordedFailure failure = listener.failed.get(0);
+ assertThat(failure.entityType).isEqualTo(ExecutionReporter.EntityTypes.TOOL);
+ assertThat(failure.entityName).isEqualTo("search");
+ assertThat(failure.entityMetadata).containsEntry("toolCallId", "call-1");
+ assertThat(failure.errorType).isEqualTo(IllegalStateException.class.getName());
+ assertThat(failure.errorMessage).isEqualTo("backend down");
+ assertThat(failure.problemCategory)
+ .isEqualTo(ExecutionReporter.ProblemCategories.TOOL_CALL_FAILED);
+ }
+ @Test
+ void throwingListenerNeverFailsTheReportingCall() throws Exception {
+ RecordingComponentListener receiver = new RecordingComponentListener();
+ ComponentExecutionListener thrower =
+ (entityType, entityName, entityMetadata, event) -> {
+ throw new IllegalStateException("listener boom");
+ };
+ RunnerContextImpl runnerContext =
+ new RunnerContextImpl(null, () -> {}, emptyAgentPlan(), null, "job");
+ switchToChatModelAction(runnerContext, List.of(thrower, receiver));
+
+ assertThatCode(
+ () -> {
+ runnerContext.reportExecutionStarted(
+ ExecutionReporter.EntityTypes.LLM, "model-a", Map.of());
+ runnerContext.reportExecutionSucceeded(
+ ExecutionReporter.EntityTypes.LLM, "model-a", Map.of());
+ runnerContext.reportExecutionFailed(
+ ExecutionReporter.EntityTypes.LLM,
+ "model-a",
+ Map.of(),
+ new IllegalStateException("call failed"),
+ null);
+ })
+ .doesNotThrowAnyException();
+
+ // The throwing listener is skipped; the remaining listener still receives every report.
+ assertThat(receiver.started).hasSize(1);
+ assertThat(receiver.succeeded).hasSize(1);
+ assertThat(receiver.failed).hasSize(1);
+ }
+
+ @Test
+ void reportingWithoutListenersIsANoOp() throws Exception {
+ RunnerContextImpl runnerContext =
+ new RunnerContextImpl(null, () -> {}, emptyAgentPlan(), null, "job");
runnerContext.switchActionContext(
- "chat_model_action", null, "business-key", actionA, activeReportsA);
- runnerContext.reportExecutionSucceeded(
- ExecutionReporter.EntityTypes.LLM, "model-a", Map.of());
-
- assertThat(reports).hasSize(3);
- RecordedReport actionAStarted = reports.get(0);
- RecordedReport actionBStarted = reports.get(1);
- RecordedReport actionAFinished = reports.get(2);
-
- assertThat(actionAFinished.traceContext().getExecutionId())
- .isEqualTo(actionAStarted.traceContext().getExecutionId());
- assertThat(actionAFinished.traceContext().getParentExecutionId())
- .isEqualTo(actionA.getExecutionId());
- assertThat(actionBStarted.traceContext().getExecutionId())
- .isNotEqualTo(actionAStarted.traceContext().getExecutionId());
+ "chat_model_action", null, new ArrayList<>(), "business-key", "obs-1", false, null);
+
+ assertThatCode(
+ () -> {
+ runnerContext.reportExecutionStarted(
+ ExecutionReporter.EntityTypes.LLM, "model-a", Map.of());
+ runnerContext.reportExecutionSucceeded(
+ ExecutionReporter.EntityTypes.LLM, "model-a", Map.of());
+ })
+ .doesNotThrowAnyException();
}
@Test
void pythonReporterBridgePreservesMetadataAndPythonErrorFields() throws Exception {
- List<RecordedReport> reports = new ArrayList<>();
+ RecordingComponentListener listener = new RecordingComponentListener();
PythonRunnerContextImpl runnerContext =
new PythonRunnerContextImpl(null, () -> {}, emptyAgentPlan(), null, "job");
- ExecutionTraceContext actionTraceContext =
- ExecutionTraceContext.forInputRun("business-key", "agent")
- .childExecution("action", "tool_call_action");
- runnerContext.setExecutionEventSink(
- (event, context) -> reports.add(new RecordedReport(event, context)));
runnerContext.switchActionContext(
- "tool_call_action", null, "business-key", actionTraceContext, new HashMap<>());
+ "tool_call_action",
+ null,
+ new ArrayList<>(),
+ "business-key",
+ "obs-1",
+ false,
+ List.of(listener));
String metadata = "{\"toolCallId\":\"call-1\",\"toolType\":\"function\"}";
runnerContext.reportExecutionStartedJson(
@@ -144,42 +159,85 @@ class RunnerContextImplExecutionReporterTest {
"bad response",
ExecutionReporter.ProblemCategories.TOOL_CALL_FAILED);
- assertThat(reports).hasSize(2);
- RecordedReport started = reports.get(0);
- RecordedReport failed = reports.get(1);
-
- assertThat(failed.traceContext().getExecutionId())
- .isEqualTo(started.traceContext().getExecutionId());
- assertThat(failed.traceContext().getEntityMetadata())
+ assertThat(listener.started).hasSize(1);
+ assertThat(listener.started.get(0).get(2))
+ .asInstanceOf(org.assertj.core.api.InstanceOfAssertFactories.MAP)
.containsEntry("toolCallId", "call-1")
.containsEntry("toolType", "function");
- assertThat(failed.event.getType())
- .isEqualTo(ExecutionLifecycleEvents.EXECUTION_FAILED_EVENT_TYPE);
- assertThat(failed.event.getAttr("errorType")).isEqualTo("builtins.ValueError");
- assertThat(failed.event.getAttr("errorMessage")).isEqualTo("bad response");
- assertThat(failed.event.getAttr(ExecutionLifecycleEvents.PROBLEM_CATEGORY_ATTRIBUTE))
+
+ assertThat(listener.failed).hasSize(1);
+ RecordedFailure failure = listener.failed.get(0);
+ // Python reports cross the bridge as strings and must reach listeners verbatim.
+ assertThat(failure.errorType).isEqualTo("builtins.ValueError");
+ assertThat(failure.errorMessage).isEqualTo("bad response");
+ assertThat(failure.problemCategory)
.isEqualTo(ExecutionReporter.ProblemCategories.TOOL_CALL_FAILED);
}
+ private static void switchToChatModelAction(
+ RunnerContextImpl runnerContext, List<ComponentExecutionListener> listeners) {
+ runnerContext.switchActionContext(
+ "chat_model_action",
+ null,
+ new ArrayList<>(),
+ "business-key",
+ "obs-1",
+ false,
+ listeners);
+ }
+
private static AgentPlan emptyAgentPlan() {
return new AgentPlan(new HashMap<>(), new HashMap<>());
}
- private static class RecordedReport {
- private final Event event;
- private final ExecutionTraceContext traceContext;
-
- private RecordedReport(Event event, ExecutionTraceContext traceContext) {
- this.event = event;
- this.traceContext = traceContext;
- }
-
- private ExecutionTraceContext traceContext() {
- return traceContext;
+ /** Records the raw arguments of every component report it receives. */
+ private static final class RecordingComponentListener implements ComponentExecutionListener {
+ private final List<List<Object>> started = new ArrayList<>();
+ private final List<List<Object>> succeeded = new ArrayList<>();
+ private final List<RecordedFailure> failed = new ArrayList<>();
+
+ @Override
+ public void onComponentExecution(
+ String entityType,
+ String entityName,
+ Map<String, Object> entityMetadata,
+ Event event) {
+ switch (event.getType()) {
+ case ExecutionLifecycleEvents.EXECUTION_STARTED_EVENT_TYPE:
+ started.add(List.of(entityType, entityName, entityMetadata));
+ break;
+ case ExecutionLifecycleEvents.EXECUTION_FINISHED_EVENT_TYPE:
+ succeeded.add(List.of(entityType, entityName, entityMetadata));
+ break;
+ case ExecutionLifecycleEvents.EXECUTION_FAILED_EVENT_TYPE:
+ failed.add(new RecordedFailure(entityType, entityName, entityMetadata, event));
+ break;
+ default:
+ throw new AssertionError("Unexpected event type " + event.getType());
+ }
}
+ }
- private String status() {
- return (String) event.getAttr(ExecutionLifecycleEvents.STATUS_ATTRIBUTE);
+ private static final class RecordedFailure {
+ private final String entityType;
+ private final String entityName;
+ private final Map<String, Object> entityMetadata;
+ private final String errorType;
+ @Nullable private final String errorMessage;
+ @Nullable private final String problemCategory;
+
+ private RecordedFailure(
+ String entityType,
+ String entityName,
+ Map<String, Object> entityMetadata,
+ Event event) {
+ this.entityType = entityType;
+ this.entityName = entityName;
+ this.entityMetadata = entityMetadata;
+ this.errorType = (String) event.getAttr("errorType");
+ this.errorMessage = (String) event.getAttr("errorMessage");
+ this.problemCategory =
+ (String) event.getAttr(ExecutionLifecycleEvents.PROBLEM_CATEGORY_ATTRIBUTE);
}
}
}
diff --git a/runtime/src/test/java/org/apache/flink/agents/runtime/context/RunnerContextPendingEventsContractTest.java b/runtime/src/test/java/org/apache/flink/agents/runtime/context/RunnerContextPendingEventsContractTest.java
new file mode 100644
index 000000000..e6f3be5f2
--- /dev/null
+++ b/runtime/src/test/java/org/apache/flink/agents/runtime/context/RunnerContextPendingEventsContractTest.java
@@ -0,0 +1,86 @@
+/*
+ * 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.context;
+
+import org.apache.flink.agents.api.Event;
+import org.apache.flink.agents.api.InputEvent;
+import org.apache.flink.agents.plan.AgentPlan;
+import org.apache.flink.agents.runtime.metrics.FlinkAgentsMetricGroupImpl;
+import org.apache.flink.runtime.metrics.groups.UnregisteredMetricGroups;
+import org.junit.jupiter.api.Test;
+
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.List;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
+/** Tests the per-task pending-event isolation contract. */
+class RunnerContextPendingEventsContractTest {
+
+ @Test
+ void emittedEventsDrainAndBufferIsClearBeforeTaskSwitch() {
+ RunnerContextImpl context = newContext();
+ RunnerContextImpl.MemoryContext memoryA = new RunnerContextImpl.MemoryContext(null, null);
+ RunnerContextImpl.MemoryContext memoryB = new RunnerContextImpl.MemoryContext(null, null);
+ List<Event> bufferA = new ArrayList<>();
+ List<Event> bufferB = new ArrayList<>();
+ Event eventA = new InputEvent(1L);
+
+ context.switchActionContext("action-a", memoryA, bufferA, "key-a", "obs-a", false, null);
+ context.sendEvent(eventA);
+ assertThat(context.drainEvents(null)).containsExactly(eventA);
+ context.checkNoPendingEvents();
+
+ context.switchActionContext("action-b", memoryB, bufferB, "key-b", "obs-b", false, null);
+ assertThat(context.drainEvents(null)).isEmpty();
+ }
+
+ @Test
+ void bufferedEventsStayIsolatedPerTaskAcrossContextSwitches() {
+ RunnerContextImpl context = newContext();
+ RunnerContextImpl.MemoryContext memoryA = new RunnerContextImpl.MemoryContext(null, null);
+ RunnerContextImpl.MemoryContext memoryB = new RunnerContextImpl.MemoryContext(null, null);
+ List<Event> bufferA = new ArrayList<>();
+ List<Event> bufferB = new ArrayList<>();
+ Event eventA = new InputEvent(1L);
+
+ context.switchActionContext("action-a", memoryA, bufferA, "key-a", "obs-a", false, null);
+ context.sendEvent(eventA);
+
+ // Switching to another action task now exposes that task's own (empty) buffer: action-a's
+ // event stays isolated in bufferA and cannot contaminate action-b, even though action-a
+ // yielded with an undrained buffer.
+ context.switchActionContext("action-b", memoryB, bufferB, "key-b", "obs-b", false, null);
+ assertThat(context.drainEvents(null)).isEmpty();
+
+ // Switching back to action-a still sees its buffered event.
+ context.switchActionContext("action-a", memoryA, bufferA, "key-a", "obs-a", false, null);
+ assertThat(context.drainEvents(null)).containsExactly(eventA);
+ }
+
+ private static RunnerContextImpl newContext() {
+ return new RunnerContextImpl(
+ new FlinkAgentsMetricGroupImpl(
+ UnregisteredMetricGroups.createUnregisteredOperatorMetricGroup()),
+ () -> {},
+ new AgentPlan(new HashMap<>(), new HashMap<>()),
+ null,
+ "job");
+ }
+}
diff --git a/runtime/src/test/java/org/apache/flink/agents/runtime/context/TestMemoryObservationFlush.java b/runtime/src/test/java/org/apache/flink/agents/runtime/context/TestMemoryObservationFlush.java
index 0bb01ad31..f3f34f99e 100644
--- a/runtime/src/test/java/org/apache/flink/agents/runtime/context/TestMemoryObservationFlush.java
+++ b/runtime/src/test/java/org/apache/flink/agents/runtime/context/TestMemoryObservationFlush.java
@@ -33,6 +33,7 @@ import org.junit.jupiter.api.Test;
import javax.annotation.Nullable;
+import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
@@ -84,9 +85,11 @@ class TestMemoryObservationFlush {
new RunnerContextImpl.MemoryContext(
new CachedMemoryStore(new ForTestMemoryMapState<>()),
new CachedMemoryStore(new ForTestMemoryMapState<>())),
+ new ArrayList<>(),
contextKey,
"observation-1",
- suppressed);
+ suppressed,
+ null);
return context;
}
@@ -222,9 +225,11 @@ class TestMemoryObservationFlush {
new RunnerContextImpl.MemoryContext(
new CachedMemoryStore(new ForTestMemoryMapState<>()),
new CachedMemoryStore(new ForTestMemoryMapState<>())),
+ new ArrayList<>(),
"user-43",
"observation-2",
- true);
+ true,
+ null);
assertThat(ltm.configureCallCount).isEqualTo(1);
assertThat(ltm.switchCallCount).isEqualTo(2);
@@ -259,9 +264,11 @@ class TestMemoryObservationFlush {
new RunnerContextImpl.MemoryContext(
new CachedMemoryStore(new ForTestMemoryMapState<>()),
new CachedMemoryStore(new ForTestMemoryMapState<>())),
+ new ArrayList<>(),
"user-42",
"observation-2",
- false);
+ false,
+ null);
ltm.record("user-42", "observation-2", "b", "from-b");
LongTermUpdateEvent bEvent =
@@ -274,9 +281,11 @@ class TestMemoryObservationFlush {
new RunnerContextImpl.MemoryContext(
new CachedMemoryStore(new ForTestMemoryMapState<>()),
new CachedMemoryStore(new ForTestMemoryMapState<>())),
+ new ArrayList<>(),
"user-42",
"observation-1",
- false);
+ false,
+ null);
context.discardMemoryObservation();
assertThat(context.drainEventsAtActionFinish(null)).isEmpty();
assertThat(ltm.pendingRecords).isEmpty();
diff --git a/runtime/src/test/java/org/apache/flink/agents/runtime/lifecycle/TaskLifecycleListenerNotificationTest.java b/runtime/src/test/java/org/apache/flink/agents/runtime/lifecycle/TaskLifecycleListenerNotificationTest.java
new file mode 100644
index 000000000..ff9b7f24e
--- /dev/null
+++ b/runtime/src/test/java/org/apache/flink/agents/runtime/lifecycle/TaskLifecycleListenerNotificationTest.java
@@ -0,0 +1,532 @@
+/*
+ * 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.lifecycle;
+
+import org.apache.flink.agents.api.Event;
+import org.apache.flink.agents.api.EventType;
+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.DurableCallable;
+import org.apache.flink.agents.api.context.RunnerContext;
+import org.apache.flink.agents.plan.AgentPlan;
+import org.apache.flink.agents.runtime.actionstate.ActionState;
+import org.apache.flink.agents.runtime.actionstate.ActionStateStore;
+import org.apache.flink.agents.runtime.actionstate.InMemoryActionStateStore;
+import org.apache.flink.agents.runtime.async.ContinuationActionExecutor;
+import org.apache.flink.agents.runtime.operator.ActionExecutionOperator;
+import org.apache.flink.agents.runtime.operator.ActionExecutionOperatorFactory;
+import org.apache.flink.agents.runtime.operator.ActionTask;
+import org.apache.flink.api.common.typeinfo.TypeInformation;
+import org.apache.flink.api.java.functions.KeySelector;
+import org.apache.flink.runtime.checkpoint.OperatorSubtaskState;
+import org.apache.flink.streaming.api.operators.StreamOperator;
+import org.apache.flink.streaming.api.operators.StreamOperatorParameters;
+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.ArrayList;
+import java.util.List;
+import java.util.concurrent.CopyOnWriteArrayList;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
+/**
+ * Tests that {@link ActionExecutionOperator} broadcasts the record/task lifecycle events to
+ * injected {@link TaskLifecycleListener}s in the expected order, independently of any particular
+ * listener implementation.
+ */
+public class TaskLifecycleListenerNotificationTest {
+
+ @BeforeEach
+ void resetRecording() {
+ RecordingListener.EVENTS.clear();
+ }
+
+ /** Plain listener that records every lifecycle notification it receives. */
+ public static class RecordingListener implements TaskLifecycleListener {
+
+ static final List<String> EVENTS = new CopyOnWriteArrayList<>();
+
+ @Override
+ public void onRecordStart(Object key) {
+ EVENTS.add("recordStart:" + key);
+ }
+
+ @Override
+ public void onActionPrepared(ActionTask task) {
+ EVENTS.add("prepared:" + task.getAction().getName());
+ }
+
+ @Override
+ public void onActionStarted(ActionTask task) {
+ EVENTS.add("started:" + task.getAction().getName());
+ }
+
+ @Override
+ public void onActionTransferred(ActionTask from, ActionTask to) {
+ EVENTS.add(
+ "transferred:" + from.getAction().getName() + "->" + to.getAction().getName());
+ }
+
+ @Override
+ public void onActionFinishing(ActionTask task) {
+ EVENTS.add("finishing:" + task.getAction().getName());
+ }
+
+ @Override
+ public void onActionFinished(ActionTask task) {
+ EVENTS.add("finished:" + task.getAction().getName());
+ }
+
+ @Override
+ public void onActionReused(ActionTask task) {
+ EVENTS.add("reused:" + task.getAction().getName());
+ }
+
+ @Override
+ public void onActionFailed(ActionTask task, Throwable error) {
+ EVENTS.add("failed:" + task.getAction().getName());
+ }
+
+ @Override
+ public void onRecordFinished(Object key) {
+ EVENTS.add("recordFinished:" + key);
+ }
+ }
+
+ /** Agent with a plain synchronous action. */
+ public static class SyncAgent extends Agent {
+
+ @org.apache.flink.agents.api.annotation.Action(EventType.InputEvent)
+ public static void handleInput(Event event, RunnerContext context) {
+ Long input = (Long) InputEvent.fromEvent(event).getInput();
+ context.sendEvent(new OutputEvent(input * 2));
+ }
+ }
+
+ /** Agent whose input action suspends on a durable async call, forcing a task transfer. */
+ public static class AsyncAgent extends Agent {
+
+ @org.apache.flink.agents.api.annotation.Action(EventType.InputEvent)
+ public static void handleInput(Event event, RunnerContext context) throws Exception {
+ Long input = (Long) InputEvent.fromEvent(event).getInput();
+ Long result =
+ context.durableExecuteAsync(
+ new DurableCallable<Long>() {
+ @Override
+ public String getId() {
+ return "lifecycle-notification";
+ }
+
+ @Override
+ public Class<Long> getResultClass() {
+ return Long.class;
+ }
+
+ @Override
+ public Long call() {
+ try {
+ // Force the action to yield before the call completes,
+ // so the task is suspended and transferred.
+ Thread.sleep(50);
+ } catch (InterruptedException e) {
+ Thread.currentThread().interrupt();
+ }
+ return input * 2;
+ }
+ });
+ context.sendEvent(new OutputEvent(result));
+ }
+ }
+
+ @Test
+ void recordStartAndFinishedPairAroundSyncTasks() throws Exception {
+ AgentPlan plan = new AgentPlan(new SyncAgent());
+ try (KeyedOneInputStreamOperatorTestHarness<Long, Long, Object> testHarness =
+ new KeyedOneInputStreamOperatorTestHarness<>(
+ new ActionExecutionOperatorFactory(plan, true),
+ (KeySelector<Long, Long>) value -> value,
+ TypeInformation.of(Long.class))) {
+ testHarness.open();
+ ActionExecutionOperator<Long, Object> operator =
+ (ActionExecutionOperator<Long, Object>) testHarness.getOperator();
+ operator.addTaskLifecycleListener(new RecordingListener());
+
+ testHarness.processElement(new StreamRecord<>(7L));
+ operator.waitInFlightEventsFinished();
+
+ assertThat(RecordingListener.EVENTS)
+ .containsExactly(
+ "recordStart:7",
+ "prepared:handleInput",
+ "started:handleInput",
+ "finishing:handleInput",
+ "finished:handleInput",
+ "recordFinished:7");
+
+ // A second record on the same key starts and finishes its own round.
+ testHarness.processElement(new StreamRecord<>(7L));
+ operator.waitInFlightEventsFinished();
+
+ assertThat(RecordingListener.EVENTS)
+ .containsExactly(
+ "recordStart:7",
+ "prepared:handleInput",
+ "started:handleInput",
+ "finishing:handleInput",
+ "finished:handleInput",
+ "recordFinished:7",
+ "recordStart:7",
+ "prepared:handleInput",
+ "started:handleInput",
+ "finishing:handleInput",
+ "finished:handleInput",
+ "recordFinished:7");
+ }
+ }
+
+ /**
+ * Exposes the test-only {@link ActionExecutionOperatorFactory} constructor, which is
+ * package-private to the operator package, to tests in this package.
+ */
+ private static class TestableOperatorFactory
+ extends ActionExecutionOperatorFactory<Long, Object> {
+
+ TestableOperatorFactory(AgentPlan agentPlan, ActionStateStore actionStateStore) {
+ super(agentPlan, true, actionStateStore);
+ }
+ }
+
+ /**
+ * Registers listeners on the operator right after creation, before {@code initializeState} and
+ * {@code open} run, so notifications emitted while resuming in-flight work during {@code open}
+ * are captured as well.
+ */
+ private static class ListenerInjectingOperatorFactory
+ extends ActionExecutionOperatorFactory<Long, Object> {
+
+ private final List<TaskLifecycleListener> listeners;
+
+ ListenerInjectingOperatorFactory(AgentPlan agentPlan, TaskLifecycleListener listener) {
+ super(agentPlan, true);
+ this.listeners = new ArrayList<>();
+ this.listeners.add(listener);
+ }
+
+ @Override
+ public <T extends StreamOperator<Object>> T createStreamOperator(
+ StreamOperatorParameters<Object> parameters) {
+ T operator = super.createStreamOperator(parameters);
+ ActionExecutionOperator<Long, Object> actionOperator =
+ (ActionExecutionOperator<Long, Object>) operator;
+ listeners.forEach(actionOperator::addTaskLifecycleListener);
+ return operator;
+ }
+ }
+
+ /**
+ * Listener that captures the durable action-state picture observed at {@code onActionFinishing}
+ * and {@code onActionFinished} time, so the test can assert the finishing notification arrives
+ * before the completed state is persisted and the finished notification after.
+ */
+ private static class CompletionStateObservingListener implements TaskLifecycleListener {
+
+ private final InMemoryActionStateStore store;
+ private final List<String> events = new CopyOnWriteArrayList<>();
+ private volatile Integer stateCountAtFinishing = null;
+ private volatile Boolean anyStateCompletedAtFinishing = null;
+ private volatile Boolean anyStateCompletedAtFinished = null;
+
+ CompletionStateObservingListener(InMemoryActionStateStore store) {
+ this.store = store;
+ }
+
+ @Override
+ public void onActionFinishing(ActionTask task) {
+ List<ActionState> states = currentStates();
+ stateCountAtFinishing = states.size();
+ anyStateCompletedAtFinishing = states.stream().anyMatch(ActionState::isCompleted);
+ events.add("finishing:" + task.getAction().getName());
+ }
+
+ @Override
+ public void onActionFinished(ActionTask task) {
+ anyStateCompletedAtFinished =
+ currentStates().stream().anyMatch(ActionState::isCompleted);
+ events.add("finished:" + task.getAction().getName());
+ }
+
+ private List<ActionState> currentStates() {
+ return store.getKeyedActionStates().values().stream()
+ .flatMap(perKey -> perKey.values().stream())
+ .collect(java.util.stream.Collectors.toList());
+ }
+ }
+
+ @Test
+ void actionFinishingIsNotifiedBeforeAndFinishedAfterCompletedStateIsDurable() throws Exception {
+ AgentPlan plan = new AgentPlan(new SyncAgent());
+ InMemoryActionStateStore actionStateStore = new InMemoryActionStateStore(false);
+ try (KeyedOneInputStreamOperatorTestHarness<Long, Long, Object> testHarness =
+ new KeyedOneInputStreamOperatorTestHarness<>(
+ new TestableOperatorFactory(plan, actionStateStore),
+ (KeySelector<Long, Long>) value -> value,
+ TypeInformation.of(Long.class))) {
+ testHarness.open();
+ ActionExecutionOperator<Long, Object> operator =
+ (ActionExecutionOperator<Long, Object>) testHarness.getOperator();
+ CompletionStateObservingListener listener =
+ new CompletionStateObservingListener(actionStateStore);
+ operator.addTaskLifecycleListener(listener);
+
+ testHarness.processElement(new StreamRecord<>(7L));
+ operator.waitInFlightEventsFinished();
+
+ assertThat(listener.events)
+ .containsExactly("finishing:handleInput", "finished:handleInput");
+ // The action state was initialized before the invocation, but the finishing
+ // notification must arrive before the completed state is persisted, so the
+ // listener still sees an uncompleted state at notification time.
+ assertThat(listener.stateCountAtFinishing).isEqualTo(1);
+ assertThat(listener.anyStateCompletedAtFinishing).isFalse();
+ // The finished notification arrives after the completed state is durable.
+ assertThat(listener.anyStateCompletedAtFinished).isTrue();
+
+ // After processing finishes, the durable state records the completion.
+ assertThat(
+ actionStateStore.getKeyedActionStates().values().stream()
+ .flatMap(perKey -> perKey.values().stream())
+ .allMatch(ActionState::isCompleted))
+ .isTrue();
+ }
+ }
+
+ @Test
+ void replayOfCompletedActionEmitsPreparedAndReusedPair() throws Exception {
+ AgentPlan plan = new AgentPlan(new SyncAgent());
+ InMemoryActionStateStore actionStateStore = new InMemoryActionStateStore(false);
+
+ // First execution runs the action and persists its completed state.
+ try (KeyedOneInputStreamOperatorTestHarness<Long, Long, Object> testHarness =
+ new KeyedOneInputStreamOperatorTestHarness<>(
+ new TestableOperatorFactory(plan, actionStateStore),
+ (KeySelector<Long, Long>) value -> value,
+ TypeInformation.of(Long.class))) {
+ testHarness.open();
+ ActionExecutionOperator<Long, Object> operator =
+ (ActionExecutionOperator<Long, Object>) testHarness.getOperator();
+
+ testHarness.processElement(new StreamRecord<>(7L));
+ operator.waitInFlightEventsFinished();
+ }
+
+ // Replay the same input against the persisted completed state: the invocation is
+ // skipped, but the prepared/reused pair must still be emitted so listener bookkeeping
+ // that opened on preparation is closed on the reuse path as well.
+ try (KeyedOneInputStreamOperatorTestHarness<Long, Long, Object> testHarness =
+ new KeyedOneInputStreamOperatorTestHarness<>(
+ new TestableOperatorFactory(plan, actionStateStore),
+ (KeySelector<Long, Long>) value -> value,
+ TypeInformation.of(Long.class))) {
+ testHarness.open();
+ ActionExecutionOperator<Long, Object> operator =
+ (ActionExecutionOperator<Long, Object>) testHarness.getOperator();
+ operator.addTaskLifecycleListener(new RecordingListener());
+
+ testHarness.processElement(new StreamRecord<>(7L));
+ operator.waitInFlightEventsFinished();
+
+ assertThat(RecordingListener.EVENTS)
+ .containsExactly(
+ "recordStart:7",
+ "prepared:handleInput",
+ "reused:handleInput",
+ "recordFinished:7");
+ }
+ }
+
+ /** Agent with no actions at all; input records trigger nothing. */
+ public static class EmptyAgent extends Agent {}
+
+ @Test
+ void resumedInFlightRecordReEmitsRecordStart() throws Exception {
+ AgentPlan plan = new AgentPlan(new SyncAgent());
+ OperatorSubtaskState snapshot;
+
+ // First execution: admit the record but snapshot before its tasks run, so the record
+ // is in flight (processing key + pending task) at snapshot time.
+ try (KeyedOneInputStreamOperatorTestHarness<Long, Long, Object> testHarness =
+ new KeyedOneInputStreamOperatorTestHarness<>(
+ new ActionExecutionOperatorFactory<>(plan, true),
+ (KeySelector<Long, Long>) value -> value,
+ TypeInformation.of(Long.class))) {
+ testHarness.open();
+ testHarness.processElement(new StreamRecord<>(7L));
+ snapshot = testHarness.snapshot(1L, 1L);
+ }
+
+ // Restore: the in-flight record resumes during open(), and its record start is
+ // re-emitted to align with the replayed task-level callbacks, giving listeners a
+ // paired bracket for the replayed round.
+ try (KeyedOneInputStreamOperatorTestHarness<Long, Long, Object> testHarness =
+ new KeyedOneInputStreamOperatorTestHarness<>(
+ new ListenerInjectingOperatorFactory(plan, new RecordingListener()),
+ (KeySelector<Long, Long>) value -> value,
+ TypeInformation.of(Long.class))) {
+ testHarness.initializeState(snapshot);
+ testHarness.open();
+ ActionExecutionOperator<Long, Object> operator =
+ (ActionExecutionOperator<Long, Object>) testHarness.getOperator();
+ operator.waitInFlightEventsFinished();
+
+ assertThat(RecordingListener.EVENTS)
+ .containsExactly(
+ "recordStart:7",
+ "prepared:handleInput",
+ "started:handleInput",
+ "finishing:handleInput",
+ "finished:handleInput",
+ "recordFinished:7");
+ }
+ }
+
+ @Test
+ void recordWithoutTriggeredActionsEmitsNoLifecycleEvents() throws Exception {
+ AgentPlan plan = new AgentPlan(new EmptyAgent());
+ try (KeyedOneInputStreamOperatorTestHarness<Long, Long, Object> testHarness =
+ new KeyedOneInputStreamOperatorTestHarness<>(
+ new ActionExecutionOperatorFactory(plan, true),
+ (KeySelector<Long, Long>) value -> value,
+ TypeInformation.of(Long.class))) {
+ testHarness.open();
+ ActionExecutionOperator<Long, Object> operator =
+ (ActionExecutionOperator<Long, Object>) testHarness.getOperator();
+ operator.addTaskLifecycleListener(new RecordingListener());
+
+ testHarness.processElement(new StreamRecord<>(3L));
+ operator.waitInFlightEventsFinished();
+
+ // No task was ever created, so the start/finished pair stays closed.
+ assertThat(RecordingListener.EVENTS).isEmpty();
+ }
+ }
+
+ @Test
+ void suspendedTaskEmitsTransferBeforeItsSuccessorIsPrepared() throws Exception {
+ AgentPlan plan = new AgentPlan(new AsyncAgent());
+ try (KeyedOneInputStreamOperatorTestHarness<Long, Long, Object> testHarness =
+ new KeyedOneInputStreamOperatorTestHarness<>(
+ new ActionExecutionOperatorFactory(plan, true),
+ (KeySelector<Long, Long>) value -> value,
+ TypeInformation.of(Long.class))) {
+ testHarness.open();
+ ActionExecutionOperator<Long, Object> operator =
+ (ActionExecutionOperator<Long, Object>) testHarness.getOperator();
+ operator.addTaskLifecycleListener(new RecordingListener());
+
+ testHarness.processElement(new StreamRecord<>(5L));
+ operator.waitInFlightEventsFinished();
+
+ List<StreamRecord<Object>> recordOutput =
+ (List<StreamRecord<Object>>) testHarness.getRecordOutput();
+ assertThat(recordOutput.size()).isEqualTo(1);
+ assertThat(recordOutput.get(0).getValue()).isEqualTo(10L);
+
+ if (ContinuationActionExecutor.isContinuationSupported()) {
+ // JDK 21+: the action suspends on the durable async call and the operator
+ // re-polls the suspended task, emitting a transferred->prepared pair per
+ // suspension round. Assert the invariant the test is named after: every
+ // suspension transfers the task before its successor is prepared, bounded by
+ // recordStart/prepared/started first and finishing/finished last. The started
+ // notification fires once per action execution, not per suspension round.
+ assertThat(RecordingListener.EVENTS)
+ .startsWith("recordStart:5", "prepared:handleInput")
+ .endsWith(
+ "finishing:handleInput",
+ "finished:handleInput",
+ "recordFinished:5");
+ List<String> middle =
+ RecordingListener.EVENTS.subList(3, RecordingListener.EVENTS.size() - 3);
+ assertThat(middle.get(0)).isEqualTo("started:handleInput");
+ List<String> suspensionRounds = middle.subList(1, middle.size());
+ assertThat(suspensionRounds)
+ .as("suspension rounds alternate transferred -> prepared")
+ .isNotEmpty();
+ assertThat(suspensionRounds.size() % 2).isZero();
+ for (int i = 0; i < suspensionRounds.size(); i += 2) {
+ assertThat(suspensionRounds.get(i))
+ .isEqualTo("transferred:handleInput->handleInput");
+ assertThat(suspensionRounds.get(i + 1)).isEqualTo("prepared:handleInput");
+ }
+ } else {
+ // JDK 11 fallback runs the action synchronously: no suspension, no transfer.
+ assertThat(RecordingListener.EVENTS)
+ .containsExactly(
+ "recordStart:5",
+ "prepared:handleInput",
+ "started:handleInput",
+ "finishing:handleInput",
+ "finished:handleInput",
+ "recordFinished:5");
+ }
+ }
+ }
+
+ /** Agent whose action always fails, exercising the failure notification path. */
+ public static class FailingAgent extends Agent {
+
+ @org.apache.flink.agents.api.annotation.Action(EventType.InputEvent)
+ public static void handleInput(Event event, RunnerContext context) {
+ throw new IllegalStateException("action boom");
+ }
+ }
+
+ @Test
+ void failedActionNotifiesFailureWithoutTerminalSuccessCallbacks() throws Exception {
+ AgentPlan plan = new AgentPlan(new FailingAgent());
+ try (KeyedOneInputStreamOperatorTestHarness<Long, Long, Object> testHarness =
+ new KeyedOneInputStreamOperatorTestHarness<>(
+ new ActionExecutionOperatorFactory(plan, true),
+ (KeySelector<Long, Long>) value -> value,
+ TypeInformation.of(Long.class))) {
+ testHarness.open();
+ ActionExecutionOperator<Long, Object> operator =
+ (ActionExecutionOperator<Long, Object>) testHarness.getOperator();
+ operator.addTaskLifecycleListener(new RecordingListener());
+
+ org.assertj.core.api.Assertions.assertThatThrownBy(
+ () -> {
+ testHarness.processElement(new StreamRecord<>(7L));
+ operator.waitInFlightEventsFinished();
+ })
+ .hasStackTraceContaining("action boom");
+
+ // The failure notification replaces the finishing/finished/recordFinished tail.
+ assertThat(RecordingListener.EVENTS)
+ .containsExactly(
+ "recordStart:7",
+ "prepared:handleInput",
+ "started:handleInput",
+ "failed:handleInput");
+ }
+ }
+}
diff --git a/runtime/src/test/java/org/apache/flink/agents/runtime/operator/ActionTaskContextManagerTest.java b/runtime/src/test/java/org/apache/flink/agents/runtime/operator/ActionTaskContextManagerTest.java
index a1ea1b0c8..827daf117 100644
--- a/runtime/src/test/java/org/apache/flink/agents/runtime/operator/ActionTaskContextManagerTest.java
+++ b/runtime/src/test/java/org/apache/flink/agents/runtime/operator/ActionTaskContextManagerTest.java
@@ -17,9 +17,10 @@
*/
package org.apache.flink.agents.runtime.operator;
+import org.apache.flink.agents.api.Event;
import org.apache.flink.agents.api.InputEvent;
+import org.apache.flink.agents.api.trace.ExecutionLifecycleEvents;
import org.apache.flink.agents.api.trace.ExecutionReporter;
-import org.apache.flink.agents.api.trace.ExecutionTraceContext;
import org.apache.flink.agents.plan.AgentPlan;
import org.apache.flink.agents.plan.actions.Action;
import org.apache.flink.agents.runtime.ResourceCache;
@@ -29,10 +30,10 @@ import org.apache.flink.agents.runtime.async.ContinuationActionExecutor;
import org.apache.flink.agents.runtime.async.ContinuationContext;
import org.apache.flink.agents.runtime.context.JavaRunnerContextImpl;
import org.apache.flink.agents.runtime.context.RunnerContextImpl;
+import org.apache.flink.agents.runtime.lifecycle.ComponentExecutionListener;
import org.apache.flink.agents.runtime.memory.InteranlBaseLongTermMemory;
import org.apache.flink.agents.runtime.memory.MemoryObjectImpl;
import org.apache.flink.agents.runtime.metrics.FlinkAgentsMetricGroupImpl;
-import org.apache.flink.agents.runtime.trace.ExecutionEventSink;
import org.apache.flink.api.common.serialization.SerializerConfigImpl;
import org.apache.flink.api.common.state.MapState;
import org.apache.flink.api.common.typeinfo.TypeInformation;
@@ -46,6 +47,8 @@ import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.function.Function;
import static org.assertj.core.api.Assertions.assertThat;
import static org.assertj.core.api.Assertions.assertThatThrownBy;
@@ -119,17 +122,21 @@ class ActionTaskContextManagerTest {
}
@Test
- void perTaskMapsAreIsolatedAcrossPutGetRemove() throws Exception {
+ void perTaskContextsAreIsolatedAcrossPutGetRemove() throws Exception {
try (ActionTaskContextManager mgr = new ActionTaskContextManager(1)) {
Action action = TestActions.noopAction();
- ActionTask t1 = new JavaActionTask("k", new InputEvent(1L), action);
- ActionTask t2 = new JavaActionTask("k", new InputEvent(2L), action);
+ ActionTask t1 = new JavaActionTask("k", new InputEvent(1L), action, 1L);
+ ActionTask t2 = new JavaActionTask("k", new InputEvent(2L), action, 1L);
+
+ // Contexts records are created explicitly; mutators never create one implicitly.
+ mgr.createContexts(t1);
+ mgr.createContexts(t2);
ContinuationContext c1 = new ContinuationContext();
mgr.putContinuationContext(t1, c1);
mgr.putPythonAwaitableRef(t2, "ref-2");
- // Cross-task isolation: each map only carries the entry it was given.
+ // Cross-task isolation: each contexts record only carries the entry it was given.
assertThat(mgr.getContinuationContext(t1)).isSameAs(c1);
assertThat(mgr.getContinuationContext(t2)).isNull();
assertThat(mgr.getPythonAwaitableRef(t1)).isNull();
@@ -137,11 +144,12 @@ class ActionTaskContextManagerTest {
assertThat(mgr.hasContinuationContext(t1)).isTrue();
assertThat(mgr.hasContinuationContext(t2)).isFalse();
- // Remove and re-check
- mgr.removeContinuationContext(t1);
- mgr.removePythonAwaitableRef(t2);
- assertThat(mgr.hasContinuationContext(t1)).isFalse();
- assertThat(mgr.getPythonAwaitableRef(t2)).isNull();
+ // Removing the whole contexts record wipes that task's contexts as a unit; the sibling
+ // is intact.
+ mgr.removeContexts(t1);
+ assertThat(mgr.getPythonAwaitableRef(t2)).isEqualTo("ref-2");
+ assertThat(mgr.hasContinuationContext(t2)).isFalse();
+ mgr.removeContexts(t2);
}
}
@@ -167,11 +175,11 @@ class ActionTaskContextManagerTest {
@Test
void createAndSetRunnerContextBuildsFreshMemoryContextOnFirstCall() throws Exception {
try (ActionTaskContextManager mgr = new ActionTaskContextManager(1)) {
- ActionTask t = new JavaActionTask("k", new InputEvent(1L), TestActions.noopAction());
+ ActionTask t =
+ new JavaActionTask("k", new InputEvent(1L), TestActions.noopAction(), 1L);
invokeCreateAndSetRunnerContext(mgr, t);
- // Production path: createAndSetRunnerContext at ActionTaskContextManager.java:210-218
- // — the else branch builds a fresh MemoryContext when the map has no entry.
+ // Production path: createAndSetRunnerContext pins the freshly created MemoryContext.
assertThat(t.getRunnerContext()).isInstanceOf(JavaRunnerContextImpl.class);
assertThat(t.getRunnerContext().getMemoryContext()).isNotNull();
}
@@ -181,12 +189,11 @@ class ActionTaskContextManagerTest {
void createAndSetRunnerContextReusesExistingMemoryContext() throws Exception {
try (ActionTaskContextManager mgr = new ActionTaskContextManager(1)) {
Action action = TestActions.noopAction();
- ActionTask from = new JavaActionTask("k", new InputEvent(1L), action);
- ActionTask to = new JavaActionTask("k", new InputEvent(2L), action);
+ ActionTask from = new JavaActionTask("k", new InputEvent(1L), action, 1L);
+ ActionTask to = new JavaActionTask("k", new InputEvent(2L), action, 1L);
- // Step 1: createAndSetRunnerContext(from) — runner context now carries a fresh
- // MemoryContext, but the map (actionTaskMemoryContexts) is still empty (production
- // code at lines 210-218 only reads from the map, never writes).
+ // Step 1: createAndSetRunnerContext(from) — runner context carries and pins a fresh
+ // MemoryContext.
invokeCreateAndSetRunnerContext(mgr, from);
RunnerContextImpl.MemoryContext fromMemCtx = from.getRunnerContext().getMemoryContext();
assertThat(fromMemCtx).isNotNull();
@@ -211,8 +218,8 @@ class ActionTaskContextManagerTest {
void sameKeyTasksSwitchLtmWithDistinctObservationIds() throws Exception {
try (ActionTaskContextManager mgr = new ActionTaskContextManager(1)) {
Action action = TestActions.noopAction();
- ActionTask suspended = new JavaActionTask("k", new InputEvent(1L), action);
- ActionTask sibling = new JavaActionTask("k", new InputEvent(1L), action);
+ ActionTask suspended = new JavaActionTask("k", new InputEvent(1L), action, 1L);
+ ActionTask sibling = new JavaActionTask("k", new InputEvent(1L), action, 1L);
InteranlBaseLongTermMemory ltm = mock(InteranlBaseLongTermMemory.class);
invokeCreateAndSetRunnerContext(mgr, suspended, ltm);
@@ -228,8 +235,8 @@ class ActionTaskContextManagerTest {
void transferContextsCopiesMemoryAndContinuationToNewTask() throws Exception {
try (ActionTaskContextManager mgr = new ActionTaskContextManager(1)) {
Action action = TestActions.noopAction();
- ActionTask from = new JavaActionTask("k", new InputEvent(1L), action);
- ActionTask to = new JavaActionTask("k", new InputEvent(2L), action);
+ ActionTask from = new JavaActionTask("k", new InputEvent(1L), action, 1L);
+ ActionTask to = new JavaActionTask("k", new InputEvent(2L), action, 1L);
// Populate `from`'s runner context with a MemoryContext and ContinuationContext.
invokeCreateAndSetRunnerContext(mgr, from);
@@ -237,87 +244,106 @@ class ActionTaskContextManagerTest {
assertThat(fromMemCtx).isNotNull();
from.markExecutionStartedEventEmitted();
- // transferContexts (ActionTaskContextManager.java:266-286) copies but does NOT
- // remove from source. The from-side continuation map is never populated (the
- // continuation lives on from's runner context until transfer copies it over for
- // `to`). Operator-side cleanup of `from`'s entries is the operator's
- // responsibility — see ActionExecutionOperator.java:366-369.
+ // Mirrors the production order: the operator removes the source record before
+ // transferring (ActionExecutionOperator). transferContexts must therefore extract
+ // everything from the source's runner context, not from its already-removed record.
+ mgr.removeContexts(from);
mgr.transferContexts(from, to, new DurableExecutionManager(null));
- // (a) The memory context entry for `to` is the same instance fromTask holds.
- RunnerContextImpl.MemoryContext toMemCtx = mgr.removeMemoryContext(to);
- assertThat(toMemCtx).isSameAs(fromMemCtx);
-
- // After remove, the map no longer has `to`'s entry.
- assertThat(mgr.removeMemoryContext(to)).isNull();
+ // (a) Preparing `to` reuses the transferred MemoryContext instance.
+ invokeCreateAndSetRunnerContext(mgr, to);
+ assertThat(to.getRunnerContext().getMemoryContext()).isSameAs(fromMemCtx);
// (b) Continuation context routed to `to`.
assertThat(mgr.hasContinuationContext(to)).isTrue();
- // (c) The `from`-side continuation map entry was never populated by the transfer
- // — the source carries its continuation on its runner context, not on the
- // manager's map.
- assertThat(mgr.hasContinuationContext(from)).isFalse();
+ // (c) The pending-event buffer is shared with the source's live buffer, so events
+ // emitted before the suspend survive into the generated task.
+ assertThat(to.getRunnerContext().getPendingEvents())
+ .isSameAs(from.getRunnerContext().getPendingEvents());
+
+ // (d) The removed source record fails fast on access.
+ assertThatThrownBy(() -> mgr.getContinuationContext(from))
+ .isInstanceOf(NullPointerException.class)
+ .hasMessageContaining("Missing contexts for action task");
- // (d) Persisted Action lifecycle state follows the continuation task.
+ // (e) Persisted Action lifecycle state follows the continuation task.
assertThat(to.hasExecutionStartedEventEmitted()).isTrue();
}
}
@Test
- void reportedExecutionStateFollowsActionExecutionAcrossContinuationTasks() throws Exception {
+ void componentListenersFollowActionExecutionAcrossContinuationTasks() throws Exception {
try (ActionTaskContextManager mgr = new ActionTaskContextManager(1)) {
Action action = TestActions.noopAction();
- ActionTask from = new JavaActionTask("k", new InputEvent(1L), action);
+ ActionTask from = new JavaActionTask("k", new InputEvent(1L), action, 1L);
ActionTask to =
- new JavaActionTask("k", new InputEvent(1L), action, from.getTraceContext());
- List<ExecutionTraceContext> reports = new ArrayList<>();
- ExecutionEventSink sink = (event, context) -> reports.add(context);
-
- invokeCreateAndSetRunnerContext(mgr, from, sink);
+ new JavaActionTask("k", new InputEvent(1L), action, 1L, from.getTraceContext());
+ RecordingComponentListener listener = new RecordingComponentListener();
+ AtomicInteger factoryInvocations = new AtomicInteger();
+ Function<ActionTask, List<ComponentExecutionListener>> factory =
+ task -> {
+ factoryInvocations.incrementAndGet();
+ return List.of(listener);
+ };
+
+ invokeCreateAndSetRunnerContext(mgr, from, factory);
from.getRunnerContext()
.reportExecutionStarted(
ExecutionReporter.EntityTypes.TOOL, "slow-tool", Map.of());
mgr.transferContexts(from, to, new DurableExecutionManager(null));
- invokeCreateAndSetRunnerContext(mgr, to, sink);
+ invokeCreateAndSetRunnerContext(mgr, to, factory);
to.getRunnerContext()
.reportExecutionSucceeded(
ExecutionReporter.EntityTypes.TOOL, "slow-tool", Map.of());
- assertThat(reports).hasSize(2);
- assertThat(reports.get(1).getExecutionId()).isEqualTo(reports.get(0).getExecutionId());
+ // The continuation task shares the execution's listener list, so the start/terminal
+ // pair reaches the same listener instance.
+ assertThat(factoryInvocations.get()).isOne();
+ assertThat(listener.started).containsExactly("slow-tool");
+ assertThat(listener.succeeded).containsExactly("slow-tool");
}
}
@Test
- void completingActionExecutionDropsReportedExecutionState() throws Exception {
+ void removingContextsDropsComponentListeners() throws Exception {
try (ActionTaskContextManager mgr = new ActionTaskContextManager(1)) {
- ActionTask task = new JavaActionTask("k", new InputEvent(1L), TestActions.noopAction());
- List<ExecutionTraceContext> reports = new ArrayList<>();
- ExecutionEventSink sink = (event, context) -> reports.add(context);
-
- invokeCreateAndSetRunnerContext(mgr, task, sink);
+ ActionTask task =
+ new JavaActionTask("k", new InputEvent(1L), TestActions.noopAction(), 1L);
+ List<RecordingComponentListener> created = new ArrayList<>();
+ Function<ActionTask, List<ComponentExecutionListener>> factory =
+ ignored -> {
+ RecordingComponentListener listener = new RecordingComponentListener();
+ created.add(listener);
+ return List.of(listener);
+ };
+
+ invokeCreateAndSetRunnerContext(mgr, task, factory);
task.getRunnerContext()
.reportExecutionStarted(ExecutionReporter.EntityTypes.LLM, "model-a", Map.of());
- mgr.completeActionExecution(task);
- invokeCreateAndSetRunnerContext(mgr, task, sink);
+ mgr.removeContexts(task);
+ invokeCreateAndSetRunnerContext(mgr, task, factory);
task.getRunnerContext()
.reportExecutionSucceeded(
ExecutionReporter.EntityTypes.LLM, "model-a", Map.of());
- assertThat(reports).hasSize(2);
- assertThat(reports.get(1).getExecutionId())
- .isNotEqualTo(reports.get(0).getExecutionId());
+ // A released contexts record starts a fresh listener list on the next preparation.
+ assertThat(created).hasSize(2);
+ assertThat(created.get(0).started).containsExactly("model-a");
+ assertThat(created.get(0).succeeded).isEmpty();
+ assertThat(created.get(1).succeeded).containsExactly("model-a");
}
}
@Test
void activeExecutionReportsDoNotEnterActionTaskState() throws Exception {
try (ActionTaskContextManager mgr = new ActionTaskContextManager(1)) {
- ActionTask task = new JavaActionTask("k", new InputEvent(1L), TestActions.noopAction());
- invokeCreateAndSetRunnerContext(mgr, task, (event, context) -> {});
+ ActionTask task =
+ new JavaActionTask("k", new InputEvent(1L), TestActions.noopAction(), 1L);
+ invokeCreateAndSetRunnerContext(
+ mgr, task, ignored -> List.of(new RecordingComponentListener()));
task.getRunnerContext()
.reportExecutionStarted(
ExecutionReporter.EntityTypes.TOOL,
@@ -344,8 +370,8 @@ class ActionTaskContextManagerTest {
try (ActionTaskContextManager mgr = new ActionTaskContextManager(1)) {
Action action = TestActions.noopAction();
InputEvent event = new InputEvent(1L);
- ActionTask from = new JavaActionTask("k", event, action);
- ActionTask to = new JavaActionTask("k", new InputEvent(2L), action);
+ ActionTask from = new JavaActionTask("k", event, action, 1L);
+ ActionTask to = new JavaActionTask("k", new InputEvent(2L), action, 1L);
invokeCreateAndSetRunnerContext(mgr, from);
@@ -380,7 +406,7 @@ class ActionTaskContextManagerTest {
void closeIsIdempotent() throws Exception {
// Not using try-with-resources here because we want to call close() explicitly twice.
ActionTaskContextManager mgr = new ActionTaskContextManager(1);
- ActionTask t = new JavaActionTask("k", new InputEvent(1L), TestActions.noopAction());
+ ActionTask t = new JavaActionTask("k", new InputEvent(1L), TestActions.noopAction(), 1L);
invokeCreateAndSetRunnerContext(mgr, t);
// First close() shuts down the runner context and the continuation executor
@@ -408,8 +434,10 @@ class ActionTaskContextManagerTest {
}
private static void invokeCreateAndSetRunnerContext(
- ActionTaskContextManager mgr, ActionTask task, ExecutionEventSink executionEventSink) {
- invokeCreateAndSetRunnerContext(mgr, task, null, executionEventSink);
+ ActionTaskContextManager mgr,
+ ActionTask task,
+ Function<ActionTask, List<ComponentExecutionListener>> componentListenerFactory) {
+ invokeCreateAndSetRunnerContext(mgr, task, null, componentListenerFactory);
}
@SuppressWarnings("unchecked")
@@ -417,7 +445,7 @@ class ActionTaskContextManagerTest {
ActionTaskContextManager mgr,
ActionTask task,
InteranlBaseLongTermMemory longTermMemory,
- ExecutionEventSink executionEventSink) {
+ Function<ActionTask, List<ComponentExecutionListener>> componentListenerFactory) {
AgentPlan plan = newEmptyAgentPlan();
ResourceCache cache = mock(ResourceCache.class);
FlinkAgentsMetricGroupImpl metricGroup =
@@ -436,7 +464,26 @@ class ActionTaskContextManagerTest {
shortTermMem,
/* pythonRunnerContext */ null,
longTermMemory,
- executionEventSink);
+ componentListenerFactory);
+ }
+
+ /** Records the entity names of the component reports it receives. */
+ private static final class RecordingComponentListener implements ComponentExecutionListener {
+ private final List<String> started = new ArrayList<>();
+ private final List<String> succeeded = new ArrayList<>();
+
+ @Override
+ public void onComponentExecution(
+ String entityType,
+ String entityName,
+ Map<String, Object> entityMetadata,
+ Event event) {
+ if (ExecutionLifecycleEvents.EXECUTION_STARTED_EVENT_TYPE.equals(event.getType())) {
+ started.add(entityName);
+ } else {
+ succeeded.add(entityName);
+ }
+ }
}
private static AgentPlan newEmptyAgentPlan() {
diff --git a/runtime/src/test/java/org/apache/flink/agents/runtime/operator/ActionTaskTest.java b/runtime/src/test/java/org/apache/flink/agents/runtime/operator/ActionTaskTest.java
index 90d451a9c..ea24d79ca 100644
--- a/runtime/src/test/java/org/apache/flink/agents/runtime/operator/ActionTaskTest.java
+++ b/runtime/src/test/java/org/apache/flink/agents/runtime/operator/ActionTaskTest.java
@@ -42,8 +42,8 @@ class ActionTaskTest {
Event triggeringEvent = new InputEvent(1L);
Action action = TestActions.noopAction();
- ActionTask first = new JavaActionTask("key", triggeringEvent, action);
- ActionTask second = new JavaActionTask("key", triggeringEvent, action);
+ ActionTask first = new JavaActionTask("key", triggeringEvent, action, 1L);
+ ActionTask second = new JavaActionTask("key", triggeringEvent, action, 1L);
assertThat(first.getObservationId()).isNotEqualTo(second.getObservationId());
assertThat(first).isNotEqualTo(second);
@@ -51,7 +51,8 @@ class ActionTaskTest {
@Test
void taskRestoredWithoutObservationIdGetsStableFallback() {
- ActionTask task = new JavaActionTask("key", new InputEvent(1L), TestActions.noopAction());
+ ActionTask task =
+ new JavaActionTask("key", new InputEvent(1L), TestActions.noopAction(), 1L);
task.observationId = null;
String restoredObservationId = task.getObservationId();
@@ -67,7 +68,7 @@ class ActionTaskTest {
"python-action",
new PythonFunction("test_module", "test_action"),
List.of(InputEvent.EVENT_TYPE));
- PythonActionTask task = new PythonActionTask("key", new InputEvent(1L), action);
+ PythonActionTask task = new PythonActionTask("key", new InputEvent(1L), action, 1L);
PythonRunnerContextImpl context = mock(PythonRunnerContextImpl.class);
PythonActionExecutor executor = mock(PythonActionExecutor.class);
task.setRunnerContext(context);
@@ -93,7 +94,8 @@ class ActionTaskTest {
new PythonFunction("test_module", "test_action"),
List.of(InputEvent.EVENT_TYPE));
PythonGeneratorActionTask task =
- new PythonGeneratorActionTask("key", new InputEvent(1L), action, "observation-id");
+ new PythonGeneratorActionTask(
+ "key", new InputEvent(1L), action, 1L, "observation-id");
PythonRunnerContextImpl context = mock(PythonRunnerContextImpl.class);
PythonActionExecutor executor = mock(PythonActionExecutor.class);
task.setRunnerContext(context);
@@ -115,7 +117,7 @@ class ActionTaskTest {
void resultFinalizesOutputLineage() {
Event triggeringEvent = new InputEvent(1L);
Action action = TestActions.noopAction();
- ActionTask task = new JavaActionTask("key", triggeringEvent, action);
+ ActionTask task = new JavaActionTask("key", triggeringEvent, action, 1L);
Event outputEvent = new Event("result");
ActionTask.ActionTaskResult result =
@@ -130,7 +132,7 @@ class ActionTaskTest {
void resultRejectsSelfLoopBeforeMutatingAnyOutput() {
Event triggeringEvent = new InputEvent(1L);
Action action = TestActions.noopAction();
- ActionTask task = new JavaActionTask("key", triggeringEvent, action);
+ ActionTask task = new JavaActionTask("key", triggeringEvent, action, 1L);
Event validOutput = new Event("result");
assertThatThrownBy(
diff --git a/runtime/src/test/java/org/apache/flink/agents/runtime/operator/TestActions.java b/runtime/src/test/java/org/apache/flink/agents/runtime/operator/TestActions.java
index ec421d683..0b67bb8b5 100644
--- a/runtime/src/test/java/org/apache/flink/agents/runtime/operator/TestActions.java
+++ b/runtime/src/test/java/org/apache/flink/agents/runtime/operator/TestActions.java
@@ -31,12 +31,12 @@ import java.util.List;
* do not need to redeclare the boilerplate around {@link JavaFunction#JavaFunction(Class, String,
* Class[])} signature checks.
*/
-final class TestActions {
+public final class TestActions {
private TestActions() {}
/** Returns a minimal noop Java action backed by {@link #noop(InputEvent, RunnerContext)}. */
- static Action noopAction() {
+ public static Action noopAction() {
try {
return new Action(
"noop",
diff --git a/runtime/src/test/java/org/apache/flink/agents/runtime/python/resource/PythonRuntimeResourceTest.java b/runtime/src/test/java/org/apache/flink/agents/runtime/python/resource/PythonRuntimeResourceTest.java
new file mode 100644
index 000000000..dbaf71edb
--- /dev/null
+++ b/runtime/src/test/java/org/apache/flink/agents/runtime/python/resource/PythonRuntimeResourceTest.java
@@ -0,0 +1,55 @@
+/*
+ * 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.python.resource;
+
+import org.apache.flink.agents.api.resource.ResourceType;
+import org.junit.jupiter.api.Test;
+import pemja.core.object.PyObject;
+
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.verifyNoInteractions;
+
+/** Tests for {@link PythonRuntimeResource}. */
+public class PythonRuntimeResourceTest {
+
+ private final PyObject pythonResource = mock(PyObject.class);
+ private final PythonRuntimeResource handle =
+ new PythonRuntimeResource(ResourceType.CHAT_MODEL, pythonResource);
+
+ @Test
+ public void testHandleReportsTheTypeItWasMaterializedFor() {
+ assertThat(handle.getResourceType()).isEqualTo(ResourceType.CHAT_MODEL);
+ }
+
+ @Test
+ public void testHandleExposesThePythonObjectItStandsFor() {
+ assertThat(handle.getPythonResource()).isSameAs(pythonResource);
+ }
+
+ // The Python runtime opened the resource and will close it, so a handle that opened or closed
+ // it again would break the invariants its owner already established.
+ @Test
+ public void testOpenAndCloseLeaveThePythonOwnedResourceUntouched() throws Exception {
+ handle.open();
+ handle.close();
+
+ verifyNoInteractions(pythonResource);
+ }
+}
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
new file mode 100644
index 000000000..3bc241387
--- /dev/null
+++ b/runtime/src/test/java/org/apache/flink/agents/runtime/subagent/BaseAsyncSubagentSetupTest.java
@@ -0,0 +1,521 @@
+/*
+ * 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.subagent;
+
+import org.apache.flink.agents.api.Event;
+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.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;
+import org.apache.flink.agents.plan.AgentConfiguration;
+import org.apache.flink.agents.plan.AgentPlan;
+import org.apache.flink.agents.runtime.operator.ActionExecutionOperator;
+import org.apache.flink.agents.runtime.operator.ActionExecutionOperatorFactory;
+import org.apache.flink.agents.runtime.resource.ResourceContextImpl;
+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.Test;
+
+import java.util.List;
+import java.util.Map;
+import java.util.concurrent.CancellationException;
+
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
+
+/**
+ * The async-job base in pub/sub mode, exercised through the example {@link MockAsyncSubagentSetup}:
+ * unit-style assertions of the durable POST and its crash-window reconciler and the await
+ * composition, failover replay equivalence, cancellation, and pipeline flows proving that the pub
+ * POST lands immediately and the handle subscribes to the run.
+ */
+public class BaseAsyncSubagentSetupTest {
+
+ private static final String RESOURCE_NAME = "ext-agent";
+
+ // ------------------------------------------------------------------------------------------
+ // The pub: one durable POST, issued immediately
+ // ------------------------------------------------------------------------------------------
+
+ @Test
+ void submitRequestPostsImmediatelyWithoutQuerying() throws Exception {
+ MockAsyncSubagentSetup setup = new MockAsyncSubagentSetup(2, false);
+
+ DurableCallable<Void> callable = setup.submitRequestForTest("ping", "sid-1", "call-1");
... 3305 lines suppressed ...