From 0b75c9273644cf4d22ae913547a63234f336d6ee Mon Sep 17 00:00:00 2001 From: "luogen.lg" Date: Wed, 5 Aug 2026 11:28:15 +0800 Subject: [PATCH 01/11] [runtime] Consolidate per-task execution contexts into a single transferable contexts record --- .../runtime/context/RunnerContextImpl.java | 13 +- .../operator/ActionExecutionOperator.java | 9 +- .../operator/ActionTaskContextManager.java | 135 ++++++++++++------ ...unnerContextImplExecutionReporterTest.java | 35 ++++- ...unnerContextPendingEventsContractTest.java | 86 +++++++++++ .../context/TestMemoryObservationFlush.java | 5 + .../ActionTaskContextManagerTest.java | 60 ++++---- 7 files changed, 259 insertions(+), 84 deletions(-) create mode 100644 runtime/src/test/java/org/apache/flink/agents/runtime/context/RunnerContextPendingEventsContractTest.java 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 dcfc34ed0..9e930ddbf 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 @@ -122,7 +122,8 @@ public CachedMemoryStore getSensoryMemStore() { private static final Logger LOG = LoggerFactory.getLogger(RunnerContextImpl.class); - protected final List pendingEvents = new ArrayList<>(); + protected List pendingEvents = new ArrayList<>(); + protected final FlinkAgentsMetricGroupImpl agentMetricGroup; protected final Runnable mailboxThreadChecker; protected final AgentPlan agentPlan; @@ -182,12 +183,14 @@ public void setLongTermMemory(InteranlBaseLongTermMemory ltm) { public void switchActionContext( String actionName, MemoryContext memoryContext, + List pendingEvents, String contextKey, String observationId, boolean observationSuppressed) { switchActionContext( actionName, memoryContext, + pendingEvents, contextKey, observationId, observationSuppressed, @@ -198,12 +201,14 @@ public void switchActionContext( public void switchActionContext( String actionName, MemoryContext memoryContext, + List pendingEvents, String contextKey, @Nullable ExecutionTraceContext actionTraceContext, @Nullable Map activeReportedExecutions) { switchActionContext( actionName, memoryContext, + pendingEvents, contextKey, null, false, @@ -214,6 +219,7 @@ public void switchActionContext( public void switchActionContext( String actionName, MemoryContext memoryContext, + List pendingEvents, String contextKey, @Nullable String observationId, boolean observationSuppressed, @@ -221,6 +227,7 @@ public void switchActionContext( @Nullable Map activeReportedExecutions) { this.actionName = actionName; this.memoryContext = memoryContext; + this.pendingEvents = pendingEvents; this.contextKey = contextKey; this.observationId = observationId; this.observationSuppressed = observationSuppressed; @@ -350,6 +357,10 @@ public void checkNoPendingEvents() { this.pendingEvents.isEmpty(), "There are pending events remaining in the context."); } + public List getPendingEvents() { + return this.pendingEvents; + } + public List getSensoryMemoryUpdates() { mailboxThreadChecker.run(); return List.copyOf(memoryContext.getSensoryMemoryUpdates()); 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..996ba14f6 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 @@ -449,6 +449,7 @@ private void processActionTaskForKey(Object key, String contextKey) throws Excep 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) { @@ -483,11 +484,11 @@ private void processActionTaskForKey(Object key, String contextKey) throws Excep 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); durableExecManager.maybePersistTaskResult( key, sequenceNumber, 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..37cdd8508 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,6 +17,7 @@ */ 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; @@ -36,10 +37,13 @@ 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; /** @@ -50,47 +54,84 @@ *
    *
  • The shared (Java) {@link RunnerContextImpl} that is reused across action tasks via {@link * RunnerContextImpl#switchActionContext}. - *
  • 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. + *
  • 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), and the Python awaitable reference, created, + * transferred, and removed as one unit. *
  • Active child-execution reports, keyed by Action execution id, that pair start and terminal * reports across continuation tasks without entering Flink state. *
  • The {@link ContinuationActionExecutor} thread pool used to run async Java continuations. *
* - *

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. + *

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. * - *

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}. - * - *

Design constraint: package-private; no manager-to-manager held references. Cross-cutting data - * flows via method parameters. + *

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 actionTaskMemoryContexts; - private final Map continuationContexts; - private final Map pythonAwaitableRefs; + private final Map actionTaskContexts; private final Map> activeReportedExecutionsByActionExecutionId; - private final ContinuationActionExecutor continuationActionExecutor; + + private ContinuationActionExecutor continuationActionExecutor; ActionTaskContextManager(int numAsyncThreads) { - this.actionTaskMemoryContexts = new HashMap<>(); - this.continuationContexts = new HashMap<>(); - this.pythonAwaitableRefs = new HashMap<>(); + this.actionTaskContexts = new HashMap<>(); this.activeReportedExecutionsByActionExecutionId = 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 pendingEvents = new ArrayList<>(); + } + + 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. * @@ -190,6 +231,12 @@ void createAndSetRunnerContext( PythonRunnerContextImpl pythonRunnerContext, @Nullable InteranlBaseLongTermMemory longTermMemory, @Nullable ExecutionEventSink executionEventSink) { + 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 = @@ -219,19 +266,19 @@ void createAndSetRunnerContext( } 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()), @@ -246,6 +293,7 @@ void createAndSetRunnerContext( continuationContext = this.getContinuationContext(actionTask); } else { continuationContext = new ContinuationContext(); + putContinuationContext(actionTask, continuationContext); } ((JavaRunnerContextImpl) context).setContinuationContext(continuationContext); } @@ -260,22 +308,19 @@ void createAndSetRunnerContext( 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. * *

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 +328,13 @@ RunnerContextImpl.MemoryContext removeMemoryContext(ActionTask actionTask) { */ 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(); RunnerContextImpl.DurableExecutionContext durableContext = fromTask.getRunnerContext().getDurableExecutionContext(); if (durableContext != null) { @@ -321,34 +371,27 @@ private Map getOrCreateActiveReport @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/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..9fa5037bc 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 @@ -47,7 +47,12 @@ void reportedExecutionReusesChildTraceContextBetweenStartAndFinish() throws Exce runnerContext.setExecutionEventSink( (event, context) -> reports.add(new RecordedReport(event, context))); runnerContext.switchActionContext( - "chat_model_action", null, "business-key", actionTraceContext, new HashMap<>()); + "chat_model_action", + null, + new ArrayList<>(), + "business-key", + actionTraceContext, + new HashMap<>()); runnerContext.reportExecutionStarted( ExecutionReporter.EntityTypes.LLM, "model-a", Map.of()); @@ -93,17 +98,32 @@ void reportedExecutionStateFollowsActionContextAcrossSwitches() throws Exception Map activeReportsB = new HashMap<>(); runnerContext.switchActionContext( - "chat_model_action", null, "business-key", actionA, activeReportsA); + "chat_model_action", + null, + new ArrayList<>(), + "business-key", + actionA, + activeReportsA); runnerContext.reportExecutionStarted( ExecutionReporter.EntityTypes.LLM, "model-a", Map.of()); runnerContext.switchActionContext( - "tool_call_action", null, "business-key", actionB, activeReportsB); + "tool_call_action", + null, + new ArrayList<>(), + "business-key", + actionB, + activeReportsB); runnerContext.reportExecutionStarted( ExecutionReporter.EntityTypes.TOOL, "search", Map.of("toolCallId", "call-1")); runnerContext.switchActionContext( - "chat_model_action", null, "business-key", actionA, activeReportsA); + "chat_model_action", + null, + new ArrayList<>(), + "business-key", + actionA, + activeReportsA); runnerContext.reportExecutionSucceeded( ExecutionReporter.EntityTypes.LLM, "model-a", Map.of()); @@ -131,7 +151,12 @@ void pythonReporterBridgePreservesMetadataAndPythonErrorFields() throws Exceptio 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", + actionTraceContext, + new HashMap<>()); String metadata = "{\"toolCallId\":\"call-1\",\"toolType\":\"function\"}"; runnerContext.reportExecutionStartedJson( 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..35b3aa2d1 --- /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 bufferA = new ArrayList<>(); + List bufferB = new ArrayList<>(); + Event eventA = new InputEvent(1L); + + context.switchActionContext("action-a", memoryA, bufferA, "key-a", "obs-a", false); + context.sendEvent(eventA); + assertThat(context.drainEvents(null)).containsExactly(eventA); + context.checkNoPendingEvents(); + + context.switchActionContext("action-b", memoryB, bufferB, "key-b", "obs-b", false); + 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 bufferA = new ArrayList<>(); + List bufferB = new ArrayList<>(); + Event eventA = new InputEvent(1L); + + context.switchActionContext("action-a", memoryA, bufferA, "key-a", "obs-a", false); + 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); + 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); + 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..14e0c37c1 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 javax.annotation.Nullable; +import java.util.ArrayList; import java.util.HashMap; import java.util.List; import java.util.Map; @@ -84,6 +85,7 @@ private RunnerContextImpl createContext( new RunnerContextImpl.MemoryContext( new CachedMemoryStore(new ForTestMemoryMapState<>()), new CachedMemoryStore(new ForTestMemoryMapState<>())), + new ArrayList<>(), contextKey, "observation-1", suppressed); @@ -222,6 +224,7 @@ void observationConfigurationIsNotRepeatedAcrossActionSwitches() throws Exceptio new RunnerContextImpl.MemoryContext( new CachedMemoryStore(new ForTestMemoryMapState<>()), new CachedMemoryStore(new ForTestMemoryMapState<>())), + new ArrayList<>(), "user-43", "observation-2", true); @@ -259,6 +262,7 @@ void interleavedSameKeyActionsKeepLtmEventsWithTheirOwningExecution() throws Exc new RunnerContextImpl.MemoryContext( new CachedMemoryStore(new ForTestMemoryMapState<>()), new CachedMemoryStore(new ForTestMemoryMapState<>())), + new ArrayList<>(), "user-42", "observation-2", false); @@ -274,6 +278,7 @@ void interleavedSameKeyActionsKeepLtmEventsWithTheirOwningExecution() throws Exc new RunnerContextImpl.MemoryContext( new CachedMemoryStore(new ForTestMemoryMapState<>()), new CachedMemoryStore(new ForTestMemoryMapState<>())), + new ArrayList<>(), "user-42", "observation-1", false); 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..ab584de9f 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 @@ -119,17 +119,21 @@ private static void setField(ActionTaskContextManager mgr, String name, Object v } @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); + // 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 +141,12 @@ void perTaskMapsAreIsolatedAcrossPutGetRemove() throws Exception { 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); } } @@ -170,8 +175,7 @@ void createAndSetRunnerContextBuildsFreshMemoryContextOnFirstCall() throws Excep ActionTask t = new JavaActionTask("k", new InputEvent(1L), TestActions.noopAction()); 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(); } @@ -184,9 +188,8 @@ void createAndSetRunnerContextReusesExistingMemoryContext() throws Exception { ActionTask from = new JavaActionTask("k", new InputEvent(1L), action); ActionTask to = new JavaActionTask("k", new InputEvent(2L), action); - // 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(); @@ -237,29 +240,30 @@ void transferContextsCopiesMemoryAndContinuationToNewTask() throws Exception { 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(); } } From 966ec52b254d548b81ef4544f4a31738586bc71e Mon Sep 17 00:00:00 2001 From: "luogen.lg" Date: Wed, 5 Aug 2026 15:36:15 +0800 Subject: [PATCH 02/11] [python] Add optional durable id keying for durable execution --- python/flink_agents/api/runner_context.py | 14 +++++++ .../flink_agents/runtime/durable_execution.py | 37 ++++++++++++++++++- .../runtime/flink_runner_context.py | 7 ++++ .../runtime/tests/test_durable_execution.py | 21 +++++++++++ 4 files changed, 78 insertions(+), 1 deletion(-) 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 @@ def durable_execute( 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 @@ def my_action(event, ctx): 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 @@ def durable_execute_async( 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 @@ async def my_action(event, ctx): 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/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..d230f6c0e 100644 --- a/python/flink_agents/runtime/flink_runner_context.py +++ b/python/flink_agents/runtime/flink_runner_context.py @@ -51,6 +51,7 @@ _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 @@ -1113,6 +1114,7 @@ def durable_execute( 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 +1128,8 @@ def durable_execute( 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 +1157,7 @@ def durable_execute_async( 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 +1172,8 @@ def durable_execute_async( 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( 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 @@ _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 From e59e5222a1268e4c65a4fa60524eaa3c29143559 Mon Sep 17 00:00:00 2001 From: "luogen.lg" Date: Wed, 5 Aug 2026 15:40:05 +0800 Subject: [PATCH 03/11] [runtime] Record the input sequence number on action tasks --- .../operator/ActionExecutionOperator.java | 17 ++++++--- .../agents/runtime/operator/ActionTask.java | 36 ++++++++++++++++--- .../runtime/operator/JavaActionTask.java | 12 ++++--- .../python/operator/PythonActionTask.java | 22 +++++++----- .../operator/PythonGeneratorActionTask.java | 19 ++++++---- .../ActionTaskContextManagerTest.java | 35 +++++++++--------- .../runtime/operator/ActionTaskTest.java | 16 +++++---- 7 files changed, 108 insertions(+), 49 deletions(-) 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 996ba14f6..670b687c5 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 @@ -317,7 +317,12 @@ private void processEvent( if (triggerActions != null && !triggerActions.isEmpty()) { for (Action triggerAction : triggerActions) { stateManager.addActionTask( - createActionTask(key, triggerAction, event, traceContext)); + createActionTask( + key, + triggerAction, + event, + stateManager.getSequenceNumber(), + traceContext)); } } } @@ -729,13 +734,17 @@ private void notifyExecutionLifecycleEvent(ExecutionTraceContext traceContext, E } 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()); 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..6841c8712 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 Object getKey() { 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 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/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 @@ */ 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 @@ protected PythonActionTask( 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 ActionTaskResult invoke(ClassLoader userCodeClassLoader, PythonActionExec ((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 @@ /** 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 ActionTaskResult invoke(ClassLoader userCodeClassLoader, PythonActionExec + "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/test/java/org/apache/flink/agents/runtime/operator/ActionTaskContextManagerTest.java b/runtime/src/test/java/org/apache/flink/agents/runtime/operator/ActionTaskContextManagerTest.java index ab584de9f..fcec0f854 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 @@ -122,8 +122,8 @@ private static void setField(ActionTaskContextManager mgr, String name, Object v 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); @@ -172,7 +172,8 @@ void createOrGetRunnerContextThrowsWhenPythonContextRequestedButNull() throws Ex @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 pins the freshly created MemoryContext. @@ -185,8 +186,8 @@ void createAndSetRunnerContextBuildsFreshMemoryContextOnFirstCall() throws Excep 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 carries and pins a fresh // MemoryContext. @@ -214,8 +215,8 @@ void createAndSetRunnerContextReusesExistingMemoryContext() throws Exception { 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); @@ -231,8 +232,8 @@ void sameKeyTasksSwitchLtmWithDistinctObservationIds() throws Exception { 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); @@ -272,9 +273,9 @@ void transferContextsCopiesMemoryAndContinuationToNewTask() throws Exception { void reportedExecutionStateFollowsActionExecutionAcrossContinuationTasks() 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()); + new JavaActionTask("k", new InputEvent(1L), action, 1L, from.getTraceContext()); List reports = new ArrayList<>(); ExecutionEventSink sink = (event, context) -> reports.add(context); @@ -297,7 +298,8 @@ void reportedExecutionStateFollowsActionExecutionAcrossContinuationTasks() throw @Test void completingActionExecutionDropsReportedExecutionState() throws Exception { try (ActionTaskContextManager mgr = new ActionTaskContextManager(1)) { - ActionTask task = new JavaActionTask("k", new InputEvent(1L), TestActions.noopAction()); + ActionTask task = + new JavaActionTask("k", new InputEvent(1L), TestActions.noopAction(), 1L); List reports = new ArrayList<>(); ExecutionEventSink sink = (event, context) -> reports.add(context); @@ -320,7 +322,8 @@ void completingActionExecutionDropsReportedExecutionState() throws Exception { @Test void activeExecutionReportsDoNotEnterActionTaskState() throws Exception { try (ActionTaskContextManager mgr = new ActionTaskContextManager(1)) { - ActionTask task = new JavaActionTask("k", new InputEvent(1L), TestActions.noopAction()); + ActionTask task = + new JavaActionTask("k", new InputEvent(1L), TestActions.noopAction(), 1L); invokeCreateAndSetRunnerContext(mgr, task, (event, context) -> {}); task.getRunnerContext() .reportExecutionStarted( @@ -348,8 +351,8 @@ void transferContextsRoutesDurableContextThroughManager() throws Exception { 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); @@ -384,7 +387,7 @@ void transferContextsRoutesDurableContextThroughManager() throws Exception { 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 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 @@ void distinctActionExecutionsHaveDistinctObservationIds() { 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 @@ void distinctActionExecutionsHaveDistinctObservationIds() { @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 @@ void pythonContinuationKeepsObservationId() throws Exception { "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 @@ void pythonRecoveredContinuationKeepsObservationId() throws Exception { 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 @@ void pythonRecoveredContinuationKeepsObservationId() throws Exception { 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 @@ void resultFinalizesOutputLineage() { 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( From afa13afff5a0f373c6dc5c1b2633f9194ffd2669 Mon Sep 17 00:00:00 2001 From: "luogen.lg" Date: Wed, 26 Aug 2026 13:54:46 +0800 Subject: [PATCH 04/11] [runtime] Broadcast record/action lifecycle events to registered listeners --- .../runtime/flink_runner_context.py | 113 ++++ .../runtime/task_lifecycle_listener.py | 96 ++++ .../tests/test_task_lifecycle_bridge.py | 177 ++++++ .../runtime/context/RunnerContextImpl.java | 103 +--- .../lifecycle/ComponentExecutionListener.java | 59 ++ .../PythonTaskLifecycleListener.java | 82 +++ .../lifecycle/TaskLifecycleListener.java | 130 +++++ .../operator/ActionExecutionOperator.java | 162 ++++-- .../agents/runtime/operator/ActionTask.java | 2 +- .../operator/ActionTaskContextManager.java | 58 +- .../python/utils/PythonActionExecutor.java | 80 +++ .../EventLogComponentExecutionListener.java | 86 +++ .../trace/EventLogTaskLifecycleListener.java | 64 +++ ...unnerContextImplExecutionReporterTest.java | 279 +++++---- ...unnerContextPendingEventsContractTest.java | 10 +- .../context/TestMemoryObservationFlush.java | 12 +- ...TaskLifecycleListenerNotificationTest.java | 532 ++++++++++++++++++ .../ActionTaskContextManagerTest.java | 90 ++- .../agents/runtime/operator/TestActions.java | 4 +- ...ventLogComponentExecutionListenerTest.java | 232 ++++++++ .../EventLogTaskLifecycleListenerTest.java | 150 +++++ 21 files changed, 2222 insertions(+), 299 deletions(-) create mode 100644 python/flink_agents/runtime/task_lifecycle_listener.py create mode 100644 python/flink_agents/runtime/tests/test_task_lifecycle_bridge.py create mode 100644 runtime/src/main/java/org/apache/flink/agents/runtime/lifecycle/ComponentExecutionListener.java create mode 100644 runtime/src/main/java/org/apache/flink/agents/runtime/lifecycle/PythonTaskLifecycleListener.java create mode 100644 runtime/src/main/java/org/apache/flink/agents/runtime/lifecycle/TaskLifecycleListener.java create mode 100644 runtime/src/main/java/org/apache/flink/agents/runtime/trace/EventLogComponentExecutionListener.java create mode 100644 runtime/src/main/java/org/apache/flink/agents/runtime/trace/EventLogTaskLifecycleListener.java create mode 100644 runtime/src/test/java/org/apache/flink/agents/runtime/lifecycle/TaskLifecycleListenerNotificationTest.java create mode 100644 runtime/src/test/java/org/apache/flink/agents/runtime/trace/EventLogComponentExecutionListenerTest.java create mode 100644 runtime/src/test/java/org/apache/flink/agents/runtime/trace/EventLogTaskLifecycleListenerTest.java diff --git a/python/flink_agents/runtime/flink_runner_context.py b/python/flink_agents/runtime/flink_runner_context.py index d230f6c0e..aba61f773 100644 --- a/python/flink_agents/runtime/flink_runner_context.py +++ b/python/flink_agents/runtime/flink_runner_context.py @@ -66,6 +66,7 @@ _failure_of, _first_or_logged, ) +from flink_agents.runtime.task_lifecycle_listener import TaskLifecycleListener logger = logging.getLogger(__name__) @@ -486,6 +487,10 @@ def __init__( 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. @@ -531,6 +536,55 @@ def get_resource( resource.set_metric_group(metric_group or self.action_metric_group) return resource + 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]: @@ -1302,6 +1356,65 @@ def close_flink_runner_context( ctx.close() +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/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_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/context/RunnerContextImpl.java b/runtime/src/main/java/org/apache/flink/agents/runtime/context/RunnerContextImpl.java index 9e930ddbf..b8f0da501 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.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.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; @@ -151,9 +149,8 @@ public CachedMemoryStore getSensoryMemStore() { /** 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 activeReportedExecutions; + /** Component execution listeners of the current action execution, fanned out best-effort. */ + @Nullable protected List componentExecutionListeners; /** Context for fine-grained durable execution, may be null if not enabled. */ @Nullable protected DurableExecutionContext durableExecutionContext; @@ -180,42 +177,6 @@ public void setLongTermMemory(InteranlBaseLongTermMemory ltm) { this.ltm = ltm; } - public void switchActionContext( - String actionName, - MemoryContext memoryContext, - List pendingEvents, - String contextKey, - String observationId, - boolean observationSuppressed) { - switchActionContext( - actionName, - memoryContext, - pendingEvents, - contextKey, - observationId, - observationSuppressed, - null, - null); - } - - public void switchActionContext( - String actionName, - MemoryContext memoryContext, - List pendingEvents, - String contextKey, - @Nullable ExecutionTraceContext actionTraceContext, - @Nullable Map activeReportedExecutions) { - switchActionContext( - actionName, - memoryContext, - pendingEvents, - contextKey, - null, - false, - actionTraceContext, - activeReportedExecutions); - } - public void switchActionContext( String actionName, MemoryContext memoryContext, @@ -223,8 +184,7 @@ public void switchActionContext( String contextKey, @Nullable String observationId, boolean observationSuppressed, - @Nullable ExecutionTraceContext actionTraceContext, - @Nullable Map activeReportedExecutions) { + @Nullable List componentExecutionListeners) { this.actionName = actionName; this.memoryContext = memoryContext; this.pendingEvents = pendingEvents; @@ -232,17 +192,12 @@ public void switchActionContext( 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; } @@ -361,6 +316,11 @@ public List getPendingEvents() { return this.pendingEvents; } + @Nullable + public List getComponentExecutionListeners() { + return this.componentExecutionListeners; + } + public List getSensoryMemoryUpdates() { mailboxThreadChecker.run(); return List.copyOf(memoryContext.getSensoryMemoryUpdates()); @@ -415,42 +375,27 @@ public void reportExecutionFailed( 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 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. + * + *

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. + * + *

Invariants a listener may rely on, and must not break: + * + *

    + *
  • The callback runs on the mailbox thread, so a listener needs no synchronization of its own. + *
  • An exception thrown by a listener is logged and swallowed, so reporting never fails the + * reporting component and never starves the remaining listeners. + *
  • 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. + *
  • The event instance is shared with every other listener, so a listener must treat it as + * read-only. + *
+ */ +@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 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}. + * + *

All callbacks run on the mailbox thread. + * + *

Event pairing semantics: + * + *

    + *
  • {@code onRecordStart}/{@code onRecordFinished} bracket the processing of one input record + * for a key; a record that triggers no actions emits neither callback. + *
  • {@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. + *
  • {@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}. + *
+ * + *

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 670b687c5..f27a06bd3 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,6 @@ 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.trace.ExecutionTraceContext; import org.apache.flink.agents.plan.AgentPlan; import org.apache.flink.agents.plan.JavaFunction; @@ -32,6 +30,8 @@ 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.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; @@ -40,6 +40,8 @@ import org.apache.flink.agents.runtime.metrics.FlinkAgentsMetricGroupImpl; import org.apache.flink.agents.runtime.python.operator.PythonActionTask; 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 +143,13 @@ public class ActionExecutionOperator extends AbstractStreamOperator taskLifecycleListeners = new ArrayList<>(); + + // Broadcast targets for component execution reports, injected per action execution. + private transient List componentExecutionListeners = + new ArrayList<>(); + public ActionExecutionOperator( AgentPlan agentPlan, Boolean inputIsJava, @@ -217,6 +226,15 @@ public void open() throws Exception { // runner context created by ActionTaskContextManager. ltm = pythonBridge.getLongTermMemory(); + if (taskLifecycleListeners == null) { + taskLifecycleListeners = new ArrayList<>(); + } + if (componentExecutionListeners == null) { + componentExecutionListeners = new ArrayList<>(); + } + + registerEventLogListeners(); + // init context manager for runner context creation and memory contexts contextManager = new ActionTaskContextManager( @@ -305,8 +323,14 @@ private void processEvent( 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); @@ -323,6 +347,10 @@ private void processEvent( event, stateManager.getSequenceNumber(), traceContext)); + if (freshRecordRound) { + notifyRecordStart(key); + freshRecordRound = false; + } } } } @@ -428,7 +456,8 @@ private void processActionTaskForKey(Object key, String contextKey) throws Excep stateManager.getShortTermMemState(), pythonBridge.getPythonRunnerContext(), ltm, - executionEventLogger); + this::createComponentListeners); + notifyActionPrepared(actionTask); long sequenceNumber = stateManager.getSequenceNumber(); boolean isFinished; @@ -465,8 +494,8 @@ private void processActionTaskForKey(Object key, String contextKey) throws Excep 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); @@ -494,6 +523,11 @@ private void processActionTaskForKey(Object key, String contextKey) throws Excep // action task, meaning it is not finished. contextManager.removeContexts(actionTask); durableExecManager.removeDurableContext(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, @@ -508,25 +542,15 @@ private void processActionTaskForKey(Object key, String contextKey) throws Excep 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; @@ -548,12 +572,14 @@ private void processActionTaskForKey(Object key, String contextKey) throws Excep // 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); @@ -707,30 +733,99 @@ private void notifyActionStarted(ActionTask actionTask) { 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 createComponentListeners(ActionTask actionTask) { + List listeners = new ArrayList<>(); + listeners.add( + new EventLogComponentExecutionListener( + actionTask.getTraceContext(), executionEventLogger)); + listeners.addAll(componentExecutionListeners); + return listeners; + } + + /** + * 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( @@ -791,6 +886,9 @@ private void tryResumeProcessActionTasks() throws Exception { } 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 6841c8712..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 @@ -153,7 +153,7 @@ public String getObservationId() { 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 37cdd8508..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 @@ -19,7 +19,6 @@ 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; @@ -28,13 +27,12 @@ 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; @@ -45,6 +43,7 @@ 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}. @@ -56,10 +55,8 @@ * RunnerContextImpl#switchActionContext}. *

  • 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), and the Python awaitable reference, created, - * transferred, and removed as one unit. - *
  • Active child-execution reports, keyed by Action execution id, that pair start and terminal - * reports across continuation tasks without entering Flink state. + * continuation context (for async Java actions), the Python awaitable reference, and the + * component execution listeners, created, transferred, and removed as one unit. *
  • The {@link ContinuationActionExecutor} thread pool used to run async Java continuations. * * @@ -76,14 +73,11 @@ class ActionTaskContextManager implements AutoCloseable { private RunnerContextImpl runnerContext; private final Map actionTaskContexts; - private final Map> - activeReportedExecutionsByActionExecutionId; private ContinuationActionExecutor continuationActionExecutor; ActionTaskContextManager(int numAsyncThreads) { this.actionTaskContexts = new HashMap<>(); - this.activeReportedExecutionsByActionExecutionId = new HashMap<>(); this.continuationActionExecutor = new ContinuationActionExecutor(numAsyncThreads); } @@ -96,6 +90,7 @@ private static final class ActionTaskContexts { @Nullable private ContinuationContext continuationContext; @Nullable private String pythonAwaitableRef; private List pendingEvents = new ArrayList<>(); + @Nullable private List componentListeners; } private boolean hasContexts(ActionTask actionTask) { @@ -197,9 +192,10 @@ RunnerContextImpl createOrGetRunnerContext( *
  • Selects a Java or Python runner context based on the action's {@code Exec} type. *
  • Reuses any existing {@link RunnerContextImpl.MemoryContext} for this task; otherwise * builds a fresh one backed by the supplied sensory/short-term memory states. - *
  • Wires the runtime-level execution event sink onto the runner context. + *
  • Creates or reuses the per-action-execution component listener list and wires it onto + * the runner context. *
  • 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. *
  • For Java contexts, attaches a continuation context (re-used if the task is resuming * from an async suspend, fresh otherwise). *
  • For Python contexts, attaches the per-task awaitable reference (or {@code null} if the @@ -230,7 +226,9 @@ void createAndSetRunnerContext( MapState shortTermMemState, PythonRunnerContextImpl pythonRunnerContext, @Nullable InteranlBaseLongTermMemory longTermMemory, - @Nullable ExecutionEventSink executionEventSink) { + @Nullable + Function> + 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 @@ -264,7 +262,6 @@ void createAndSetRunnerContext( throw new IllegalStateException( "Unsupported action type: " + actionTask.action.getExec().getClass()); } - context.setExecutionEventSink(executionEventSink); RunnerContextImpl.MemoryContext memoryContext = getMemoryContext(actionTask); if (memoryContext == null) { @@ -282,8 +279,7 @@ void createAndSetRunnerContext( contextKey, actionTask.getObservationId(), MemoryEvent.isMemoryType(actionTask.event.getType()), - actionTask.getTraceContext(), - getOrCreateActiveReportedExecutions(actionTask)); + getOrCreateComponentListeners(actionTask, componentListenerFactory)); if (context instanceof JavaRunnerContextImpl) { ContinuationContext continuationContext; @@ -335,6 +331,11 @@ void transferContexts( // 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) { @@ -354,19 +355,20 @@ void transferContexts( } } - void completeActionExecution(ActionTask actionTask) { - activeReportedExecutionsByActionExecutionId.remove( - actionTask.getTraceContext().getExecutionId()); - } - - private Map getOrCreateActiveReportedExecutions( - ActionTask actionTask) { - String executionId = actionTask.getTraceContext().getExecutionId(); - if (executionId == null) { - throw new IllegalStateException("Action execution id must not be null."); + @Nullable + private List getOrCreateComponentListeners( + ActionTask actionTask, + @Nullable + Function> + componentListenerFactory) { + if (componentListenerFactory == null) { + return null; + } + ActionTaskContexts contexts = requireContexts(actionTask); + if (contexts.componentListeners == null) { + contexts.componentListeners = componentListenerFactory.apply(actionTask); } - return activeReportedExecutionsByActionExecutionId.computeIfAbsent( - executionId, ignored -> new HashMap<>()); + return contexts.componentListeners; } @Nullable 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..73c82fe71 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 @@ -23,6 +23,7 @@ import org.apache.flink.agents.api.agents.AgentExecutionOptions; 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.types.Row; import org.apache.flink.util.ExceptionUtils; @@ -49,6 +50,25 @@ public class PythonActionExecutor implements AutoCloseable { private static final String CLOSE_FLINK_RUNNER_CONTEXT = "flink_runner_context.close_flink_runner_context"; + // =========== 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 +117,66 @@ public PyObject getPythonRunnerContext() { return pythonRunnerContext; } + /** + * 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/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 activeReportedExecutions = + new HashMap<>(); + + public EventLogComponentExecutionListener( + ExecutionTraceContext actionTraceContext, ExecutionEventSink executionEventSink) { + this.actionTraceContext = actionTraceContext; + this.executionEventSink = executionEventSink; + } + + @Override + public void onComponentExecution( + String entityType, String entityName, Map 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/context/RunnerContextImplExecutionReporterTest.java b/runtime/src/test/java/org/apache/flink/agents/runtime/context/RunnerContextImplExecutionReporterTest.java index 9fa5037bc..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,143 +20,133 @@ 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 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, - new ArrayList<>(), - "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 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 activeReportsA = new HashMap<>(); - Map activeReportsB = new HashMap<>(); + switchToChatModelAction(runnerContext, List.of(listener)); - runnerContext.switchActionContext( - "chat_model_action", - null, - new ArrayList<>(), - "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, - new ArrayList<>(), - "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, - new ArrayList<>(), - "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 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, new ArrayList<>(), "business-key", - actionTraceContext, - new HashMap<>()); + "obs-1", + false, + List.of(listener)); String metadata = "{\"toolCallId\":\"call-1\",\"toolType\":\"function\"}"; runnerContext.reportExecutionStartedJson( @@ -169,42 +159,85 @@ void pythonReporterBridgePreservesMetadataAndPythonErrorFields() throws Exceptio "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 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> started = new ArrayList<>(); + private final List> succeeded = new ArrayList<>(); + private final List failed = new ArrayList<>(); + + @Override + public void onComponentExecution( + String entityType, + String entityName, + Map 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 entityMetadata; + private final String errorType; + @Nullable private final String errorMessage; + @Nullable private final String problemCategory; + + private RecordedFailure( + String entityType, + String entityName, + Map 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 index 35b3aa2d1..e6f3be5f2 100644 --- 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 @@ -42,12 +42,12 @@ void emittedEventsDrainAndBufferIsClearBeforeTaskSwitch() { List bufferB = new ArrayList<>(); Event eventA = new InputEvent(1L); - context.switchActionContext("action-a", memoryA, bufferA, "key-a", "obs-a", false); + 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); + context.switchActionContext("action-b", memoryB, bufferB, "key-b", "obs-b", false, null); assertThat(context.drainEvents(null)).isEmpty(); } @@ -60,17 +60,17 @@ void bufferedEventsStayIsolatedPerTaskAcrossContextSwitches() { List bufferB = new ArrayList<>(); Event eventA = new InputEvent(1L); - context.switchActionContext("action-a", memoryA, bufferA, "key-a", "obs-a", false); + 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); + 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); + context.switchActionContext("action-a", memoryA, bufferA, "key-a", "obs-a", false, null); assertThat(context.drainEvents(null)).containsExactly(eventA); } 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 14e0c37c1..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 @@ -88,7 +88,8 @@ private RunnerContextImpl createContext( new ArrayList<>(), contextKey, "observation-1", - suppressed); + suppressed, + null); return context; } @@ -227,7 +228,8 @@ void observationConfigurationIsNotRepeatedAcrossActionSwitches() throws Exceptio new ArrayList<>(), "user-43", "observation-2", - true); + true, + null); assertThat(ltm.configureCallCount).isEqualTo(1); assertThat(ltm.switchCallCount).isEqualTo(2); @@ -265,7 +267,8 @@ void interleavedSameKeyActionsKeepLtmEventsWithTheirOwningExecution() throws Exc new ArrayList<>(), "user-42", "observation-2", - false); + false, + null); ltm.record("user-42", "observation-2", "b", "from-b"); LongTermUpdateEvent bEvent = @@ -281,7 +284,8 @@ void interleavedSameKeyActionsKeepLtmEventsWithTheirOwningExecution() throws Exc 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 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() { + @Override + public String getId() { + return "lifecycle-notification"; + } + + @Override + public Class 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 testHarness = + new KeyedOneInputStreamOperatorTestHarness<>( + new ActionExecutionOperatorFactory(plan, true), + (KeySelector) value -> value, + TypeInformation.of(Long.class))) { + testHarness.open(); + ActionExecutionOperator operator = + (ActionExecutionOperator) 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 { + + 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 { + + private final List listeners; + + ListenerInjectingOperatorFactory(AgentPlan agentPlan, TaskLifecycleListener listener) { + super(agentPlan, true); + this.listeners = new ArrayList<>(); + this.listeners.add(listener); + } + + @Override + public > T createStreamOperator( + StreamOperatorParameters parameters) { + T operator = super.createStreamOperator(parameters); + ActionExecutionOperator actionOperator = + (ActionExecutionOperator) 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 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 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 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 testHarness = + new KeyedOneInputStreamOperatorTestHarness<>( + new TestableOperatorFactory(plan, actionStateStore), + (KeySelector) value -> value, + TypeInformation.of(Long.class))) { + testHarness.open(); + ActionExecutionOperator operator = + (ActionExecutionOperator) 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 testHarness = + new KeyedOneInputStreamOperatorTestHarness<>( + new TestableOperatorFactory(plan, actionStateStore), + (KeySelector) value -> value, + TypeInformation.of(Long.class))) { + testHarness.open(); + ActionExecutionOperator operator = + (ActionExecutionOperator) 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 testHarness = + new KeyedOneInputStreamOperatorTestHarness<>( + new TestableOperatorFactory(plan, actionStateStore), + (KeySelector) value -> value, + TypeInformation.of(Long.class))) { + testHarness.open(); + ActionExecutionOperator operator = + (ActionExecutionOperator) 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 testHarness = + new KeyedOneInputStreamOperatorTestHarness<>( + new ActionExecutionOperatorFactory<>(plan, true), + (KeySelector) 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 testHarness = + new KeyedOneInputStreamOperatorTestHarness<>( + new ListenerInjectingOperatorFactory(plan, new RecordingListener()), + (KeySelector) value -> value, + TypeInformation.of(Long.class))) { + testHarness.initializeState(snapshot); + testHarness.open(); + ActionExecutionOperator operator = + (ActionExecutionOperator) 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 testHarness = + new KeyedOneInputStreamOperatorTestHarness<>( + new ActionExecutionOperatorFactory(plan, true), + (KeySelector) value -> value, + TypeInformation.of(Long.class))) { + testHarness.open(); + ActionExecutionOperator operator = + (ActionExecutionOperator) 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 testHarness = + new KeyedOneInputStreamOperatorTestHarness<>( + new ActionExecutionOperatorFactory(plan, true), + (KeySelector) value -> value, + TypeInformation.of(Long.class))) { + testHarness.open(); + ActionExecutionOperator operator = + (ActionExecutionOperator) testHarness.getOperator(); + operator.addTaskLifecycleListener(new RecordingListener()); + + testHarness.processElement(new StreamRecord<>(5L)); + operator.waitInFlightEventsFinished(); + + List> recordOutput = + (List>) 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 middle = + RecordingListener.EVENTS.subList(3, RecordingListener.EVENTS.size() - 3); + assertThat(middle.get(0)).isEqualTo("started:handleInput"); + List 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 testHarness = + new KeyedOneInputStreamOperatorTestHarness<>( + new ActionExecutionOperatorFactory(plan, true), + (KeySelector) value -> value, + TypeInformation.of(Long.class))) { + testHarness.open(); + ActionExecutionOperator operator = + (ActionExecutionOperator) 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 fcec0f854..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.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.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; @@ -270,52 +273,67 @@ void transferContextsCopiesMemoryAndContinuationToNewTask() throws Exception { } @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, 1L); ActionTask to = new JavaActionTask("k", new InputEvent(1L), action, 1L, from.getTraceContext()); - List reports = new ArrayList<>(); - ExecutionEventSink sink = (event, context) -> reports.add(context); - - invokeCreateAndSetRunnerContext(mgr, from, sink); + RecordingComponentListener listener = new RecordingComponentListener(); + AtomicInteger factoryInvocations = new AtomicInteger(); + Function> 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(), 1L); - List reports = new ArrayList<>(); - ExecutionEventSink sink = (event, context) -> reports.add(context); - - invokeCreateAndSetRunnerContext(mgr, task, sink); + List created = new ArrayList<>(); + Function> 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"); } } @@ -324,7 +342,8 @@ void activeExecutionReportsDoNotEnterActionTaskState() throws Exception { try (ActionTaskContextManager mgr = new ActionTaskContextManager(1)) { ActionTask task = new JavaActionTask("k", new InputEvent(1L), TestActions.noopAction(), 1L); - invokeCreateAndSetRunnerContext(mgr, task, (event, context) -> {}); + invokeCreateAndSetRunnerContext( + mgr, task, ignored -> List.of(new RecordingComponentListener())); task.getRunnerContext() .reportExecutionStarted( ExecutionReporter.EntityTypes.TOOL, @@ -415,8 +434,10 @@ private static void invokeCreateAndSetRunnerContext( } private static void invokeCreateAndSetRunnerContext( - ActionTaskContextManager mgr, ActionTask task, ExecutionEventSink executionEventSink) { - invokeCreateAndSetRunnerContext(mgr, task, null, executionEventSink); + ActionTaskContextManager mgr, + ActionTask task, + Function> componentListenerFactory) { + invokeCreateAndSetRunnerContext(mgr, task, null, componentListenerFactory); } @SuppressWarnings("unchecked") @@ -424,7 +445,7 @@ private static void invokeCreateAndSetRunnerContext( ActionTaskContextManager mgr, ActionTask task, InteranlBaseLongTermMemory longTermMemory, - ExecutionEventSink executionEventSink) { + Function> componentListenerFactory) { AgentPlan plan = newEmptyAgentPlan(); ResourceCache cache = mock(ResourceCache.class); FlinkAgentsMetricGroupImpl metricGroup = @@ -443,7 +464,26 @@ private static void invokeCreateAndSetRunnerContext( 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 started = new ArrayList<>(); + private final List succeeded = new ArrayList<>(); + + @Override + public void onComponentExecution( + String entityType, + String entityName, + Map 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/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 @@ * 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/trace/EventLogComponentExecutionListenerTest.java b/runtime/src/test/java/org/apache/flink/agents/runtime/trace/EventLogComponentExecutionListenerTest.java new file mode 100644 index 000000000..eea2fb074 --- /dev/null +++ b/runtime/src/test/java/org/apache/flink/agents/runtime/trace/EventLogComponentExecutionListenerTest.java @@ -0,0 +1,232 @@ +/* + * 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.EventContext; +import org.apache.flink.agents.api.logger.EventLogger; +import org.apache.flink.agents.api.logger.EventLoggerOpenParams; +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.runtime.eventlog.EventLogWriter; +import org.junit.jupiter.api.Test; + +import javax.annotation.Nullable; + +import java.util.ArrayList; +import java.util.List; +import java.util.Map; + +import static org.assertj.core.api.Assertions.assertThat; + +/** + * Tests for {@link EventLogComponentExecutionListener}, the per-action-execution adapter that keeps + * the event log's start/terminal report pairing. + */ +class EventLogComponentExecutionListenerTest { + + @Test + void startAndTerminalReportsShareOneChildExecution() { + CapturingEventLogger logger = new CapturingEventLogger(); + ExecutionTraceContext actionContext = actionTraceContext(); + EventLogComponentExecutionListener listener = + new EventLogComponentExecutionListener(actionContext, sink(logger)); + + listener.onComponentExecution( + ExecutionReporter.EntityTypes.LLM, + "model-a", + Map.of(), + ExecutionLifecycleEvents.executionStarted()); + listener.onComponentExecution( + ExecutionReporter.EntityTypes.LLM, + "model-a", + Map.of(), + ExecutionLifecycleEvents.executionFinished()); + + assertThat(logger.records).hasSize(2); + Event start = logger.records.get(0).event; + Event terminal = logger.records.get(1).event; + assertThat(start.getType()) + .isEqualTo(ExecutionLifecycleEvents.EXECUTION_STARTED_EVENT_TYPE); + assertThat(terminal.getType()) + .isEqualTo(ExecutionLifecycleEvents.EXECUTION_FINISHED_EVENT_TYPE); + assertThat(status(start)).isEqualTo(ExecutionLifecycleEvents.STATUS_STARTED); + assertThat(status(terminal)).isEqualTo(ExecutionLifecycleEvents.STATUS_SUCCESS); + + ExecutionTraceContext startContext = logger.records.get(0).traceContext; + ExecutionTraceContext terminalContext = logger.records.get(1).traceContext; + assertThat(terminalContext.getExecutionId()).isEqualTo(startContext.getExecutionId()); + assertThat(startContext.getParentExecutionId()).isEqualTo(actionContext.getExecutionId()); + assertThat(startContext.getEntityType()).isEqualTo(ExecutionReporter.EntityTypes.LLM); + assertThat(startContext.getEntityName()).isEqualTo("model-a"); + } + + @Test + void pairingSurvivesWhenReportsUseSeparateListenerAccesses() { + CapturingEventLogger logger = new CapturingEventLogger(); + EventLogComponentExecutionListener listener = + new EventLogComponentExecutionListener(actionTraceContext(), sink(logger)); + Map metadata = Map.of("toolCallId", "call-1"); + + // Mirrors a continuation: the start is reported first, the terminal arrives later + // through the same per-execution listener instance. + listener.onComponentExecution( + ExecutionReporter.EntityTypes.TOOL, + "search", + metadata, + ExecutionLifecycleEvents.executionStarted()); + listener.onComponentExecution( + ExecutionReporter.EntityTypes.TOOL, + "search", + metadata, + ExecutionLifecycleEvents.executionFailed( + "builtins.ValueError", + "bad response", + ExecutionReporter.ProblemCategories.TOOL_CALL_FAILED)); + + assertThat(logger.records).hasSize(2); + Event failed = logger.records.get(1).event; + assertThat(failed.getType()) + .isEqualTo(ExecutionLifecycleEvents.EXECUTION_FAILED_EVENT_TYPE); + assertThat(failed.getAttr("errorType")).isEqualTo("builtins.ValueError"); + assertThat(failed.getAttr("errorMessage")).isEqualTo("bad response"); + assertThat(failed.getAttr(ExecutionLifecycleEvents.PROBLEM_CATEGORY_ATTRIBUTE)) + .isEqualTo(ExecutionReporter.ProblemCategories.TOOL_CALL_FAILED); + assertThat(logger.records.get(1).traceContext.getExecutionId()) + .isEqualTo(logger.records.get(0).traceContext.getExecutionId()); + assertThat(logger.records.get(1).traceContext.getEntityMetadata()) + .containsEntry("toolCallId", "call-1"); + } + + @Test + void terminalReportWithoutStartGetsAFreshExecutionId() { + CapturingEventLogger logger = new CapturingEventLogger(); + ExecutionTraceContext actionContext = actionTraceContext(); + EventLogComponentExecutionListener listener = + new EventLogComponentExecutionListener(actionContext, sink(logger)); + + listener.onComponentExecution( + ExecutionReporter.EntityTypes.PARSER, + "json-parser", + Map.of(), + ExecutionLifecycleEvents.executionFinished()); + + assertThat(logger.records).hasSize(1); + ExecutionTraceContext context = logger.records.get(0).traceContext; + assertThat(context.getExecutionId()).isNotBlank(); + assertThat(context.getParentExecutionId()).isEqualTo(actionContext.getExecutionId()); + } + + @Test + void repeatedStartReportReplacesTheActiveReport() { + CapturingEventLogger logger = new CapturingEventLogger(); + EventLogComponentExecutionListener listener = + new EventLogComponentExecutionListener(actionTraceContext(), sink(logger)); + + listener.onComponentExecution( + ExecutionReporter.EntityTypes.LLM, + "model-a", + Map.of(), + ExecutionLifecycleEvents.executionStarted()); + listener.onComponentExecution( + ExecutionReporter.EntityTypes.LLM, + "model-a", + Map.of(), + ExecutionLifecycleEvents.executionStarted()); + listener.onComponentExecution( + ExecutionReporter.EntityTypes.LLM, + "model-a", + Map.of(), + ExecutionLifecycleEvents.executionFinished()); + + assertThat(logger.records).hasSize(3); + // The terminal pairs with the second start; the first start stays unpaired. + assertThat(logger.records.get(2).traceContext.getExecutionId()) + .isEqualTo(logger.records.get(1).traceContext.getExecutionId()) + .isNotEqualTo(logger.records.get(0).traceContext.getExecutionId()); + } + + @Test + void disabledTraceSwitchSuppressesExecutionRecords() { + CapturingEventLogger logger = new CapturingEventLogger(); + EventLogComponentExecutionListener listener = + new EventLogComponentExecutionListener( + actionTraceContext(), + ExecutionEventLogger.forEventLogWriter( + EventLogWriter.forEventLogger(logger, false))); + + listener.onComponentExecution( + ExecutionReporter.EntityTypes.LLM, + "model-a", + Map.of(), + ExecutionLifecycleEvents.executionStarted()); + + assertThat(logger.records).isEmpty(); + } + + private static ExecutionTraceContext actionTraceContext() { + return ExecutionTraceContext.forInputRun("business-key", "agent") + .childExecution("action", "chat_model_action"); + } + + private static ExecutionEventSink sink(EventLogger logger) { + return ExecutionEventLogger.forEventLogWriter(EventLogWriter.forEventLogger(logger)); + } + + private static String status(Event event) { + return (String) event.getAttr(ExecutionLifecycleEvents.STATUS_ATTRIBUTE); + } + + /** Records every appended (event, trace context) pair. */ + private static final class CapturingEventLogger implements EventLogger { + private final List records = new ArrayList<>(); + + @Override + public void open(EventLoggerOpenParams params) {} + + @Override + public void append(EventContext eventContext, Event event) { + append(eventContext, event, null); + } + + @Override + public void append( + EventContext eventContext, + Event event, + @Nullable ExecutionTraceContext traceContext) { + records.add(new Appended(event, traceContext)); + } + + @Override + public void flush() {} + + @Override + public void close() {} + } + + private static final class Appended { + private final Event event; + private final ExecutionTraceContext traceContext; + + private Appended(Event event, ExecutionTraceContext traceContext) { + this.event = event; + this.traceContext = traceContext; + } + } +} diff --git a/runtime/src/test/java/org/apache/flink/agents/runtime/trace/EventLogTaskLifecycleListenerTest.java b/runtime/src/test/java/org/apache/flink/agents/runtime/trace/EventLogTaskLifecycleListenerTest.java new file mode 100644 index 000000000..c6ead8f57 --- /dev/null +++ b/runtime/src/test/java/org/apache/flink/agents/runtime/trace/EventLogTaskLifecycleListenerTest.java @@ -0,0 +1,150 @@ +/* + * 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.EventContext; +import org.apache.flink.agents.api.InputEvent; +import org.apache.flink.agents.api.logger.EventLogger; +import org.apache.flink.agents.api.logger.EventLoggerOpenParams; +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.runtime.eventlog.EventLogWriter; +import org.apache.flink.agents.runtime.operator.ActionTask; +import org.apache.flink.agents.runtime.operator.JavaActionTask; +import org.apache.flink.agents.runtime.operator.TestActions; +import org.junit.jupiter.api.Test; + +import javax.annotation.Nullable; + +import java.util.ArrayList; +import java.util.List; + +import static org.assertj.core.api.Assertions.assertThat; + +/** + * Tests for {@link EventLogTaskLifecycleListener}, the event log's view of the action lifecycle. + * These pin the record sequence and attributes of the pre-listener direct-emission path. + */ +class EventLogTaskLifecycleListenerTest { + + @Test + void lifecycleCallbacksEmitTheOriginalExecutionLifecycleSequence() { + CapturingEventLogger logger = new CapturingEventLogger(); + EventLogTaskLifecycleListener listener = + new EventLogTaskLifecycleListener( + ExecutionEventLogger.forEventLogWriter( + EventLogWriter.forEventLogger(logger))); + ActionTask task = new JavaActionTask("k", new InputEvent(1L), TestActions.noopAction(), 1L); + + listener.onActionStarted(task); + listener.onActionFinished(task); + + assertThat(logger.records).hasSize(2); + assertThat(logger.records.get(0).event.getType()) + .isEqualTo(ExecutionLifecycleEvents.EXECUTION_STARTED_EVENT_TYPE); + assertThat(logger.records.get(1).event.getType()) + .isEqualTo(ExecutionLifecycleEvents.EXECUTION_FINISHED_EVENT_TYPE); + assertThat(logger.records.get(0).traceContext).isEqualTo(task.getTraceContext()); + assertThat(logger.records.get(1).traceContext).isEqualTo(task.getTraceContext()); + } + + @Test + void reuseEmitsExecutionReusedOnTheTaskTraceContext() { + CapturingEventLogger logger = new CapturingEventLogger(); + EventLogTaskLifecycleListener listener = + new EventLogTaskLifecycleListener( + ExecutionEventLogger.forEventLogWriter( + EventLogWriter.forEventLogger(logger))); + ActionTask task = new JavaActionTask("k", new InputEvent(1L), TestActions.noopAction(), 1L); + + listener.onActionReused(task); + + assertThat(logger.records).hasSize(1); + assertThat(logger.records.get(0).event.getType()) + .isEqualTo(ExecutionLifecycleEvents.EXECUTION_REUSED_EVENT_TYPE); + assertThat(status(logger.records.get(0).event)) + .isEqualTo(ExecutionLifecycleEvents.STATUS_REUSED); + assertThat(logger.records.get(0).traceContext).isEqualTo(task.getTraceContext()); + } + + @Test + void failureCarriesRootCauseDetailsAndTheActionExecutionCategory() { + CapturingEventLogger logger = new CapturingEventLogger(); + EventLogTaskLifecycleListener listener = + new EventLogTaskLifecycleListener( + ExecutionEventLogger.forEventLogWriter( + EventLogWriter.forEventLogger(logger))); + ActionTask task = new JavaActionTask("k", new InputEvent(1L), TestActions.noopAction(), 1L); + + listener.onActionFailed( + task, new RuntimeException(new IllegalStateException("action boom"))); + + assertThat(logger.records).hasSize(1); + Event failed = logger.records.get(0).event; + assertThat(failed.getType()) + .isEqualTo(ExecutionLifecycleEvents.EXECUTION_FAILED_EVENT_TYPE); + assertThat(failed.getAttr("errorType")).isEqualTo(IllegalStateException.class.getName()); + assertThat(failed.getAttr("errorMessage")).isEqualTo("action boom"); + assertThat(failed.getAttr(ExecutionLifecycleEvents.PROBLEM_CATEGORY_ATTRIBUTE)) + .isEqualTo(ExecutionReporter.ProblemCategories.ACTION_EXECUTION_FAILED); + assertThat(logger.records.get(0).traceContext).isEqualTo(task.getTraceContext()); + } + + private static String status(Event event) { + return (String) event.getAttr(ExecutionLifecycleEvents.STATUS_ATTRIBUTE); + } + + /** Records every appended (event, trace context) pair. */ + private static final class CapturingEventLogger implements EventLogger { + private final List records = new ArrayList<>(); + + @Override + public void open(EventLoggerOpenParams params) {} + + @Override + public void append(EventContext eventContext, Event event) { + append(eventContext, event, null); + } + + @Override + public void append( + EventContext eventContext, + Event event, + @Nullable ExecutionTraceContext traceContext) { + records.add(new Appended(event, traceContext)); + } + + @Override + public void flush() {} + + @Override + public void close() {} + } + + private static final class Appended { + private final Event event; + private final ExecutionTraceContext traceContext; + + private Appended(Event event, ExecutionTraceContext traceContext) { + this.event = event; + this.traceContext = traceContext; + } + } +} From 6cb1aa09ff3bbfcae739e333feaf0c5a80d1624c Mon Sep 17 00:00:00 2001 From: "luogen.lg" Date: Thu, 20 Aug 2026 15:42:37 +0800 Subject: [PATCH 05/11] [runtime] Materialize resources of a type through their owning runtime --- .../resourceprovider/ResourceProvider.java | 9 ++ python/flink_agents/plan/resource_provider.py | 11 ++ .../runtime/flink_runner_context.py | 27 ++++ .../runtime/tests/test_eager_materialize.py | 129 ++++++++++++++++++ .../flink/agents/runtime/ResourceCache.java | 61 +++++++++ .../runtime/operator/PythonBridgeManager.java | 30 ++-- .../resource/PythonRuntimeResource.java | 57 ++++++++ .../python/utils/PythonActionExecutor.java | 30 ++++ .../agents/runtime/ResourceCacheTest.java | 54 ++++++++ .../resource/PythonRuntimeResourceTest.java | 55 ++++++++ 10 files changed, 448 insertions(+), 15 deletions(-) create mode 100644 python/flink_agents/runtime/tests/test_eager_materialize.py create mode 100644 runtime/src/main/java/org/apache/flink/agents/runtime/python/resource/PythonRuntimeResource.java create mode 100644 runtime/src/test/java/org/apache/flink/agents/runtime/python/resource/PythonRuntimeResourceTest.java 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 ResourceType getType() { 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/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 @@ def provide( "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/runtime/flink_runner_context.py b/python/flink_agents/runtime/flink_runner_context.py index aba61f773..4e9fefa88 100644 --- a/python/flink_agents/runtime/flink_runner_context.py +++ b/python/flink_agents/runtime/flink_runner_context.py @@ -536,6 +536,26 @@ def get_resource( 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) @@ -1356,6 +1376,13 @@ 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. 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/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..3d5661655 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,18 @@ 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.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 +52,7 @@ public class ResourceCache implements AutoCloseable { private final Map> resourceProviders; private final Map> cache = new ConcurrentHashMap<>(); private volatile PythonResourceAdapter pythonResourceAdapter; + private volatile PythonActionExecutor pythonActionExecutor; private final ResourceContextImpl resourceContext; /** @@ -86,6 +92,15 @@ void setPythonResourceAdapter(PythonResourceAdapter adapter) { 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; } @@ -157,6 +172,52 @@ public void put(String name, ResourceType type, Resource resource) { 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 eagerMaterialize(ResourceType type) throws Exception { + Map providers = resourceProviders.get(type); + List materialized = new ArrayList<>(); + if (providers == null) { + return materialized; + } + boolean hasPythonOwned = false; + for (Map.Entry 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 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/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.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 @@ *
      *
    • The {@link PythonEnvironmentManager} that prepares dependencies and the Pemja runtime. *
    • The {@link PythonInterpreter} obtained from that environment. - *
    • The {@link PythonActionExecutor} (when the plan contains Python actions or Mem0). + *
    • The {@link PythonActionExecutor} (when the plan contains Python actions, Python-owned + * resources, or Mem0). *
    • The {@link PythonRunnerContextImpl} consumed by Python actions. *
    • The Java/Python resource adapters that bridge resource lookups across languages. *
    • 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. * - *

      Scans the agent plan for any {@link PythonFunction} action or {@link - * PythonResourceProvider}. If neither is present, this method is a no-op and {@link + *

      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 @@ void open( .anyMatch( resourceProviderMap -> resourceProviderMap.values().stream() - .anyMatch( - resourceProvider -> - resourceProvider - instanceof - PythonResourceProvider)); + .anyMatch(ResourceProvider::isPythonOwned)); boolean mem0Configured = isMem0Configured(agentPlan); @@ -189,8 +188,9 @@ void open( 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 @@ private void initPythonResourceAdapter(AgentPlan agentPlan, ResourceCache resour } /** - * @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/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. + * + *

      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. + * + *

      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 73c82fe71..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,16 +21,22 @@ 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; @@ -50,6 +56,9 @@ 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"; @@ -117,6 +126,27 @@ public PyObject getPythonRunnerContext() { return pythonRunnerContext; } + /** + * Materializes every resource of the given type that the Python runtime owns and returns one + * handle per resource, keyed by resource name. + * + *

      See {@link PythonRuntimeResource} for what the returned handle may and may not do. + */ + @SuppressWarnings("unchecked") + public Map eagerMaterialize(ResourceType type) { + Object pythonResources = + interpreter.invoke(EAGER_MATERIALIZE, pythonRunnerContext, type.getValue()); + if (pythonResources == null) { + return Collections.emptyMap(); + } + Map handles = new HashMap<>(); + ((Map) 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 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..5db942909 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 @@ -40,6 +40,7 @@ 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.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; @@ -49,6 +50,7 @@ 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 +61,7 @@ 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 +223,57 @@ public Object invokePythonTool(String module, String qualName, Map 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 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(); 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); + } +} From 45341b98659bc43b1e97a0384d603c04a2ca8132 Mon Sep 17 00:00:00 2001 From: "luogen.lg" Date: Thu, 6 Aug 2026 01:31:52 +0800 Subject: [PATCH 06/11] [api] Add sub-agent resource definitions --- api/pom.xml | 17 ++ .../agents/api/resource/ResourceType.java | 3 +- .../agents/api/subagent/SubagentFuture.java | 57 ++++++ .../agents/api/subagent/SubagentFutures.java | 44 +++++ .../agents/api/subagent/SubagentResult.java | 119 +++++++++++ .../agents/api/subagent/SubagentSetup.java | 58 ++++++ .../flink/agents/api/yaml/YamlLoader.java | 2 + .../flink/agents/api/yaml/spec/AgentSpec.java | 9 +- .../api/yaml/spec/YamlAgentsDocument.java | 9 +- .../api/subagent/SubagentRegisterTest.java | 68 +++++++ .../api/subagent/SubagentResultTest.java | 75 +++++++ .../api/subagent/TestSubagentSetup.java | 91 +++++++++ .../assets/yaml-contracts.yaml | 2 +- .../flink-agents-dev/assets/yaml-schema.json | 14 ++ docs/yaml-schema.json | 14 ++ python/flink_agents/api/resource.py | 3 +- python/flink_agents/api/subagent.py | 186 ++++++++++++++++++ .../api/tests/subagent_test_utils.py | 49 +++++ .../flink_agents/api/tests/test_subagent.py | 60 ++++++ python/flink_agents/api/yaml/loader.py | 1 + python/flink_agents/api/yaml/specs.py | 2 + 21 files changed, 878 insertions(+), 5 deletions(-) create mode 100644 api/src/main/java/org/apache/flink/agents/api/subagent/SubagentFuture.java create mode 100644 api/src/main/java/org/apache/flink/agents/api/subagent/SubagentFutures.java create mode 100644 api/src/main/java/org/apache/flink/agents/api/subagent/SubagentResult.java create mode 100644 api/src/main/java/org/apache/flink/agents/api/subagent/SubagentSetup.java create mode 100644 api/src/test/java/org/apache/flink/agents/api/subagent/SubagentRegisterTest.java create mode 100644 api/src/test/java/org/apache/flink/agents/api/subagent/SubagentResultTest.java create mode 100644 api/src/test/java/org/apache/flink/agents/api/subagent/TestSubagentSetup.java create mode 100644 python/flink_agents/api/subagent.py create mode 100644 python/flink_agents/api/tests/subagent_test_utils.py create mode 100644 python/flink_agents/api/tests/test_subagent.py 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. + + + + org.apache.maven.plugins + maven-jar-plugin + 3.4.2 + + + + test-jar + + + + + + + \ 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 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. + * + *

      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}. + * + *

      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 getResult(Class 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 static LoadedFile buildAgents(Path path) { 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 @@ private static Agent buildAgent(AgentSpec spec) { 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 embeddingModelSetups; private final List vectorStores; private final List mcpServers; + private final List subagents; @JsonCreator public AgentSpec( @@ -55,7 +56,8 @@ public AgentSpec( List embeddingModelConnections, @JsonProperty("embedding_model_setups") List embeddingModelSetups, @JsonProperty("vector_stores") List vectorStores, - @JsonProperty("mcp_servers") List mcpServers) { + @JsonProperty("mcp_servers") List mcpServers, + @JsonProperty("subagents") List subagents) { this.name = name; this.description = description; this.prompts = orEmpty(prompts); @@ -68,6 +70,7 @@ public AgentSpec( this.embeddingModelSetups = orEmpty(embeddingModelSetups); this.vectorStores = orEmpty(vectorStores); this.mcpServers = orEmpty(mcpServers); + this.subagents = orEmpty(subagents); } private static List orEmpty(List list) { @@ -121,4 +124,8 @@ public List getVectorStores() { public List getMcpServers() { return mcpServers; } + + public List 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 embeddingModelSetups; private final List vectorStores; private final List mcpServers; + private final List subagents; @JsonCreator public YamlAgentsDocument( @@ -53,7 +54,8 @@ public YamlAgentsDocument( List embeddingModelConnections, @JsonProperty("embedding_model_setups") List embeddingModelSetups, @JsonProperty("vector_stores") List vectorStores, - @JsonProperty("mcp_servers") List mcpServers) { + @JsonProperty("mcp_servers") List mcpServers, + @JsonProperty("subagents") List subagents) { this.agents = orEmpty(agents); this.prompts = orEmpty(prompts); this.tools = orEmpty(tools); @@ -65,6 +67,7 @@ public YamlAgentsDocument( this.embeddingModelSetups = orEmpty(embeddingModelSetups); this.vectorStores = orEmpty(vectorStores); this.mcpServers = orEmpty(mcpServers); + this.subagents = orEmpty(subagents); } private static List orEmpty(List list) { @@ -114,4 +117,8 @@ public List getVectorStores() { public List getMcpServers() { return mcpServers; } + + public List 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 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 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/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/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 @@ "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: From 0ec451bc8dad1f4c3581de6085a131fb1e23358c Mon Sep 17 00:00:00 2001 From: "luogen.lg" Date: Thu, 6 Aug 2026 01:32:05 +0800 Subject: [PATCH 07/11] [plan] Register sub-agent setups as AGENT resources in the planner --- plan/pom.xml | 7 ++ .../apache/flink/agents/plan/AgentPlan.java | 25 ++++++ .../plan/AgentPlanSubagentResourceTest.java | 87 +++++++++++++++++++ python/flink_agents/plan/agent_plan.py | 21 +++++ .../test_agent_plan_subagent_resources.py | 72 +++++++++++++++ 5 files changed, 212 insertions(+) create mode 100644 plan/src/test/java/org/apache/flink/agents/plan/AgentPlanSubagentResourceTest.java create mode 100644 python/flink_agents/plan/tests/test_agent_plan_subagent_resources.py 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. flink-agents-api ${project.version} + + org.apache.flink + flink-agents-api + ${project.version} + test-jar + test + org.apache.flink flink-agents-integrations-mcp 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.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 @@ private void extractResourceProvidersFromAgent(Agent agent) throws Exception { + " method on your Agent class so its tools and prompts can be" + " discovered."); } + } else if (type == ResourceType.AGENT) { + for (Map.Entry 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 kv : entry.getValue().entrySet()) { ResourceDescriptor descriptor = 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 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 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/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 @@ 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/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()) From 64a022c593ec518e947a92c91fd6080b8ded265a Mon Sep 17 00:00:00 2001 From: "luogen.lg" Date: Wed, 19 Aug 2026 03:48:10 +0800 Subject: [PATCH 08/11] [runtime][python] Add sub-agent setup base with deterministic id assignment --- python/flink_agents/runtime/base_subagent.py | 294 +++++++++++++++++ python/flink_agents/runtime/resource_cache.py | 7 + .../flink_agents/runtime/subagent_handles.py | 125 +++++++ .../runtime/tests/test_base_subagent.py | 305 ++++++++++++++++++ .../flink/agents/runtime/ResourceCache.java | 7 + .../operator/ActionExecutionOperator.java | 33 ++ .../runtime/subagent/BaseSubagentSetup.java | 153 +++++++++ .../subagent/CompletedSubagentFuture.java | 49 +++ .../subagent/PendingSubagentCallRegistry.java | 67 ++++ .../runtime/subagent/SubagentFutureGroup.java | 81 +++++ .../runtime/subagent/SubagentIdAllocator.java | 143 ++++++++ .../agents/runtime/ResourceCacheTest.java | 32 ++ .../subagent/BaseSubagentSetupTest.java | 292 +++++++++++++++++ .../subagent/SubagentIdAllocatorTest.java | 251 ++++++++++++++ 14 files changed, 1839 insertions(+) create mode 100644 python/flink_agents/runtime/base_subagent.py create mode 100644 python/flink_agents/runtime/subagent_handles.py create mode 100644 python/flink_agents/runtime/tests/test_base_subagent.py create mode 100644 runtime/src/main/java/org/apache/flink/agents/runtime/subagent/BaseSubagentSetup.java create mode 100644 runtime/src/main/java/org/apache/flink/agents/runtime/subagent/CompletedSubagentFuture.java create mode 100644 runtime/src/main/java/org/apache/flink/agents/runtime/subagent/PendingSubagentCallRegistry.java create mode 100644 runtime/src/main/java/org/apache/flink/agents/runtime/subagent/SubagentFutureGroup.java create mode 100644 runtime/src/main/java/org/apache/flink/agents/runtime/subagent/SubagentIdAllocator.java create mode 100644 runtime/src/test/java/org/apache/flink/agents/runtime/subagent/BaseSubagentSetupTest.java create mode 100644 runtime/src/test/java/org/apache/flink/agents/runtime/subagent/SubagentIdAllocatorTest.java 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/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 @@ def get_resource(self, name: str, type: ResourceType) -> Resource: ) 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..bf767cc4c --- /dev/null +++ b/python/flink_agents/runtime/subagent_handles.py @@ -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. +################################################################################ +"""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.""" + outcomes = [] + for future in self._futures: + outcome = yield from future.__await__() + outcomes.append(outcome) + return outcomes 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/runtime/src/main/java/org/apache/flink/agents/runtime/ResourceCache.java b/runtime/src/main/java/org/apache/flink/agents/runtime/ResourceCache.java index 3d5661655..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 @@ -26,6 +26,7 @@ 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; @@ -152,6 +153,12 @@ public synchronized Resource getResource(String name, ResourceType type) throws 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); } 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 f27a06bd3..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,6 +21,8 @@ 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.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; @@ -31,6 +33,7 @@ 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; @@ -39,6 +42,7 @@ 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; @@ -234,6 +238,7 @@ public void open() throws Exception { } registerEventLogListeners(); + registerSubagentSetups(); // init context manager for runner context creation and memory contexts contextManager = @@ -756,6 +761,34 @@ private List createComponentListeners(ActionTask act 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. + * + *

      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 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 perTaskAllocators = new HashMap<>(); + private final Map 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/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 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..303e44cd2 --- /dev/null +++ b/runtime/src/main/java/org/apache/flink/agents/runtime/subagent/SubagentFutureGroup.java @@ -0,0 +1,81 @@ +/* + * 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 futures; + + SubagentFutureGroup(SubagentFuture first, SubagentFuture[] others) { + this(withFirst(first, others)); + } + + private static List withFirst(SubagentFuture first, SubagentFuture[] others) { + List all = new ArrayList<>(1 + others.length); + all.add(first); + all.addAll(Arrays.asList(others)); + return all; + } + + private SubagentFutureGroup(List futures) { + this.futures = futures; + } + + @Override + public boolean isDone() { + for (SubagentFuture future : futures) { + if (!future.isDone()) { + return false; + } + } + return true; + } + + @Override + public List awaitAll() throws Exception { + List 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 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 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. + * + *

      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 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/test/java/org/apache/flink/agents/runtime/ResourceCacheTest.java b/runtime/src/test/java/org/apache/flink/agents/runtime/ResourceCacheTest.java index 5db942909..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,16 +36,20 @@ 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; @@ -521,4 +525,32 @@ public void close() throws Exception { } } } + + /** 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> providers = new HashMap<>(); + Map agentProviders = new HashMap<>(); + agentProviders.put( + "reviewer", + JavaSerializableResourceProvider.createResourceProvider( + "reviewer", ResourceType.AGENT, new TestAgentSetup())); + providers.put(ResourceType.AGENT, agentProviders); + + ResourceCache cache = new ResourceCache(providers); + List 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/subagent/BaseSubagentSetupTest.java b/runtime/src/test/java/org/apache/flink/agents/runtime/subagent/BaseSubagentSetupTest.java new file mode 100644 index 000000000..8cca1acf5 --- /dev/null +++ b/runtime/src/test/java/org/apache/flink/agents/runtime/subagent/BaseSubagentSetupTest.java @@ -0,0 +1,292 @@ +/* + * 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.RunnerContext; +import org.apache.flink.agents.api.resource.ResourceType; +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.plan.AgentConfiguration; +import org.apache.flink.agents.plan.AgentPlan; +import org.apache.flink.agents.plan.JavaFunction; +import org.apache.flink.agents.plan.actions.Action; +import org.apache.flink.agents.runtime.ResourceCache; +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.agents.runtime.operator.JavaActionTask; +import org.apache.flink.api.java.functions.KeySelector; +import org.apache.flink.api.java.typeutils.TypeExtractor; +import org.apache.flink.streaming.runtime.streamrecord.StreamRecord; +import org.apache.flink.streaming.util.KeyedOneInputStreamOperatorTestHarness; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; + +import java.util.ArrayList; +import java.util.Collections; +import java.util.List; +import java.util.UUID; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; + +/** + * The short {@code submit} forms on {@link BaseSubagentSetup}: the setup observes the task + * lifecycle and assigns the missing ids deterministically from the executing task's caller-side + * facts, then delegates to the full {@code submit} provided by the subclass. + */ +public class BaseSubagentSetupTest { + + private static final String RESOURCE_NAME = "allocating"; + + @BeforeEach + public void resetCaptures() { + AllocatingCaptureSetup.reset(); + } + + /** Captures every assigned {@code (sessionId, callId)} pair like a collecting setup. */ + public static class AllocatingCaptureSetup extends BaseSubagentSetup { + + private static final List CAPTURES = + Collections.synchronizedList(new ArrayList<>()); + + /** Clears all captures. Call before each independent scenario. */ + public static void reset() { + CAPTURES.clear(); + } + + /** Snapshot of every assignment since the last {@link #reset()}, in creation order. */ + public static List captures() { + synchronized (CAPTURES) { + return new ArrayList<>(CAPTURES); + } + } + + @Override + public SubagentFuture submit( + RunnerContext ctx, Object prompt, String sessionId, String callId) { + CAPTURES.add(new String[] {sessionId, callId, String.valueOf(prompt)}); + return new SubagentFuture(sessionId, callId) { + @Override + public boolean isDone() { + return true; + } + + @Override + public SubagentResult await() { + return SubagentResult.ok(sessionId + "|" + callId + "|" + prompt); + } + + @Override + public SubagentFutures combine(SubagentFuture... others) { + throw new UnsupportedOperationException("batching is not under test"); + } + }; + } + } + + @SuppressWarnings("unused") + public static void shortForms(Event event, RunnerContext ctx) throws Exception { + BaseSubagentSetup setup = + (BaseSubagentSetup) ctx.getResource(RESOURCE_NAME, ResourceType.AGENT); + SubagentFuture noIds = setup.submit(ctx, "a"); + SubagentFuture sessionOnly = setup.submit(ctx, "b", "given-session"); + ctx.sendEvent( + new OutputEvent(noIds.await().getResult() + "|" + sessionOnly.await().getResult())); + } + + @SuppressWarnings("unused") + public static void twoStreams(Event event, RunnerContext ctx) throws Exception { + BaseSubagentSetup first = + (BaseSubagentSetup) ctx.getResource("stream-a", ResourceType.AGENT); + BaseSubagentSetup second = + (BaseSubagentSetup) ctx.getResource("stream-b", ResourceType.AGENT); + first.submit(ctx, "a"); + second.submit(ctx, "b"); + ctx.sendEvent(new OutputEvent("done")); + } + + @Test + void shortFormsAssignIdsFromTheExecutingTask() throws Exception { + try (KeyedOneInputStreamOperatorTestHarness harness = harness(plan())) { + harness.open(); + run(harness, 1L); + run(harness, 2L); + + List captures = AllocatingCaptureSetup.captures(); + assertThat(captures).hasSize(4); + + // The no-id form assigns a fresh session and the first call within it. + String[] first = captures.get(0); + assertThat(first[2]).isEqualTo("a"); + assertThat(first[0]).endsWith("-0"); + assertThat(first[1]).isEqualTo(first[0] + "-1"); + + // The session-only form assigns the call id under the given session. + String[] second = captures.get(1); + assertThat(second[0]).isEqualTo("given-session"); + assertThat(second[1]).isEqualTo("given-session-1"); + + // Another key runs under another namespace, so the assigned session differs. + String[] third = captures.get(2); + assertThat(third[0]).endsWith("-0"); + assertThat(third[0]).isNotEqualTo(first[0]); + + assertThat(harness.getRecordOutput()).hasSize(2); + } + } + + @Test + void setupsSharingOneActionAllocateDisjointIds() throws Exception { + Agent agent = new Agent(); + agent.addResource("stream-a", ResourceType.AGENT, new AllocatingCaptureSetup()); + agent.addResource("stream-b", ResourceType.AGENT, new AllocatingCaptureSetup()); + agent.addAction( + new String[] {InputEvent.EVENT_TYPE}, + BaseSubagentSetupTest.class.getMethod( + "twoStreams", Event.class, RunnerContext.class)); + try (KeyedOneInputStreamOperatorTestHarness harness = + harness(new AgentPlan(agent, new AgentConfiguration()))) { + harness.open(); + run(harness, 1L); + + // Both sub-agents count in the same action execution but carry distinct injected + // subagent names, so their first sessions must not collide. + List captures = AllocatingCaptureSetup.captures(); + assertThat(captures).hasSize(2); + assertThat(captures.get(0)[0]).isNotEqualTo(captures.get(1)[0]); + } + } + + @Test + void transferMovesBookkeepingOntoADifferentContinuation() throws Exception { + // The generated task of a suspended execution may carry different caller-side facts, + // so the bookkeeping must move onto it rather than assume the two tasks equal. + RegistryExposingSetup setup = new RegistryExposingSetup(); + setup.setSubagentName(RESOURCE_NAME); + ActionTask from = task("act", 1L); + ActionTask to = task("act-next", 2L); + + setup.onActionPrepared(from); + setup.submit(null, "a"); + setup.exposedRegistry().trackPendingSubagentCall("sid#call-1"); + setup.onActionTransferred(from, to); + + setup.onActionPrepared(to); + setup.submit(null, "b"); + + // The allocator moved with the execution: the session ordinal continues instead of + // restarting under the continuation's own facts. + List captures = AllocatingCaptureSetup.captures(); + assertThat(captures.get(0)[0]).endsWith("-0"); + assertThat(captures.get(1)[0]).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. + assertThatThrownBy(() -> setup.onActionFinishing(to)) + .isInstanceOf(IllegalStateException.class) + .hasMessageContaining("act-next") + .hasMessageContaining("sid#call-1"); + } + + @Test + void reuseSharesTheFinalizationHookWithFinishing() throws Exception { + // The replay-reuse path must close the same bookkeeping the finishing path closes, + // otherwise a prepared task replayed as completed would leak its allocator/registry. + RegistryExposingSetup setup = new RegistryExposingSetup(); + setup.setSubagentName(RESOURCE_NAME); + ActionTask task = task("act", 1L); + + setup.onActionPrepared(task); + setup.submit(null, "a"); + setup.exposedRegistry().trackPendingSubagentCall("sid#call-1"); + + assertThatThrownBy(() -> setup.onActionReused(task)) + .isInstanceOf(IllegalStateException.class) + .hasMessageContaining("act") + .hasMessageContaining("sid#call-1"); + } + + /** Capture setup exposing the current task's registry for direct tracking. */ + public static class RegistryExposingSetup extends AllocatingCaptureSetup { + + PendingSubagentCallRegistry exposedRegistry() { + return currentTaskRegistry(); + } + } + + private static ActionTask task(String actionName, long sequenceNumber) throws Exception { + Action action = + new Action( + actionName, + new JavaFunction( + BaseSubagentSetupTest.class.getName(), + "shortForms", + new Class[] {Event.class, RunnerContext.class}), + Collections.singletonList(InputEvent.EVENT_TYPE)); + return new JavaActionTask( + "k", + new Event(UUID.randomUUID(), "TestEvent", Collections.emptyMap()), + action, + sequenceNumber); + } + + @Test + void materializationInjectsSubagentNames() throws Exception { + Agent rootAgent = new Agent(); + rootAgent.addResource("root-setup", ResourceType.AGENT, new AllocatingCaptureSetup()); + ResourceCache rootCache = + new ResourceCache(new AgentPlan(rootAgent).getResourceProviders()); + BaseSubagentSetup rootSetup = + (BaseSubagentSetup) rootCache.getResource("root-setup", ResourceType.AGENT); + assertThat(rootSetup.getSubagentName()).isEqualTo("root-setup"); + } + + @SuppressWarnings("unchecked") + private static void run( + KeyedOneInputStreamOperatorTestHarness harness, long value) + throws Exception { + harness.processElement(new StreamRecord<>(value)); + ((ActionExecutionOperator) harness.getOperator()) + .waitInFlightEventsFinished(); + } + + private static AgentPlan plan() throws Exception { + Agent agent = new Agent(); + agent.addResource(RESOURCE_NAME, ResourceType.AGENT, new AllocatingCaptureSetup()); + agent.addAction( + new String[] {InputEvent.EVENT_TYPE}, + BaseSubagentSetupTest.class.getMethod( + "shortForms", Event.class, RunnerContext.class)); + return new AgentPlan(agent, new AgentConfiguration()); + } + + private static KeyedOneInputStreamOperatorTestHarness harness( + AgentPlan plan) throws Exception { + return new KeyedOneInputStreamOperatorTestHarness<>( + new ActionExecutionOperatorFactory<>(plan, true), + (KeySelector) value -> value, + TypeExtractor.getForClass(Long.class)); + } +} diff --git a/runtime/src/test/java/org/apache/flink/agents/runtime/subagent/SubagentIdAllocatorTest.java b/runtime/src/test/java/org/apache/flink/agents/runtime/subagent/SubagentIdAllocatorTest.java new file mode 100644 index 000000000..86d6c5dc9 --- /dev/null +++ b/runtime/src/test/java/org/apache/flink/agents/runtime/subagent/SubagentIdAllocatorTest.java @@ -0,0 +1,251 @@ +/* + * 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.junit.jupiter.api.Test; + +import java.util.HashSet; +import java.util.LinkedHashMap; +import java.util.Map; +import java.util.Set; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertNotEquals; +import static org.junit.jupiter.api.Assertions.assertTrue; + +/** + * Unit tests for {@link SubagentIdAllocator}: same namespace inputs plus the same call sequence + * must reproduce byte-for-byte identical ids, and any difference in the caller-side namespace + * inputs must separate the id stream. + */ +class SubagentIdAllocatorTest { + + private static SubagentIdAllocator newAllocator( + Object key, long sequenceNumber, String actionName, Object eventInput) { + return new SubagentIdAllocator( + key, sequenceNumber, actionName, new InputEvent(eventInput), "agent"); + } + + @Test + void eventInstanceIdDoesNotAffectNamespace() { + // The namespace keeps only the event's type and attributes, so a failover replay that + // re-creates the event object (with a new event id) reproduces the same ids. + SubagentIdAllocator ctx1 = + new SubagentIdAllocator("key-1", 7L, "actionA", new InputEvent("payload"), "agent"); + SubagentIdAllocator ctx2 = + new SubagentIdAllocator("key-1", 7L, "actionA", new InputEvent("payload"), "agent"); + + assertEquals(ctx1.nextSessionId(), ctx2.nextSessionId()); + } + + @Test + void attributeMapIterationOrderDoesNotAffectNamespace() { + // A replay on another JVM may iterate the rebuilt attribute map differently, so the digest + // mapper sorts entries: identical contents in reversed order must yield the same id. + Map forward = new LinkedHashMap<>(); + forward.put("alpha", "1"); + forward.put("beta", "2"); + forward.put("gamma", "3"); + Map reversed = new LinkedHashMap<>(); + reversed.put("gamma", "3"); + reversed.put("beta", "2"); + reversed.put("alpha", "1"); + + SubagentIdAllocator ctx1 = + new SubagentIdAllocator( + "key-1", 7L, "actionA", new Event("my.EventType", forward), "agent"); + SubagentIdAllocator ctx2 = + new SubagentIdAllocator( + "key-1", 7L, "actionA", new Event("my.EventType", reversed), "agent"); + + assertEquals(ctx1.nextSessionId(), ctx2.nextSessionId()); + } + + @Test + void deterministicAcrossTwoInstancesWithSameNamespace() { + SubagentIdAllocator ctx1 = newAllocator("key-1", 7L, "actionA", "e-1"); + SubagentIdAllocator ctx2 = newAllocator("key-1", 7L, "actionA", "e-1"); + + // Same call sequence replayed on two independently constructed instances that share an + // identical namespace must reproduce byte-for-byte identical ids -- this is the core + // property that makes failover replay reproduce the same sub-agent identities. + String session1a = ctx1.nextSessionId(); + String session2a = ctx2.nextSessionId(); + assertEquals(session1a, session2a); + + String call1a = ctx1.nextCallId(session1a); + String call2a = ctx2.nextCallId(session2a); + assertEquals(call1a, call2a); + + String session1b = ctx1.nextSessionId(); + String session2b = ctx2.nextSessionId(); + assertEquals(session1b, session2b); + assertNotEquals(session1a, session1b); + } + + @Test + void namespaceSeparatesOnKey() { + SubagentIdAllocator ctx1 = newAllocator("key-1", 7L, "actionA", "e-1"); + SubagentIdAllocator ctx2 = newAllocator("key-2", 7L, "actionA", "e-1"); + + // The ordinal is 0 on the first call for both instances, so any difference in the + // returned id must come from the namespace digest alone. + assertNotEquals(ctx1.nextSessionId(), ctx2.nextSessionId()); + } + + @Test + void namespaceSeparatesOnSequenceNumber() { + SubagentIdAllocator ctx1 = newAllocator("key-1", 7L, "actionA", "e-1"); + SubagentIdAllocator ctx2 = newAllocator("key-1", 8L, "actionA", "e-1"); + + assertNotEquals(ctx1.nextSessionId(), ctx2.nextSessionId()); + } + + @Test + void namespaceSeparatesOnActionName() { + SubagentIdAllocator ctx1 = newAllocator("key-1", 7L, "actionA", "e-1"); + SubagentIdAllocator ctx2 = newAllocator("key-1", 7L, "actionB", "e-1"); + + assertNotEquals(ctx1.nextSessionId(), ctx2.nextSessionId()); + } + + @Test + void namespaceSeparatesOnTriggeringEventAloneSiblingTaskScenario() { + // Two sibling tasks sharing key/sequenceNumber/actionName but triggered by different + // events: the event attributes alone must keep their identities from colliding. + SubagentIdAllocator ctx1 = newAllocator("key-1", 7L, "actionA", "e-1"); + SubagentIdAllocator ctx2 = newAllocator("key-1", 7L, "actionA", "e-2"); + + assertNotEquals(ctx1.nextSessionId(), ctx2.nextSessionId()); + } + + @Test + void namespaceSeparatesOnAgentName() { + // Two sub-agents of the same action execution share the caller's counting range but must + // never hand out the same ids: the agent name alone separates their namespaces. + SubagentIdAllocator ctx1 = + new SubagentIdAllocator("key-1", 7L, "actionA", new InputEvent("e-1"), "agent-a"); + SubagentIdAllocator ctx2 = + new SubagentIdAllocator("key-1", 7L, "actionA", new InputEvent("e-1"), "agent-b"); + + assertNotEquals(ctx1.nextSessionId(), ctx2.nextSessionId()); + } + + @Test + void sessionOrdinalsIncreaseAndAreUnique() { + SubagentIdAllocator ctx = newAllocator("key-1", 1L, "actionA", "e-1"); + + String s0 = ctx.nextSessionId(); + String s1 = ctx.nextSessionId(); + String s2 = ctx.nextSessionId(); + + // The namespace digest is a fixed-length UUID string, so "-" unambiguously identifies + // the ordinal suffix. + assertTrue(s0.endsWith("-0")); + assertTrue(s1.endsWith("-1")); + assertTrue(s2.endsWith("-2")); + + Set unique = new HashSet<>(); + unique.add(s0); + unique.add(s1); + unique.add(s2); + assertEquals(3, unique.size()); + } + + @Test + void perSessionCallOrdinalStartsAtOneAndIncrementsPerSession() { + SubagentIdAllocator ctx = newAllocator("key-1", 1L, "actionA", "e-1"); + + String sessionA = ctx.nextSessionId(); + String sessionB = ctx.nextSessionId(); + + String callA1 = ctx.nextCallId(sessionA); + String callA2 = ctx.nextCallId(sessionA); + String callB1 = ctx.nextCallId(sessionB); + + assertTrue(callA1.endsWith("-1")); + assertTrue(callA2.endsWith("-2")); + // A different session's ordinal is tracked independently and also starts at 1, rather + // than continuing sessionA's running count. + assertTrue(callB1.endsWith("-1")); + + assertNotEquals(callA1, callA2); + assertNotEquals(callA1, callB1); + } + + @Test + void callIdIsSessionIdPlusOrdinal() { + SubagentIdAllocator ctx = newAllocator("key-1", 1L, "actionA", "e-1"); + + // The call id is formed by appending the per-session ordinal to the session id, which + // already carries the namespace digest, so no further hashing is involved. + String sessionId = ctx.nextSessionId(); + assertEquals(sessionId + "-1", ctx.nextCallId(sessionId)); + assertEquals(sessionId + "-2", ctx.nextCallId(sessionId)); + } + + @Test + void explicitSessionIdIsUsedVerbatimNotParsed() { + SubagentIdAllocator ctx = newAllocator("key-1", 1L, "actionA", "e-1"); + + // Caller-supplied session ids need not follow the "{digest}-{ordinal}" shape; any string + // is legal and distinct ids never collide (they must not be reused across executions). + String callForExplicit1 = ctx.nextCallId("checkout-session-42"); + String callForExplicit2 = ctx.nextCallId("checkout-session-43"); + assertEquals("checkout-session-42-1", callForExplicit1); + assertEquals("checkout-session-43-1", callForExplicit2); + assertNotEquals(callForExplicit1, callForExplicit2); + + // Determinism holds for explicit session ids too: a second, freshly constructed allocator + // sharing the same namespace reproduces the same id for the same explicit session id. + SubagentIdAllocator ctx2 = newAllocator("key-1", 1L, "actionA", "e-1"); + assertEquals(callForExplicit1, ctx2.nextCallId("checkout-session-42")); + } + + @Test + void reusedSessionIdRestartsCallOrdinalInANewActionExecution() { + // Contract pinned here: ids assigned without an explicit id are only valid within one + // action execution. Each task constructs its own allocator, so when a later action reuses + // a session id, the per-session call ordinal starts at 1 again and the same call ids + // recur instead of continuing the earlier action's count. + SubagentIdAllocator firstTask = newAllocator("key-1", 1L, "actionA", "e-1"); + assertEquals("shared-session-1", firstTask.nextCallId("shared-session")); + assertEquals("shared-session-2", firstTask.nextCallId("shared-session")); + + // A different task (other sequence number / triggering event) allocates under the same + // session id and hands out the identical ids again. + SubagentIdAllocator secondTask = newAllocator("key-1", 2L, "actionA", "e-2"); + assertEquals("shared-session-1", secondTask.nextCallId("shared-session")); + } + + @Test + void idsAreOpaqueAndDoNotLeakPlainKeyOrActionName() { + SubagentIdAllocator ctx = newAllocator("super-secret-key", 1L, "myCustomActionName", "e-1"); + + String sessionId = ctx.nextSessionId(); + String callId = ctx.nextCallId(sessionId); + + assertFalse(sessionId.contains("super-secret-key")); + assertFalse(sessionId.contains("myCustomActionName")); + assertFalse(callId.contains("super-secret-key")); + assertFalse(callId.contains("myCustomActionName")); + } +} From 3024260ee0e82ccb57666f92ef788be757f0a3a1 Mon Sep 17 00:00:00 2001 From: "luogen.lg" Date: Wed, 5 Aug 2026 20:00:36 +0800 Subject: [PATCH 09/11] [runtime][python] Add deferred execution mode for external sub-agent setups --- .../flink_agents/runtime/deferred_subagent.py | 180 +++++++ .../flink_agents/runtime/subagent_handles.py | 21 +- .../runtime/tests/test_deferred_subagent.py | 452 ++++++++++++++++++ .../subagent/BaseDeferredSubagentSetup.java | 66 +++ .../subagent/DeferredSubagentFuture.java | 129 +++++ .../runtime/subagent/SubagentFutureGroup.java | 12 + .../subagent/DeferredSubagentSubmitTest.java | 322 +++++++++++++ .../subagent/MockDeferredSubagentSetup.java | 108 +++++ 8 files changed, 1289 insertions(+), 1 deletion(-) create mode 100644 python/flink_agents/runtime/deferred_subagent.py create mode 100644 python/flink_agents/runtime/tests/test_deferred_subagent.py create mode 100644 runtime/src/main/java/org/apache/flink/agents/runtime/subagent/BaseDeferredSubagentSetup.java create mode 100644 runtime/src/main/java/org/apache/flink/agents/runtime/subagent/DeferredSubagentFuture.java create mode 100644 runtime/src/test/java/org/apache/flink/agents/runtime/subagent/DeferredSubagentSubmitTest.java create mode 100644 runtime/src/test/java/org/apache/flink/agents/runtime/subagent/MockDeferredSubagentSetup.java 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/subagent_handles.py b/python/flink_agents/runtime/subagent_handles.py index bf767cc4c..07d4f91e9 100644 --- a/python/flink_agents/runtime/subagent_handles.py +++ b/python/flink_agents/runtime/subagent_handles.py @@ -117,7 +117,26 @@ def combine(self, *others: SubagentFuture) -> SubagentFutures: return SubagentFutureGroup((*self._futures, *others)) def __await__(self) -> Any: - """Wait for every handle in submission order.""" + """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__() 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/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. + * + *

      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. + * + *

      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. + * + *

      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 prepare( + RunnerContext ctx, Object prompt, String sessionId, String callId); +} 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> preparedSupplier; + + @Nullable private DurableCallable prepared; + private boolean done; + private boolean cancelled; + @Nullable private SubagentResult value; + + public DeferredSubagentFuture( + String sessionId, + String callId, + RunnerContext ctx, + @Nullable PendingSubagentCallRegistry registry, + Supplier> 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 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. + * + *

      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/SubagentFutureGroup.java b/runtime/src/main/java/org/apache/flink/agents/runtime/subagent/SubagentFutureGroup.java index 303e44cd2..c7312de46 100644 --- 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 @@ -58,6 +58,18 @@ public boolean isDone() { @Override public List 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 outcomes = new ArrayList<>(futures.size()); for (SubagentFuture future : futures) { outcomes.add(future.await()); diff --git a/runtime/src/test/java/org/apache/flink/agents/runtime/subagent/DeferredSubagentSubmitTest.java b/runtime/src/test/java/org/apache/flink/agents/runtime/subagent/DeferredSubagentSubmitTest.java new file mode 100644 index 000000000..cfc46c508 --- /dev/null +++ b/runtime/src/test/java/org/apache/flink/agents/runtime/subagent/DeferredSubagentSubmitTest.java @@ -0,0 +1,322 @@ +/* + * 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.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.api.java.functions.KeySelector; +import org.apache.flink.api.java.typeutils.TypeExtractor; +import org.apache.flink.streaming.runtime.streamrecord.StreamRecord; +import org.apache.flink.streaming.util.KeyedOneInputStreamOperatorTestHarness; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; + +import java.util.List; +import java.util.concurrent.CancellationException; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; +import static org.junit.jupiter.api.Assertions.assertThrows; + +/** + * The deferred execution mode of {@link BaseDeferredSubagentSetup}: submit only registers the + * invocation, the request is issued when the handle is resolved, several handles can be resolved + * together, and a dropped handle fails the action through the base's per-task registry. Ids are + * supplied explicitly by the caller actions. + */ +public class DeferredSubagentSubmitTest { + + private static final String RESOURCE_NAME = "capturing"; + + private final RunnerContext ctx = null; + + @BeforeEach + public void resetCaptures() { + MockDeferredSubagentSetup.reset(); + } + + // The short forms inherit the base's deterministic assignment, which needs a prepared task. + + @Test + void shortFormsRequireAPreparedTask() throws Exception { + BaseDeferredSubagentSetup setup = new MockDeferredSubagentSetup(); + + assertThrows(IllegalStateException.class, () -> setup.submit(ctx, "ping", "sid-1")); + assertThrows(IllegalStateException.class, () -> setup.submit(ctx, "ping")); + } + + // Batched resolve: every pending deferred handle is prepared up front, then the prepared + // calls are executed one by one in submission order. + + @SuppressWarnings("unused") + public static void batched(Event event, RunnerContext ctx) throws Exception { + BaseDeferredSubagentSetup setup = + (BaseDeferredSubagentSetup) ctx.getResource(RESOURCE_NAME, ResourceType.AGENT); + SubagentFuture first = setup.submit(ctx, "batch-a", "session", "call-1"); + SubagentFuture second = setup.submit(ctx, "batch-b", "session", "call-2"); + // Nothing has been issued yet: submit only created the deferred handles. + if (!MockDeferredSubagentSetup.captures().isEmpty()) { + throw new IllegalStateException( + "deferred submit issued the request too early: " + + MockDeferredSubagentSetup.captures()); + } + List results = first.combine(second).awaitAll(); + ctx.sendEvent( + new OutputEvent(results.get(0).getResult() + "|" + results.get(1).getResult())); + } + + @Test + void batchedResolveResolvesEveryHandleInSubmissionOrder() throws Exception { + try (KeyedOneInputStreamOperatorTestHarness harness = + harness(plan("batched", new MockDeferredSubagentSetup()))) { + harness.open(); + run(harness, 1L); + + assertThat(MockDeferredSubagentSetup.captures()).hasSize(2); + assertThat(MockDeferredSubagentSetup.executionCount()).isEqualTo(2); + assertThat(harness.getRecordOutput()).hasSize(1); + } + } + + // A dropped handle fails the action instead of silently skipping the call: the base records + // every handle in its per-task registry and checks it when the task finishes. + + @SuppressWarnings("unused") + public static void dropsHandle(Event event, RunnerContext ctx) throws Exception { + BaseDeferredSubagentSetup setup = + (BaseDeferredSubagentSetup) ctx.getResource(RESOURCE_NAME, ResourceType.AGENT); + setup.submit(ctx, "dropped", "session", "call-1"); + ctx.sendEvent(new OutputEvent("done")); + } + + @Test + void droppingADeferredHandleFailsTheAction() throws Exception { + try (KeyedOneInputStreamOperatorTestHarness harness = + harness(plan("dropsHandle", new MockDeferredSubagentSetup()))) { + harness.open(); + + assertThatThrownBy(() -> run(harness, 1L)) + .rootCause() + .isInstanceOf(IllegalStateException.class) + .hasMessageContaining("without resolving the sub-agent calls it submitted"); + assertThat(MockDeferredSubagentSetup.executionCount()).isZero(); + } + } + + // An already-resolved handle simply contributes its value to the batch. + + @SuppressWarnings("unused") + public static void batchesResolvedHandle(Event event, RunnerContext ctx) throws Exception { + BaseDeferredSubagentSetup setup = + (BaseDeferredSubagentSetup) ctx.getResource(RESOURCE_NAME, ResourceType.AGENT); + SubagentFuture resolved = new CompletedSubagentFuture("s", "c", SubagentResult.ok("x")); + SubagentFuture pending = setup.submit(ctx, "pending", "session", "call-1"); + List results = resolved.combine(pending).awaitAll(); + ctx.sendEvent( + new OutputEvent(results.get(0).getResult() + "|" + results.get(1).getResult())); + } + + @Test + void batchingAnAlreadyResolvedHandleJoinsTheBatch() throws Exception { + try (KeyedOneInputStreamOperatorTestHarness harness = + harness(plan("batchesResolvedHandle", new MockDeferredSubagentSetup()))) { + harness.open(); + run(harness, 1L); + + // Only the pending handle issued a request; the resolved one contributed its value. + assertThat(MockDeferredSubagentSetup.captures()).hasSize(1); + assertThat(MockDeferredSubagentSetup.executionCount()).isEqualTo(1); + assertThat(harness.getRecordOutput()).hasSize(1); + } + } + + // Cancellation: the request was never issued, so cancelling discards it and resolving the + // handle fails with a CancellationException. + + @SuppressWarnings("unused") + public static void cancelsBeforeResolve(Event event, RunnerContext ctx) throws Exception { + BaseDeferredSubagentSetup setup = + (BaseDeferredSubagentSetup) ctx.getResource(RESOURCE_NAME, ResourceType.AGENT); + SubagentFuture handle = setup.submit(ctx, "cancelled", "session", "call-1"); + handle.cancel(); + if (!MockDeferredSubagentSetup.captures().isEmpty()) { + throw new IllegalStateException("cancelled handle created its callable"); + } + try { + handle.await(); + throw new IllegalStateException("cancelled handle resolved"); + } catch (CancellationException expected) { + // The request was never issued; cancellation fails the resolve. + } + ctx.sendEvent(new OutputEvent("cancelled")); + } + + @Test + void cancellingBeforeResolveDiscardsTheRequest() throws Exception { + try (KeyedOneInputStreamOperatorTestHarness harness = + harness(plan("cancelsBeforeResolve", new MockDeferredSubagentSetup()))) { + harness.open(); + run(harness, 1L); + + assertThat(MockDeferredSubagentSetup.captures()).isEmpty(); + assertThat(MockDeferredSubagentSetup.executionCount()).isZero(); + assertThat(harness.getRecordOutput()).hasSize(1); + } + } + + @SuppressWarnings("unused") + public static void cancelsThroughTheGroup(Event event, RunnerContext ctx) throws Exception { + BaseDeferredSubagentSetup setup = + (BaseDeferredSubagentSetup) ctx.getResource(RESOURCE_NAME, ResourceType.AGENT); + SubagentFuture first = setup.submit(ctx, "batch-a", "session", "call-1"); + SubagentFuture second = setup.submit(ctx, "batch-b", "session", "call-2"); + first.combine(second).cancel(); + try { + first.combine(second).awaitAll(); + throw new IllegalStateException("cancelled batch resolved"); + } catch (CancellationException expected) { + // Every handle in the batch received the cancellation. + } + ctx.sendEvent(new OutputEvent("cancelled")); + } + + @Test + void cancellingThroughTheGroupPropagatesToEveryHandle() throws Exception { + try (KeyedOneInputStreamOperatorTestHarness harness = + harness(plan("cancelsThroughTheGroup", new MockDeferredSubagentSetup()))) { + harness.open(); + run(harness, 1L); + + assertThat(MockDeferredSubagentSetup.captures()).isEmpty(); + assertThat(MockDeferredSubagentSetup.executionCount()).isZero(); + assertThat(harness.getRecordOutput()).hasSize(1); + } + } + + // A cancelled handle unregisters from the base's per-task registry, so the built-in check + // does not fail the action over it. + + @SuppressWarnings("unused") + public static void cancelsTrackedHandle(Event event, RunnerContext ctx) throws Exception { + BaseDeferredSubagentSetup setup = + (BaseDeferredSubagentSetup) ctx.getResource(RESOURCE_NAME, ResourceType.AGENT); + setup.submit(ctx, "cancelled", "session", "call-1").cancel(); + ctx.sendEvent(new OutputEvent("done")); + } + + @Test + void cancellingATrackedHandleDoesNotFailTheAction() throws Exception { + try (KeyedOneInputStreamOperatorTestHarness harness = + harness(plan("cancelsTrackedHandle", new MockDeferredSubagentSetup()))) { + harness.open(); + run(harness, 1L); + + assertThat(MockDeferredSubagentSetup.executionCount()).isZero(); + assertThat(harness.getRecordOutput()).hasSize(1); + } + } + + // A system-level failure escaping the prepared callable propagates and fails the action + // instead of being folded into an error result. + + @SuppressWarnings("unused") + public static void awaitsSystemFailingHandle(Event event, RunnerContext ctx) throws Exception { + BaseDeferredSubagentSetup setup = + (BaseDeferredSubagentSetup) ctx.getResource(RESOURCE_NAME, ResourceType.AGENT); + setup.submit(ctx, "boom", "session", "call-1").await(); + ctx.sendEvent(new OutputEvent("done")); + } + + @Test + void systemLevelFailurePropagatesInsteadOfFolding() throws Exception { + try (KeyedOneInputStreamOperatorTestHarness harness = + harness(plan("awaitsSystemFailingHandle", new ThrowingDeferredSubagentSetup()))) { + harness.open(); + + assertThatThrownBy(() -> run(harness, 1L)) + .rootCause() + .isInstanceOf(IllegalStateException.class) + .hasMessageContaining("durable execution crashed"); + assertThat(harness.getRecordOutput()).isEmpty(); + } + } + + /** Deferred setup whose prepared callable throws a system-level failure instead of folding. */ + private static final class ThrowingDeferredSubagentSetup extends BaseDeferredSubagentSetup { + @Override + protected DurableCallable prepare( + RunnerContext ctx, Object prompt, String sessionId, String callId) { + return new DurableCallable() { + @Override + public String getId() { + return sessionId + "#" + callId; + } + + @Override + public Class getResultClass() { + return SubagentResult.class; + } + + @Override + public SubagentResult call() { + throw new IllegalStateException("durable execution crashed"); + } + }; + } + } + + @SuppressWarnings("unchecked") + private static void run( + KeyedOneInputStreamOperatorTestHarness harness, long value) + throws Exception { + harness.processElement(new StreamRecord<>(value)); + ((ActionExecutionOperator) harness.getOperator()) + .waitInFlightEventsFinished(); + } + + private static AgentPlan plan(String actionMethod, BaseDeferredSubagentSetup setup) + throws Exception { + Agent agent = new Agent(); + agent.addResource(RESOURCE_NAME, ResourceType.AGENT, setup); + agent.addAction( + new String[] {InputEvent.EVENT_TYPE}, + DeferredSubagentSubmitTest.class.getMethod( + actionMethod, Event.class, RunnerContext.class)); + return new AgentPlan(agent, new AgentConfiguration()); + } + + private static KeyedOneInputStreamOperatorTestHarness harness( + AgentPlan plan) throws Exception { + return new KeyedOneInputStreamOperatorTestHarness<>( + new ActionExecutionOperatorFactory<>(plan, true), + (KeySelector) value -> value, + TypeExtractor.getForClass(Long.class)); + } +} diff --git a/runtime/src/test/java/org/apache/flink/agents/runtime/subagent/MockDeferredSubagentSetup.java b/runtime/src/test/java/org/apache/flink/agents/runtime/subagent/MockDeferredSubagentSetup.java new file mode 100644 index 000000000..0a8b69a71 --- /dev/null +++ b/runtime/src/test/java/org/apache/flink/agents/runtime/subagent/MockDeferredSubagentSetup.java @@ -0,0 +1,108 @@ +/* + * 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.SubagentResult; + +import java.util.ArrayList; +import java.util.Collections; +import java.util.List; +import java.util.concurrent.atomic.AtomicInteger; + +/** + * Mock integration of {@link BaseDeferredSubagentSetup}: echoes the prompt back as the resolved + * value. Records every {@code (sessionId, callId)} pair assigned to it. + * + *

      Capture happens at prepare time rather than in the call body, so tests can assert on assigned + * ids even when the durable call is served from cache; {@link #executionCount()} separately tracks + * real executions. State is static because one setup instance is shared by every resolving task — + * call {@link #reset()} before each independent scenario. + */ +public class MockDeferredSubagentSetup extends BaseDeferredSubagentSetup { + + /** One {@code (sessionId, callId)} assignment captured at prepare time. */ + public static final class Capture { + public final String sessionId; + public final String callId; + public final Object prompt; + + Capture(String sessionId, String callId, Object prompt) { + this.sessionId = sessionId; + this.callId = callId; + this.prompt = prompt; + } + + @Override + public String toString() { + return "Capture{sessionId=" + + sessionId + + ", callId=" + + callId + + ", prompt=" + + prompt + + "}"; + } + } + + private static final List CAPTURES = Collections.synchronizedList(new ArrayList<>()); + private static final AtomicInteger EXECUTION_COUNT = new AtomicInteger(); + + /** Clears all captures and the execution counter. Call before each independent scenario. */ + public static void reset() { + CAPTURES.clear(); + EXECUTION_COUNT.set(0); + } + + /** Snapshot of every assignment captured since the last {@link #reset()}, in creation order. */ + public static List captures() { + synchronized (CAPTURES) { + return new ArrayList<>(CAPTURES); + } + } + + /** Number of times a prepared call body actually ran. */ + public static int executionCount() { + return EXECUTION_COUNT.get(); + } + + @Override + protected DurableCallable prepare( + RunnerContext ctx, Object prompt, String sessionId, String callId) { + CAPTURES.add(new Capture(sessionId, callId, prompt)); + return new DurableCallable() { + @Override + public String getId() { + return sessionId + "#" + callId; + } + + @Override + public Class getResultClass() { + return SubagentResult.class; + } + + @Override + public SubagentResult call() { + EXECUTION_COUNT.incrementAndGet(); + return SubagentResult.ok(sessionId + "|" + callId + "|" + prompt); + } + }; + } +} From d1b249043c273fbfa2f027bf2a733765cbd80656 Mon Sep 17 00:00:00 2001 From: "luogen.lg" Date: Wed, 5 Aug 2026 14:15:16 +0800 Subject: [PATCH 10/11] [runtime][python] Add async external sub-agent base running in durable pub/sub mode --- python/flink_agents/runtime/async_subagent.py | 347 +++++++++++ .../runtime/tests/test_async_subagent.py | 571 ++++++++++++++++++ .../runtime/subagent/AsyncSubagentFuture.java | 125 ++++ .../subagent/BaseAsyncSubagentSetup.java | 287 +++++++++ .../subagent/BaseAsyncSubagentSetupTest.java | 521 ++++++++++++++++ .../subagent/MockAsyncSubagentSetup.java | 154 +++++ 6 files changed, 2005 insertions(+) create mode 100644 python/flink_agents/runtime/async_subagent.py create mode 100644 python/flink_agents/runtime/tests/test_async_subagent.py create mode 100644 runtime/src/main/java/org/apache/flink/agents/runtime/subagent/AsyncSubagentFuture.java create mode 100644 runtime/src/main/java/org/apache/flink/agents/runtime/subagent/BaseAsyncSubagentSetup.java create mode 100644 runtime/src/test/java/org/apache/flink/agents/runtime/subagent/BaseAsyncSubagentSetupTest.java create mode 100644 runtime/src/test/java/org/apache/flink/agents/runtime/subagent/MockAsyncSubagentSetup.java 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/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/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 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 submitRequest( + RunnerContext ctx, String sessionId, String callId, Object prompt) { + return new DurableCallable() { + @Override + public String getId() { + return sessionId + "#" + callId; + } + + @Override + public Class getResultClass() { + return Void.class; + } + + @Override + public Void call() throws Exception { + callSubmitRequest(sessionId, callId, prompt); + return null; + } + + @Override + public Callable 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 awaitResult( + RunnerContext ctx, String sessionId, String callId) { + return new DurableCallable() { + @Override + public String getId() { + return sessionId + "#" + callId + "#await"; + } + + @Override + public Class 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}. + * + *

      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/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 callable = setup.submitRequestForTest("ping", "sid-1", "call-1"); + + assertThat(callable.getId()).isEqualTo("sid-1#call-1"); + callable.call(); + + assertThat(setup.postCount()).isEqualTo(1); + assertThat(setup.statusQueryCount()).isZero(); + assertThat(setup.fetchCount()).isZero(); + } + + @Test + void postFailureFailsTheSubmit() throws Exception { + MockAsyncSubagentSetup setup = new MockAsyncSubagentSetup(2, true); + + DurableCallable callable = setup.submitRequestForTest("ping", "sid-1", "call-1"); + + assertThatThrownBy(callable::call) + .isInstanceOf(IllegalStateException.class) + .hasMessageContaining("post failed"); + assertThat(setup.statusQueryCount()).isZero(); + assertThat(setup.fetchCount()).isZero(); + } + + // ------------------------------------------------------------------------------------------ + // The crash-window reconciler of the POST + // ------------------------------------------------------------------------------------------ + + @Test + void reconcilerRepostsWhenTheRunIsNotOnRecord() throws Exception { + MockAsyncSubagentSetup setup = new MockAsyncSubagentSetup(1, false); + + setup.submitRequestForTest("ping", "sid-1", "call-1").reconciler().call(); + + // Probe reported NOT_STARTED, so the missing POST was issued exactly once. + assertThat(setup.postCount()).isEqualTo(1); + assertThat(setup.statusQueryCount()).isEqualTo(1); + } + + @Test + void reconcilerDoesNotRepostARunningRun() throws Exception { + MockAsyncSubagentSetup setup = new MockAsyncSubagentSetup(1, false); + setup.seedRun("sid-1", "call-1", "done:ping", null, 1); + + setup.submitRequestForTest("ping", "sid-1", "call-1").reconciler().call(); + + assertThat(setup.postCount()).isZero(); + assertThat(setup.statusQueryCount()).isEqualTo(1); + } + + @Test + void reconcilerDoesNotRepostATerminalRun() throws Exception { + MockAsyncSubagentSetup setup = new MockAsyncSubagentSetup(0, false); + setup.seedRun("sid-1", "call-1", "done:ping", null, 0); + + setup.submitRequestForTest("ping", "sid-1", "call-1").reconciler().call(); + + assertThat(setup.postCount()).isZero(); + assertThat(setup.statusQueryCount()).isEqualTo(1); + } + + @Test + void reconcilerTreatsAFailedRunAsLanded() throws Exception { + MockAsyncSubagentSetup setup = new MockAsyncSubagentSetup(0, false); + setup.seedRun("sid-1", "call-1", null, "run exploded", 0); + + setup.submitRequestForTest("ping", "sid-1", "call-1").reconciler().call(); + + assertThat(setup.postCount()).isZero(); + assertThat(setup.statusQueryCount()).isEqualTo(1); + } + + // ------------------------------------------------------------------------------------------ + // The sub: the await composition + // ------------------------------------------------------------------------------------------ + + @Test + void awaitPollsUntilTerminalThenFetches() throws Exception { + MockAsyncSubagentSetup setup = new MockAsyncSubagentSetup(2, false); + setup.seedRun("sid-1", "call-1", "done:ping", null, 2); + + DurableCallable await = setup.awaitResultForTest("sid-1", "call-1"); + SubagentResult result = await.call(); + + assertThat(await.getId()).isEqualTo("sid-1#call-1#await"); + assertThat(result.isSuccess()).isTrue(); + assertThat(result.getResult()).isEqualTo("done:ping"); + // Two RUNNING probes, the terminal one, then the separate fetch. + assertThat(setup.statusQueryCount()).isEqualTo(3); + assertThat(setup.fetchCount()).isEqualTo(1); + } + + @Test + void awaitSurfacesAFailedRunWithoutFetching() throws Exception { + MockAsyncSubagentSetup setup = new MockAsyncSubagentSetup(0, false); + setup.seedRun("sid-1", "call-1", null, "run exploded", 0); + + SubagentResult result = setup.awaitResultForTest("sid-1", "call-1").call(); + + assertThat(result.isSuccess()).isFalse(); + assertThat(result.getErrorMessage()).contains("run exploded"); + assertThat(setup.fetchCount()).isZero(); + } + + // ------------------------------------------------------------------------------------------ + // Failover replay: fresh probes may take a different path to the same result + // ------------------------------------------------------------------------------------------ + + @Test + void replayAfterTheRunCompletedTakesFewerProbes() throws Exception { + // Original execution: the run completes only after two RUNNING probes. + MockAsyncSubagentSetup original = new MockAsyncSubagentSetup(2, false); + original.seedRun("sid-1", "call-1", "done:ping", null, 2); + SubagentResult before = original.awaitResultForTest("sid-1", "call-1").call(); + assertThat(original.statusQueryCount()).isEqualTo(3); + + // Replay: the run has already reached a terminal state, so the same await takes a + // shorter path — fewer probes — to the same result. + MockAsyncSubagentSetup replay = new MockAsyncSubagentSetup(2, false); + replay.seedRun("sid-1", "call-1", "done:ping", null, 0); + SubagentResult after = replay.awaitResultForTest("sid-1", "call-1").call(); + + assertThat(after.isSuccess()).isTrue(); + assertThat(after.getResult()).isEqualTo(before.getResult()); + assertThat(replay.statusQueryCount()).isEqualTo(1); + } + + // ------------------------------------------------------------------------------------------ + // Pipeline flows: pub lands immediately, the handle subscribes to the run + // ------------------------------------------------------------------------------------------ + + /** Submits through the short form and resolves directly, without probing isDone first. */ + @SuppressWarnings("unused") + public static void delegated(Event event, RunnerContext ctx) throws Exception { + MockAsyncSubagentSetup setup = + (MockAsyncSubagentSetup) ctx.getResource(RESOURCE_NAME, ResourceType.AGENT); + SubagentFuture future = setup.submit(ctx, "ping"); + SubagentResult result = future.await(); + ctx.sendEvent(new OutputEvent(result.getResult())); + } + + @Test + void pipelineSubmitsThenResolvesDirectly() throws Exception { + MockAsyncSubagentSetup setup = new MockAsyncSubagentSetup(2, false); + try (KeyedOneInputStreamOperatorTestHarness harness = + harness(plan("delegated", setup))) { + harness.open(); + run(harness, 1L); + + assertThat(setup.postCount()).isEqualTo(1); + assertThat(setup.fetchCount()).isEqualTo(1); + assertThat(harness.getRecordOutput()) + .singleElement() + .extracting(StreamRecord::getValue) + .isEqualTo("done:ping"); + } + } + + /** Probes isDone until the run turns terminal, then resolves. */ + @SuppressWarnings("unused") + public static void pollsThenResolves(Event event, RunnerContext ctx) throws Exception { + MockAsyncSubagentSetup setup = + (MockAsyncSubagentSetup) ctx.getResource(RESOURCE_NAME, ResourceType.AGENT); + SubagentFuture future = setup.submit(ctx, "ping"); + int pendingProbes = 0; + while (!future.isDone()) { + pendingProbes++; + } + SubagentResult result = future.await(); + ctx.sendEvent(new OutputEvent("seen:" + pendingProbes + "|" + result.getResult())); + } + + @Test + void pipelineProbesStatusDirectlyUntilTerminal() throws Exception { + MockAsyncSubagentSetup setup = new MockAsyncSubagentSetup(2, false); + try (KeyedOneInputStreamOperatorTestHarness harness = + harness(plan("pollsThenResolves", setup))) { + harness.open(); + run(harness, 1L); + + // Two RUNNING probes from isDone, the terminal probe and fetch from await. + assertThat(setup.statusQueryCount()).isEqualTo(4); + assertThat(setup.fetchCount()).isEqualTo(1); + assertThat(harness.getRecordOutput()) + .singleElement() + .extracting(StreamRecord::getValue) + .isEqualTo("seen:2|done:ping"); + } + } + + /** Submits and fails when the POST endpoint rejects the run. */ + @SuppressWarnings("unused") + public static void postFails(Event event, RunnerContext ctx) throws Exception { + MockAsyncSubagentSetup setup = + (MockAsyncSubagentSetup) ctx.getResource(RESOURCE_NAME, ResourceType.AGENT); + setup.submit(ctx, "ping"); + ctx.sendEvent(new OutputEvent("unreachable")); + } + + @Test + void postFailureFailsTheAction() throws Exception { + MockAsyncSubagentSetup setup = new MockAsyncSubagentSetup(2, true); + try (KeyedOneInputStreamOperatorTestHarness harness = + harness(plan("postFails", setup))) { + harness.open(); + + assertThatThrownBy(() -> run(harness, 1L)) + .rootCause() + .isInstanceOf(IllegalStateException.class) + .hasMessageContaining("post failed"); + assertThat(harness.getRecordOutput()).isEmpty(); + } + } + + /** Cancels after the pub and resolves: the default disposition is a CancellationException. */ + @SuppressWarnings("unused") + public static void cancelsThenResolves(Event event, RunnerContext ctx) throws Exception { + MockAsyncSubagentSetup setup = + (MockAsyncSubagentSetup) ctx.getResource(RESOURCE_NAME, ResourceType.AGENT); + SubagentFuture future = setup.submit(ctx, "ping"); + future.cancel(); + try { + future.await(); + throw new IllegalStateException("cancelled handle resolved"); + } catch (CancellationException expected) { + // The cancel hook returned nothing, so the resolve fails as cancelled. + } + ctx.sendEvent(new OutputEvent("cancelled")); + } + + @Test + void cancelBeforeResolveThrowsCancellationByDefault() throws Exception { + MockAsyncSubagentSetup setup = new MockAsyncSubagentSetup(2, false); + try (KeyedOneInputStreamOperatorTestHarness harness = + harness(plan("cancelsThenResolves", setup))) { + harness.open(); + run(harness, 1L); + + // The pub landed, but the cancelled resolve never awaited nor fetched. + assertThat(setup.postCount()).isEqualTo(1); + assertThat(setup.cancelCount()).isEqualTo(1); + assertThat(setup.statusQueryCount()).isZero(); + assertThat(setup.fetchCount()).isZero(); + assertThat(harness.getRecordOutput()).hasSize(1); + } + } + + /** Cancels twice: a repeated cancel on the same handle is a local no-op. */ + @SuppressWarnings("unused") + public static void cancelsTwice(Event event, RunnerContext ctx) throws Exception { + MockAsyncSubagentSetup setup = + (MockAsyncSubagentSetup) ctx.getResource(RESOURCE_NAME, ResourceType.AGENT); + SubagentFuture future = setup.submit(ctx, "ping"); + future.cancel(); + future.cancel(); + try { + future.await(); + throw new IllegalStateException("cancelled handle resolved"); + } catch (CancellationException expected) { + // Repeated cancellations are harmless; remote cancels are idempotent. + } + ctx.sendEvent(new OutputEvent("cancelled")); + } + + @Test + void repeatedCancelIsALocalNoOp() throws Exception { + MockAsyncSubagentSetup setup = new MockAsyncSubagentSetup(2, false); + try (KeyedOneInputStreamOperatorTestHarness harness = + harness(plan("cancelsTwice", setup))) { + harness.open(); + run(harness, 1L); + + // The second cancel on the same handle does not propagate again; a failover + // replay creates a fresh handle and may propagate again (idempotent remotely). + assertThat(setup.cancelCount()).isEqualTo(1); + assertThat(setup.fetchCount()).isZero(); + assertThat(harness.getRecordOutput()).hasSize(1); + } + } + + /** Continues a session: the second invocation reuses the first handle's session id. */ + @SuppressWarnings("unused") + public static void continuesASession(Event event, RunnerContext ctx) throws Exception { + MockAsyncSubagentSetup setup = + (MockAsyncSubagentSetup) ctx.getResource(RESOURCE_NAME, ResourceType.AGENT); + SubagentFuture first = setup.submit(ctx, "turn-1"); + SubagentResult firstOutcome = first.await(); + SubagentFuture second = setup.submit(ctx, "turn-2", first.getSessionId()); + SubagentResult secondOutcome = second.await(); + ctx.sendEvent( + new OutputEvent( + first.getSessionId().equals(second.getSessionId()) + + "|" + + !first.getCallId().equals(second.getCallId()) + + "|" + + firstOutcome.getResult() + + "|" + + secondOutcome.getResult())); + } + + @Test + void multiTurnContinuationReusesTheSessionFromTheHandle() throws Exception { + MockAsyncSubagentSetup setup = new MockAsyncSubagentSetup(0, false); + try (KeyedOneInputStreamOperatorTestHarness harness = + harness(plan("continuesASession", setup))) { + harness.open(); + run(harness, 1L); + + assertThat(setup.postCount()).isEqualTo(2); + assertThat(harness.getRecordOutput()) + .singleElement() + .extracting(StreamRecord::getValue) + .isEqualTo("true|true|done:turn-1|done:turn-2"); + } + } + + /** Batches two open handles and resolves them together. */ + @SuppressWarnings("unused") + public static void batchesOpenHandles(Event event, RunnerContext ctx) throws Exception { + MockAsyncSubagentSetup setup = + (MockAsyncSubagentSetup) ctx.getResource(RESOURCE_NAME, ResourceType.AGENT); + SubagentFuture first = setup.submit(ctx, "a"); + SubagentFuture second = setup.submit(ctx, "b"); + List results = first.combine(second).awaitAll(); + ctx.sendEvent( + new OutputEvent(results.get(0).getResult() + "|" + results.get(1).getResult())); + } + + @Test + void combineResolvesEveryOpenHandle() throws Exception { + MockAsyncSubagentSetup setup = new MockAsyncSubagentSetup(0, false); + try (KeyedOneInputStreamOperatorTestHarness harness = + harness(plan("batchesOpenHandles", setup))) { + harness.open(); + run(harness, 1L); + + assertThat(setup.postCount()).isEqualTo(2); + assertThat(setup.fetchCount()).isEqualTo(2); + assertThat(harness.getRecordOutput()) + .singleElement() + .extracting(StreamRecord::getValue) + .isEqualTo("done:a|done:b"); + } + } + + /** Submits a handle and never resolves it. */ + @SuppressWarnings("unused") + public static void dropsHandle(Event event, RunnerContext ctx) throws Exception { + MockAsyncSubagentSetup setup = + (MockAsyncSubagentSetup) ctx.getResource(RESOURCE_NAME, ResourceType.AGENT); + setup.submit(ctx, "dropped"); + ctx.sendEvent(new OutputEvent("done")); + } + + @Test + void droppingAnOpenHandleFailsTheAction() throws Exception { + MockAsyncSubagentSetup setup = new MockAsyncSubagentSetup(2, false); + try (KeyedOneInputStreamOperatorTestHarness harness = + harness(plan("dropsHandle", setup))) { + harness.open(); + + // The pub landed at submit, but the open handle was dropped without collecting + // its outcome: the base's per-task registry fails the action over it. + assertThatThrownBy(() -> run(harness, 1L)) + .rootCause() + .isInstanceOf(IllegalStateException.class) + .hasMessageContaining("without resolving the sub-agent calls it submitted"); + assertThat(setup.postCount()).isEqualTo(1); + assertThat(setup.fetchCount()).isZero(); + } + } + + // ------------------------------------------------------------------------------------------ + // The descriptor argument: status_poll_interval_millis + // ------------------------------------------------------------------------------------------ + + /** Setup constructed like YAML-declared entries: through the descriptor constructor. */ + private static final class DescriptorAsyncSetup extends BaseAsyncSubagentSetup { + + private DescriptorAsyncSetup(ResourceDescriptor descriptor, ResourceContext context) { + super(descriptor, context); + } + + @Override + protected void callSubmitRequest(String sessionId, String callId, Object prompt) {} + + @Override + protected RunStatus callQueryStatus(String sessionId, String callId) { + return RunStatus.completed(); + } + + @Override + protected SubagentResult callFetchResult(String sessionId, String callId) { + return SubagentResult.ok("done"); + } + + private long pollIntervalMillis() { + return statusPollIntervalMillis; + } + } + + @Test + void descriptorArgumentSetsTheStatusPollInterval() { + DescriptorAsyncSetup setup = + new DescriptorAsyncSetup( + new ResourceDescriptor( + DescriptorAsyncSetup.class.getName(), + Map.of("status_poll_interval_millis", 123)), + new ResourceContextImpl((name, type) -> null)); + + assertThat(setup.pollIntervalMillis()).isEqualTo(123); + } + + @Test + void descriptorWithoutTheArgumentKeepsTheDefaultPollInterval() { + DescriptorAsyncSetup setup = + new DescriptorAsyncSetup( + new ResourceDescriptor(DescriptorAsyncSetup.class.getName(), Map.of()), + new ResourceContextImpl((name, type) -> null)); + + assertThat(setup.pollIntervalMillis()).isEqualTo(500); + } + + // ------------------------------------------------------------------------------------------ + // Harness plumbing + // ------------------------------------------------------------------------------------------ + + @SuppressWarnings("unchecked") + private static void run( + KeyedOneInputStreamOperatorTestHarness harness, long value) + throws Exception { + harness.processElement(new StreamRecord<>(value)); + ((ActionExecutionOperator) harness.getOperator()) + .waitInFlightEventsFinished(); + } + + private static AgentPlan plan(String actionMethod, MockAsyncSubagentSetup setup) + throws Exception { + Agent agent = new Agent(); + agent.addResource(RESOURCE_NAME, ResourceType.AGENT, setup); + agent.addAction( + new String[] {InputEvent.EVENT_TYPE}, + BaseAsyncSubagentSetupTest.class.getMethod( + actionMethod, Event.class, RunnerContext.class)); + return new AgentPlan(agent, new AgentConfiguration()); + } + + private static KeyedOneInputStreamOperatorTestHarness harness( + AgentPlan plan) throws Exception { + return new KeyedOneInputStreamOperatorTestHarness<>( + new ActionExecutionOperatorFactory<>(plan, true), + (KeySelector) value -> value, + TypeExtractor.getForClass(Long.class)); + } +} diff --git a/runtime/src/test/java/org/apache/flink/agents/runtime/subagent/MockAsyncSubagentSetup.java b/runtime/src/test/java/org/apache/flink/agents/runtime/subagent/MockAsyncSubagentSetup.java new file mode 100644 index 000000000..f5d9818ec --- /dev/null +++ b/runtime/src/test/java/org/apache/flink/agents/runtime/subagent/MockAsyncSubagentSetup.java @@ -0,0 +1,154 @@ +/* + * 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.subagent.SubagentResult; + +import javax.annotation.Nullable; + +import java.util.HashMap; +import java.util.Map; + +/** + * Example integration of {@link BaseAsyncSubagentSetup}: an in-memory asynchronous agent service. + * Demonstrates that an integration only supplies the transport primitives — {@link + * #callSubmitRequest}, {@link #callQueryStatus} and {@link #callFetchResult} — plus the optional + * cancel hook; all durable composition and recovery knowledge stays in the base. Counters let tests + * assert how many times each endpoint was hit. + */ +public class MockAsyncSubagentSetup extends BaseAsyncSubagentSetup { + + /** One recorded remote run, keyed by {@code sessionId#callId}. */ + private static final class Run { + private final Object result; + @Nullable private final String error; + private int queriesRemaining; + + private Run(Object result, @Nullable String error, int queriesRemaining) { + this.result = result; + this.error = error; + this.queriesRemaining = queriesRemaining; + } + } + + private final Map runs = new HashMap<>(); + private final int queriesUntilComplete; + private final boolean failOnPost; + + private int postCount; + private int statusQueryCount; + private int fetchCount; + private int cancelCount; + + public MockAsyncSubagentSetup() { + this(2, false); + } + + /** + * Creates a setup whose runs need {@code queriesUntilComplete} RUNNING probes before turning + * terminal; {@code failOnPost} makes every submission fail. + */ + public MockAsyncSubagentSetup(int queriesUntilComplete, boolean failOnPost) { + this.queriesUntilComplete = queriesUntilComplete; + this.failOnPost = failOnPost; + // Runs turn terminal after a fixed number of probes rather than after elapsed time, so + // probing without a delay keeps the counts identical and the tests fast. + this.statusPollIntervalMillis = 0; + } + + @Override + protected void callSubmitRequest(String sessionId, String callId, Object prompt) { + postCount++; + if (failOnPost) { + throw new IllegalStateException("post failed"); + } + runs.put(sessionId + "#" + callId, new Run("done:" + prompt, null, queriesUntilComplete)); + } + + @Override + protected RunStatus callQueryStatus(String sessionId, String callId) { + statusQueryCount++; + Run run = runs.get(sessionId + "#" + callId); + if (run == null) { + return RunStatus.notStarted(); + } + if (run.queriesRemaining > 0) { + run.queriesRemaining--; + return RunStatus.running(); + } + return run.error == null ? RunStatus.completed() : RunStatus.failed(run.error); + } + + @Override + protected SubagentResult callFetchResult(String sessionId, String callId) { + fetchCount++; + Run run = runs.get(sessionId + "#" + callId); + if (run == null) { + return SubagentResult.error("no run on record"); + } + return run.error == null ? SubagentResult.ok(run.result) : SubagentResult.error(run.error); + } + + @Override + protected void callCancelRequest(String sessionId, String callId) { + cancelCount++; + } + + /** Test hook: injects a run that already exists remotely, exercising reconciler reuse. */ + public void seedRun( + String sessionId, + String callId, + Object result, + @Nullable String error, + int queriesUntilComplete) { + runs.put(sessionId + "#" + callId, new Run(result, error, queriesUntilComplete)); + } + + /** Number of times the POST endpoint has been hit. */ + public int postCount() { + return postCount; + } + + /** Number of times the status endpoint has been probed. */ + public int statusQueryCount() { + return statusQueryCount; + } + + /** Number of times the result endpoint has been fetched. */ + public int fetchCount() { + return fetchCount; + } + + /** Number of times the cancel hook has been invoked. */ + public int cancelCount() { + return cancelCount; + } + + /** Exposes the pub durable call for unit-style POST and reconciler assertions. */ + public DurableCallable submitRequestForTest( + Object prompt, String sessionId, String callId) { + return submitRequest(null, sessionId, callId, prompt); + } + + /** Exposes the await durable call for unit-style assertions. */ + public DurableCallable awaitResultForTest(String sessionId, String callId) { + return awaitResult(null, sessionId, callId); + } +} From 9e142003aac5192bb2c539b5a4f3ec1489f2cb3a Mon Sep 17 00:00:00 2001 From: "luogen.lg" Date: Wed, 19 Aug 2026 23:45:57 +0800 Subject: [PATCH 11/11] [runtime][e2e] Add sub-agent integration tests --- .../subagent_external_integration_agent.py | 148 ++++++ .../subagent_external_integration_test.py | 96 ++++ .../subagent_integration_agent.py | 87 ++++ .../subagent_integration_test.py | 89 ++++ .../external/ExternalAgentClient.java | 289 +++++++++++ .../external/ExternalAgentStubService.java | 249 +++++++++ .../ExternalAgentSubagentSetupTest.java | 480 ++++++++++++++++++ .../external/ExternalAsyncSubagentSetup.java | 123 +++++ .../ExternalDeferredSubagentSetup.java | 130 +++++ 9 files changed, 1691 insertions(+) create mode 100644 python/flink_agents/e2e_tests/e2e_tests_integration/subagent_external_integration_agent.py create mode 100644 python/flink_agents/e2e_tests/e2e_tests_integration/subagent_external_integration_test.py create mode 100644 python/flink_agents/e2e_tests/e2e_tests_integration/subagent_integration_agent.py create mode 100644 python/flink_agents/e2e_tests/e2e_tests_integration/subagent_integration_test.py create mode 100644 runtime/src/test/java/org/apache/flink/agents/runtime/subagent/external/ExternalAgentClient.java create mode 100644 runtime/src/test/java/org/apache/flink/agents/runtime/subagent/external/ExternalAgentStubService.java create mode 100644 runtime/src/test/java/org/apache/flink/agents/runtime/subagent/external/ExternalAgentSubagentSetupTest.java create mode 100644 runtime/src/test/java/org/apache/flink/agents/runtime/subagent/external/ExternalAsyncSubagentSetup.java create mode 100644 runtime/src/test/java/org/apache/flink/agents/runtime/subagent/external/ExternalDeferredSubagentSetup.java 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/runtime/src/test/java/org/apache/flink/agents/runtime/subagent/external/ExternalAgentClient.java b/runtime/src/test/java/org/apache/flink/agents/runtime/subagent/external/ExternalAgentClient.java new file mode 100644 index 000000000..78ff05924 --- /dev/null +++ b/runtime/src/test/java/org/apache/flink/agents/runtime/subagent/external/ExternalAgentClient.java @@ -0,0 +1,289 @@ +/* + * 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.external; + +import com.fasterxml.jackson.databind.JsonNode; +import com.fasterxml.jackson.databind.ObjectMapper; +import com.fasterxml.jackson.databind.node.ObjectNode; +import org.apache.flink.agents.api.subagent.SubagentResult; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import javax.annotation.Nullable; + +import java.net.URI; +import java.net.http.HttpClient; +import java.net.http.HttpRequest; +import java.net.http.HttpResponse; +import java.nio.charset.StandardCharsets; +import java.time.Duration; +import java.util.UUID; + +/** + * Minimal HTTP client for the external async-task agent demo service. The service assigns its own + * {@code task_id} per submission, so the setups keep their own {@code (sessionId, callId)} to + * {@code task_id} mapping. + * + *

      Protocol contract: + * + * 

      {@code
+ * POST /tasks              -> 202 {"task_id", "status": "pending"} when a new task is created;
+ *                             200 {"task_id", "status"} when the idempotency key already exists
+ *                             body: {"prompt", "session_id", "task_id" (optional idempotency key)}
+ * GET  /tasks/{id}         -> 200 {"task_id", "status", "session_id", "created_at",
+ *                             "updated_at", "error"}; 404 when the task is unknown
+ * GET  /tasks/{id}/result  -> 200 {"task_id", "status", "result", "error"} once terminal;
+ *                             409 {"detail", "status"} while not finished;
+ *                             404 when the task is unknown
+ * GET  /health             -> 200 {"status": "ok", "llm_backend", "task_count"}
+ * }
      + * + *

      Status lifecycle: {@code pending} {@literal ->} {@code running} {@literal ->} {@code + * succeeded} | {@code failed}. {@code error} carries {@code "ExceptionType: message"} on failure. + * + *

      The {@code task_id} body field is the idempotency key: resubmitting the same key returns the + * existing task instead of creating a duplicate, so a reconciled POST after a crash-window never + * starts a second run. {@link #taskIdFor} derives that key deterministically from the {@code + * (sessionId, callId)} pair, which makes every remote task traceable across failovers without any + * client-side cache. + */ +public class ExternalAgentClient { + + private static final Logger LOG = LoggerFactory.getLogger(ExternalAgentClient.class); + + /** Lifecycle status values reported by the service. */ + public static final String PENDING = "pending"; + + public static final String RUNNING = "running"; + public static final String SUCCEEDED = "succeeded"; + public static final String FAILED = "failed"; + + private final String baseUrl; + private final HttpClient http; + private final ObjectMapper mapper = new ObjectMapper(); + + public ExternalAgentClient(String baseUrl) { + this.baseUrl = baseUrl; + // HTTP/1.1 only: the demo service (uvicorn without websockets) warns on h2c upgrades. + this.http = + HttpClient.newBuilder() + .version(HttpClient.Version.HTTP_1_1) + .connectTimeout(Duration.ofSeconds(5)) + .build(); + } + + /** True when {@code GET /health} answers 200. */ + public boolean reachable() { + try { + HttpResponse response = send(builder("/health").GET().build(), "GET /health"); + return response.statusCode() == 200; + } catch (Exception e) { + return false; + } + } + + /** The {@code llm_backend} reported by {@code GET /health}. */ + public String llmBackend() throws Exception { + JsonNode body = parse(expect(get("/health"), 200, "/health")); + return body.get("llm_backend").asText(); + } + + /** The {@code task_count} reported by {@code GET /health}. */ + public int taskCount() throws Exception { + JsonNode body = parse(expect(get("/health"), 200, "/health")); + return body.get("task_count").asInt(); + } + + /** + * A per-JVM namespace keeping consecutive test runs against a long-lived service isolated: + * without it, a second run would deterministically hit the terminal tasks of the first. + */ + private static final String RUN_NAMESPACE = UUID.randomUUID().toString(); + + /** + * The deterministic remote task id of one logical invocation: a name-based (version 3) UUID of + * {@code sessionId#callId} within {@link #RUN_NAMESPACE}, stable across in-process failovers + * and replays. + */ + public static String taskIdFor(String sessionId, String callId) { + return UUID.nameUUIDFromBytes( + (RUN_NAMESPACE + "#" + sessionId + "#" + callId) + .getBytes(StandardCharsets.UTF_8)) + .toString(); + } + + /** + * Submits a task under the given idempotency key {@code taskId}; returns the remote {@code + * task_id}. The service answers 202 for a new task and 200 when the key already exists, so a + * resubmission after a crash never creates a duplicate. + */ + public String submit(String prompt, @Nullable String sessionId, String taskId) + throws Exception { + ObjectNode requestBody = mapper.createObjectNode().put("prompt", prompt); + if (sessionId != null) { + requestBody.put("session_id", sessionId); + } + requestBody.put("task_id", taskId); + HttpRequest request = + builder("/tasks") + .header("Content-Type", "application/json") + .POST( + HttpRequest.BodyPublishers.ofString( + mapper.writeValueAsString(requestBody))) + .build(); + HttpResponse response = send(request, "POST /tasks"); + if (response.statusCode() != 202 && response.statusCode() != 200) { + throw new IllegalStateException( + "POST /tasks expected 202/200 but got " + + response.statusCode() + + ": " + + response.body()); + } + JsonNode body = parse(response); + String returnedTaskId = body.get("task_id").asText(); + LOG.info( + "submit(prompt={}, session={}, taskId={}) -> task {} ({})", + prompt, + sessionId, + taskId, + returnedTaskId, + response.statusCode() == 202 ? "created" : "idempotent replay"); + return returnedTaskId; + } + + /** + * Probes the status of one task; {@code null} when the service has no record of the id (a 404). + */ + @Nullable + public TaskStatus status(String taskId) throws Exception { + HttpRequest request = builder("/tasks/" + taskId).GET().build(); + HttpResponse response = send(request, "GET /tasks/" + taskId); + if (response.statusCode() == 404) { + return null; + } + expectStatus(response, 200, "GET /tasks/" + taskId); + JsonNode body = parse(response); + TaskStatus taskStatus = + new TaskStatus( + body.get("status").asText(), + body.hasNonNull("error") ? body.get("error").asText() : null); + LOG.info("status(task={}) -> {}", taskId, taskStatus.getStatus()); + return taskStatus; + } + + /** + * Fetches the terminal result of one task. A non-terminal task (409) surfaces as an error + * result instead of throwing. + */ + public SubagentResult fetchResult(String taskId) throws Exception { + HttpRequest request = builder("/tasks/" + taskId + "/result").GET().build(); + HttpResponse response = send(request, "GET /tasks/" + taskId + "/result"); + if (response.statusCode() == 404) { + return SubagentResult.error("task not found: " + taskId); + } + if (response.statusCode() == 409) { + return SubagentResult.error("task not finished: " + taskId); + } + expectStatus(response, 200, "GET /tasks/" + taskId + "/result"); + JsonNode body = parse(response); + String status = body.get("status").asText(); + if (FAILED.equals(status)) { + String error = body.hasNonNull("error") ? body.get("error").asText() : "run failed"; + LOG.info("fetchResult(task={}) -> failed: {}", taskId, error); + return SubagentResult.error(error); + } + JsonNode result = body.get("result"); + SubagentResult outcome = + SubagentResult.ok( + result == null || result.isNull() ? null : result.get(0).asText()); + LOG.info("fetchResult(task={}) -> {}: {}", taskId, status, outcome.getResult()); + return outcome; + } + + /** The status snapshot of one remote task. */ + public static final class TaskStatus { + private final String status; + @Nullable private final String error; + + TaskStatus(String status, @Nullable String error) { + this.status = status; + this.error = error; + } + + public String getStatus() { + return status; + } + + @Nullable + public String getError() { + return error; + } + } + + // ------------------------------------------------------------------------------------------ + // Plumbing + // ------------------------------------------------------------------------------------------ + + private HttpRequest.Builder builder(String path) { + return HttpRequest.newBuilder(uri(path)).timeout(Duration.ofSeconds(10)); + } + + private HttpResponse get(String path) throws Exception { + return send(builder(path).GET().build(), "GET " + path); + } + + private HttpResponse send(HttpRequest request, String what) throws Exception { + long started = System.currentTimeMillis(); + HttpResponse response = http.send(request, HttpResponse.BodyHandlers.ofString()); + LOG.info( + "{} -> {} ({} ms)", + what, + response.statusCode(), + System.currentTimeMillis() - started); + return response; + } + + private HttpResponse expect(HttpResponse response, int expected, String what) + throws Exception { + expectStatus(response, expected, what); + return response; + } + + private void expectStatus(HttpResponse response, int expected, String what) + throws Exception { + if (response.statusCode() != expected) { + throw new IllegalStateException( + what + + " expected " + + expected + + " but got " + + response.statusCode() + + ": " + + response.body()); + } + } + + private JsonNode parse(HttpResponse response) throws Exception { + return mapper.readTree(response.body()); + } + + private URI uri(String path) { + return URI.create(baseUrl + path); + } +} diff --git a/runtime/src/test/java/org/apache/flink/agents/runtime/subagent/external/ExternalAgentStubService.java b/runtime/src/test/java/org/apache/flink/agents/runtime/subagent/external/ExternalAgentStubService.java new file mode 100644 index 000000000..ab2b1e0cd --- /dev/null +++ b/runtime/src/test/java/org/apache/flink/agents/runtime/subagent/external/ExternalAgentStubService.java @@ -0,0 +1,249 @@ +/* + * 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.external; + +import com.fasterxml.jackson.databind.ObjectMapper; +import com.fasterxml.jackson.databind.node.ObjectNode; +import com.sun.net.httpserver.HttpExchange; +import com.sun.net.httpserver.HttpServer; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.io.IOException; +import java.io.OutputStream; +import java.net.InetSocketAddress; +import java.nio.charset.StandardCharsets; +import java.util.Map; +import java.util.concurrent.ConcurrentHashMap; + +/** + * In-process stand-in for the external async-task agent service, speaking the same HTTP protocol on + * loopback so the integration tests need no service installed anywhere. + * + *

      A run turns terminal after a configured delay measured in wall-clock time, not after a number + * of probes, which is what lets the tests assert the real polling pacing of both execution modes. A + * prompt containing {@code "fail"} produces a failed run, and any other prompt echoes the prompt + * exactly as the offline mock backend of the demo service does. + * + *

      Submissions are idempotent in the {@code task_id} the caller supplies: a repeated submission + * answers 200 with the existing run instead of starting a second one, which is what the reconcile + * paths depend on. + */ +final class ExternalAgentStubService implements AutoCloseable { + + private static final Logger LOG = LoggerFactory.getLogger(ExternalAgentStubService.class); + + /** + * The backend name reported by {@code GET /health}, matching the demo service's offline mode. + */ + static final String LLM_BACKEND = "mock"; + + /** A run stays {@code pending} for this long before it starts running. */ + private static final long PENDING_MILLIS = 100; + + private static final ObjectMapper MAPPER = new ObjectMapper(); + + private final HttpServer server; + private final Map runs = new ConcurrentHashMap<>(); + private final long taskDelayMillis; + + private ExternalAgentStubService(HttpServer server, long taskDelayMillis) { + this.server = server; + this.taskDelayMillis = taskDelayMillis; + } + + /** Starts the service on an ephemeral loopback port with the given per-run delay. */ + static ExternalAgentStubService start(long taskDelayMillis) throws IOException { + HttpServer server = HttpServer.create(new InetSocketAddress("127.0.0.1", 0), 0); + ExternalAgentStubService service = new ExternalAgentStubService(server, taskDelayMillis); + server.createContext("/health", service::handleHealth); + server.createContext("/tasks", service::handleTasks); + server.setExecutor(null); + server.start(); + LOG.info( + "external agent stub listening on {} with a {} ms task delay", + service.baseUrl(), + taskDelayMillis); + return service; + } + + String baseUrl() { + return "http://127.0.0.1:" + server.getAddress().getPort(); + } + + @Override + public void close() { + server.stop(0); + LOG.info("external agent stub stopped after serving {} runs", runs.size()); + } + + // ------------------------------------------------------------------------------------------ + // Endpoints + // ------------------------------------------------------------------------------------------ + + private void handleHealth(HttpExchange exchange) throws IOException { + ObjectNode body = + MAPPER.createObjectNode() + .put("status", "ok") + .put("llm_backend", LLM_BACKEND) + .put("task_count", runs.size()); + respond(exchange, 200, body); + } + + private void handleTasks(HttpExchange exchange) throws IOException { + String path = exchange.getRequestURI().getPath(); + if ("POST".equals(exchange.getRequestMethod())) { + handleSubmit(exchange); + } else if (path.endsWith("/result")) { + handleResult(exchange, taskIdOf(path, "/result")); + } else { + handleStatus(exchange, taskIdOf(path, "")); + } + } + + private void handleSubmit(HttpExchange exchange) throws IOException { + ObjectNode request = + (ObjectNode) + MAPPER.readTree( + new String( + exchange.getRequestBody().readAllBytes(), + StandardCharsets.UTF_8)); + String prompt = request.path("prompt").asText(""); + String sessionId = request.path("session_id").asText(null); + String taskId = request.path("task_id").asText(null); + if (taskId == null || taskId.isEmpty()) { + respond(exchange, 400, MAPPER.createObjectNode().put("detail", "task_id is required")); + return; + } + Run existing = runs.putIfAbsent(taskId, new Run(prompt, sessionId)); + ObjectNode body = + MAPPER.createObjectNode() + .put("task_id", taskId) + .put("status", existing == null ? "pending" : existing.status(now())); + respond(exchange, existing == null ? 202 : 200, body); + } + + private void handleStatus(HttpExchange exchange, String taskId) throws IOException { + Run run = runs.get(taskId); + if (run == null) { + respond(exchange, 404, MAPPER.createObjectNode().put("detail", "unknown task")); + return; + } + String status = run.status(now()); + ObjectNode body = + MAPPER.createObjectNode() + .put("task_id", taskId) + .put("status", status) + .put("session_id", run.sessionId) + .put("created_at", run.createdAt) + .put("updated_at", now()); + if (ExternalAgentClient.FAILED.equals(status)) { + body.put("error", run.error()); + } else { + body.putNull("error"); + } + respond(exchange, 200, body); + } + + private void handleResult(HttpExchange exchange, String taskId) throws IOException { + Run run = runs.get(taskId); + if (run == null) { + respond(exchange, 404, MAPPER.createObjectNode().put("detail", "unknown task")); + return; + } + String status = run.status(now()); + if (!ExternalAgentClient.SUCCEEDED.equals(status) + && !ExternalAgentClient.FAILED.equals(status)) { + respond( + exchange, + 409, + MAPPER.createObjectNode().put("detail", "not finished").put("status", status)); + return; + } + ObjectNode body = MAPPER.createObjectNode().put("task_id", taskId).put("status", status); + if (ExternalAgentClient.FAILED.equals(status)) { + body.putNull("result").put("error", run.error()); + } else { + // The service reports the answer as the message list of the finished run. + body.putNull("error"); + body.putArray("result").add(run.answer()); + } + respond(exchange, 200, body); + } + + // ------------------------------------------------------------------------------------------ + // Plumbing + // ------------------------------------------------------------------------------------------ + + private static String taskIdOf(String path, String suffix) { + String trimmed = path.substring("/tasks/".length()); + return suffix.isEmpty() + ? trimmed + : trimmed.substring(0, trimmed.length() - suffix.length()); + } + + private static long now() { + return System.currentTimeMillis(); + } + + private static void respond(HttpExchange exchange, int status, ObjectNode body) + throws IOException { + byte[] payload = MAPPER.writeValueAsBytes(body); + exchange.getResponseHeaders().add("Content-Type", "application/json"); + exchange.sendResponseHeaders(status, payload.length); + try (OutputStream out = exchange.getResponseBody()) { + out.write(payload); + } + } + + /** One submitted run, whose state is a pure function of the time since its submission. */ + private final class Run { + private final String prompt; + private final String sessionId; + private final long createdAt = now(); + + private Run(String prompt, String sessionId) { + this.prompt = prompt; + this.sessionId = sessionId; + } + + private String status(long now) { + long elapsed = now - createdAt; + if (elapsed < PENDING_MILLIS) { + return ExternalAgentClient.PENDING; + } + if (elapsed < taskDelayMillis) { + return ExternalAgentClient.RUNNING; + } + return failing() ? ExternalAgentClient.FAILED : ExternalAgentClient.SUCCEEDED; + } + + private boolean failing() { + return prompt.toLowerCase().contains("fail"); + } + + private String answer() { + return "[offline-mock] echo: " + prompt; + } + + private String error() { + return "RuntimeError: mock agent failed on demand"; + } + } +} diff --git a/runtime/src/test/java/org/apache/flink/agents/runtime/subagent/external/ExternalAgentSubagentSetupTest.java b/runtime/src/test/java/org/apache/flink/agents/runtime/subagent/external/ExternalAgentSubagentSetupTest.java new file mode 100644 index 000000000..e42b9a2f1 --- /dev/null +++ b/runtime/src/test/java/org/apache/flink/agents/runtime/subagent/external/ExternalAgentSubagentSetupTest.java @@ -0,0 +1,480 @@ +/* + * 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.external; + +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.RunnerContext; +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.subagent.BaseAsyncSubagentSetup; +import org.apache.flink.agents.runtime.subagent.BaseDeferredSubagentSetup; +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.AfterAll; +import org.junit.jupiter.api.Assumptions; +import org.junit.jupiter.api.BeforeAll; +import org.junit.jupiter.api.Test; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import javax.annotation.Nullable; + +import java.util.List; +import java.util.UUID; + +import static org.assertj.core.api.Assertions.assertThat; + +/** + * Exercises both sub-agent execution modes against an external async-task agent service over real + * HTTP: {@link BaseAsyncSubagentSetup} (pub/sub: durable POST, then probe and fetch) and {@link + * BaseDeferredSubagentSetup} (deferred: the request is issued when the handle resolves). Each mode + * runs a success prompt and a failure prompt — the prompt describes the expected behavior, and the + * mock backend fails on demand when it contains "fail". Both modes start probing right after the + * submit and are asserted to check every call at least {@link #MIN_EXPECTED_CHECKS} times before it + * turns terminal, which is what the probe interval is derived from. + * + *

      By default the test starts {@link ExternalAgentStubService} on a loopback port, so it runs + * anywhere with no service installed. Point {@code -Dexternal.agent.url=...} at a running service + * to exercise the same suite against the real demo deployment instead; the suite is then skipped if + * that endpoint is unreachable. + */ +public class ExternalAgentSubagentSetupTest { + + private static final Logger LOG = LoggerFactory.getLogger(ExternalAgentSubagentSetupTest.class); + + /** Set to run against a service of your own instead of the in-process stub. */ + @Nullable private static final String EXTERNAL_URL = System.getProperty("external.agent.url"); + + private static final String RESOURCE_NAME = "ext-agent"; + + /** Prompt describing the expected successful behavior; the mock backend echoes it. */ + private static final String SUCCESS_PROMPT = "echo the greeting; expect a success"; + + /** Prompt describing the expected failing behavior; the mock backend fails on "fail". */ + private static final String FAILURE_PROMPT = "please fail this run"; + + /** + * Every call must be checked at least this many times before reaching a terminal state, which + * is what makes the polling of both modes observable rather than incidental. + */ + private static final int MIN_EXPECTED_CHECKS = 5; + + /** How long a run of the stub service takes; the real demo service uses about five seconds. */ + private static final long STUB_TASK_DELAY_MILLIS = 1_500; + + /** The real demo service's task delay, used to pace the probes when running against it. */ + private static final long EXTERNAL_TASK_DELAY_MILLIS = 5_000; + + @Nullable private static ExternalAgentStubService stub; + + private static String baseUrl; + + /** + * The probe interval, derived from the service's task delay so that a run is always checked + * more than {@link #MIN_EXPECTED_CHECKS} times. + */ + private static long probeIntervalMillis; + + /** The backend reported by the service; the mock backend answers deterministically. */ + private static String llmBackend; + + @BeforeAll + static void startService() throws Exception { + long taskDelayMillis; + if (EXTERNAL_URL == null) { + stub = ExternalAgentStubService.start(STUB_TASK_DELAY_MILLIS); + baseUrl = stub.baseUrl(); + taskDelayMillis = STUB_TASK_DELAY_MILLIS; + } else { + baseUrl = EXTERNAL_URL; + taskDelayMillis = EXTERNAL_TASK_DELAY_MILLIS; + } + probeIntervalMillis = taskDelayMillis / (MIN_EXPECTED_CHECKS + 1); + ExternalAgentClient client = new ExternalAgentClient(baseUrl); + Assumptions.assumeTrue( + client.reachable(), "external agent service not reachable at " + baseUrl); + llmBackend = client.llmBackend(); + LOG.info( + "external agent service at {} is up, llm_backend={}, probing every {} ms", + baseUrl, + llmBackend, + probeIntervalMillis); + } + + @AfterAll + static void stopService() { + if (stub != null) { + stub.close(); + stub = null; + } + } + + // ------------------------------------------------------------------------------------------ + // Mode 1: async pub/sub through BaseAsyncSubagentSetup + // ------------------------------------------------------------------------------------------ + + /** Submits the success prompt through the short form and awaits the run's outcome. */ + @SuppressWarnings("unused") + public static void asyncSubmitAndAwait(Event event, RunnerContext ctx) throws Exception { + BaseAsyncSubagentSetup setup = + (BaseAsyncSubagentSetup) ctx.getResource(RESOURCE_NAME, ResourceType.AGENT); + SubagentResult result = setup.submit(ctx, SUCCESS_PROMPT).await(); + ctx.sendEvent(new OutputEvent(result.isSuccess() + "|" + result.getResult())); + } + + @Test + void asyncModeSucceedsAndChecksAtLeastFiveTimes() throws Exception { + LOG.info("[test] async mode, success prompt: {}", SUCCESS_PROMPT); + ExternalAsyncSubagentSetup setup = asyncSetup(); + long started = System.currentTimeMillis(); + try (KeyedOneInputStreamOperatorTestHarness harness = + harness(plan("asyncSubmitAndAwait", setup))) { + harness.open(); + run(harness, 1L); + + assertThat(harness.getRecordOutput()) + .singleElement() + .extracting(StreamRecord::getValue) + .satisfies( + value -> assertSuccessfulEcho(String.valueOf(value), SUCCESS_PROMPT)); + // The probes are spread over the run rather than issued back-to-back. + List probes = setup.probeTimestamps(); + assertThat(probes).hasSizeGreaterThanOrEqualTo(MIN_EXPECTED_CHECKS); + for (int i = 1; i < probes.size(); i++) { + assertThat(probes.get(i) - probes.get(i - 1)) + .describedAs("gap between check #%d and #%d", i, i + 1) + .isGreaterThanOrEqualTo(probeIntervalMillis / 2); + } + LOG.info( + "[test] async success done in {} ms, checks={}", + System.currentTimeMillis() - started, + probes.size()); + } + } + + /** Submits the failure prompt; the remote run fails and the error surfaces in the Result. */ + @SuppressWarnings("unused") + public static void asyncSubmitFailingPrompt(Event event, RunnerContext ctx) throws Exception { + BaseAsyncSubagentSetup setup = + (BaseAsyncSubagentSetup) ctx.getResource(RESOURCE_NAME, ResourceType.AGENT); + SubagentResult result = setup.submit(ctx, FAILURE_PROMPT).await(); + ctx.sendEvent(new OutputEvent(result.isSuccess() + "|" + result.getErrorMessage())); + } + + @Test + void asyncModeSurfacesARemoteFailureWithoutFetching() throws Exception { + LOG.info("[test] async mode, failure prompt: {}", FAILURE_PROMPT); + ExternalAsyncSubagentSetup setup = asyncSetup(); + long started = System.currentTimeMillis(); + try (KeyedOneInputStreamOperatorTestHarness harness = + harness(plan("asyncSubmitFailingPrompt", setup))) { + harness.open(); + run(harness, 1L); + + assertThat(harness.getRecordOutput()) + .singleElement() + .extracting(StreamRecord::getValue) + .satisfies( + value -> { + String output = String.valueOf(value); + assertThat(output).startsWith("false|"); + assertThat(output).contains("mock agent failed on demand"); + }); + // The failing run is paced the same way, so it is checked as often. + assertThat(setup.probeTimestamps()).hasSizeGreaterThanOrEqualTo(MIN_EXPECTED_CHECKS); + LOG.info("[test] async failure done in {} ms", System.currentTimeMillis() - started); + } + } + + /** Batches one success and one failure handle and resolves them together. */ + @SuppressWarnings("unused") + public static void asyncBatchSuccessAndFailure(Event event, RunnerContext ctx) + throws Exception { + BaseAsyncSubagentSetup setup = + (BaseAsyncSubagentSetup) ctx.getResource(RESOURCE_NAME, ResourceType.AGENT); + SubagentFuture success = setup.submit(ctx, SUCCESS_PROMPT); + SubagentFuture failure = setup.submit(ctx, FAILURE_PROMPT); + List results = success.combine(failure).awaitAll(); + ctx.sendEvent( + new OutputEvent( + results.get(0).isSuccess() + + "|" + + results.get(0).getResult() + + "||" + + results.get(1).isSuccess() + + "|" + + results.get(1).getErrorMessage())); + } + + @Test + void asyncModeCombineResolvesSuccessAndFailureTogether() throws Exception { + LOG.info("[test] async mode, batched success + failure prompts"); + long started = System.currentTimeMillis(); + try (KeyedOneInputStreamOperatorTestHarness harness = + harness(plan("asyncBatchSuccessAndFailure", asyncSetup()))) { + harness.open(); + run(harness, 1L); + + assertThat(harness.getRecordOutput()) + .singleElement() + .extracting(StreamRecord::getValue) + .satisfies( + value -> { + String[] halves = String.valueOf(value).split("\\|\\|"); + assertSuccessfulEcho(halves[0], SUCCESS_PROMPT); + assertThat(halves[1]).startsWith("false|"); + assertThat(halves[1]).contains("mock agent failed on demand"); + }); + LOG.info("[test] async batch done in {} ms", System.currentTimeMillis() - started); + } + } + + // ------------------------------------------------------------------------------------------ + // Mode 2: deferred execution through BaseDeferredSubagentSetup + // ------------------------------------------------------------------------------------------ + + /** Resolves one deferred handle with the success prompt, issued at resolve time. */ + @SuppressWarnings("unused") + public static void deferredSubmitAndAwait(Event event, RunnerContext ctx) throws Exception { + BaseDeferredSubagentSetup setup = + (BaseDeferredSubagentSetup) ctx.getResource(RESOURCE_NAME, ResourceType.AGENT); + SubagentResult result = + setup.submit(ctx, SUCCESS_PROMPT, "session", "call-success").await(); + ctx.sendEvent(new OutputEvent(result.isSuccess() + "|" + result.getResult())); + } + + @Test + void deferredModeSucceedsAndChecksAtLeastFiveTimes() throws Exception { + LOG.info("[test] deferred mode, success prompt: {}", SUCCESS_PROMPT); + ExternalDeferredSubagentSetup setup = deferredSetup(); + long started = System.currentTimeMillis(); + try (KeyedOneInputStreamOperatorTestHarness harness = + harness(plan("deferredSubmitAndAwait", setup))) { + harness.open(); + run(harness, 1L); + + assertThat(harness.getRecordOutput()) + .singleElement() + .extracting(StreamRecord::getValue) + .satisfies( + value -> assertSuccessfulEcho(String.valueOf(value), SUCCESS_PROMPT)); + // The polls are spread over the run, so it is checked at least as often. + assertThat(setup.pollCount()).isGreaterThanOrEqualTo(MIN_EXPECTED_CHECKS); + LOG.info( + "[test] deferred success done in {} ms, checks={}", + System.currentTimeMillis() - started, + setup.pollCount()); + } + } + + /** Resolves one deferred handle with the failure prompt. */ + @SuppressWarnings("unused") + public static void deferredSubmitFailingPrompt(Event event, RunnerContext ctx) + throws Exception { + BaseDeferredSubagentSetup setup = + (BaseDeferredSubagentSetup) ctx.getResource(RESOURCE_NAME, ResourceType.AGENT); + SubagentResult result = + setup.submit(ctx, FAILURE_PROMPT, "session", "call-failure").await(); + ctx.sendEvent(new OutputEvent(result.isSuccess() + "|" + result.getErrorMessage())); + } + + @Test + void deferredModeSurfacesARemoteFailure() throws Exception { + LOG.info("[test] deferred mode, failure prompt: {}", FAILURE_PROMPT); + ExternalDeferredSubagentSetup setup = deferredSetup(); + long started = System.currentTimeMillis(); + try (KeyedOneInputStreamOperatorTestHarness harness = + harness(plan("deferredSubmitFailingPrompt", setup))) { + harness.open(); + run(harness, 1L); + + assertThat(harness.getRecordOutput()) + .singleElement() + .extracting(StreamRecord::getValue) + .satisfies( + value -> { + String output = String.valueOf(value); + assertThat(output).startsWith("false|"); + assertThat(output).contains("mock agent failed on demand"); + }); + // The failing run is paced the same way, so it is checked as often. + assertThat(setup.pollCount()).isGreaterThanOrEqualTo(MIN_EXPECTED_CHECKS); + LOG.info("[test] deferred failure done in {} ms", System.currentTimeMillis() - started); + } + } + + // ------------------------------------------------------------------------------------------ + // Mode 3: reconciliation across the crash window (idempotent resubmission) + // ------------------------------------------------------------------------------------------ + + /** + * Simulates a failover after the POST landed but before its durable record was persisted: a + * fresh setup instance, with no state from the crashed process, must find the original remote + * task through the deterministic id, and its reconciler must not start a duplicate. + */ + @Test + void reconcileAfterCrashResumesTheOriginalTask() throws Exception { + String sessionId = "reco-crash-" + UUID.randomUUID(); + String callId = "call-1"; + LOG.info("[test] reconcile after crash, session={}", sessionId); + ExternalAgentClient client = new ExternalAgentClient(baseUrl); + int tasksBefore = client.taskCount(); + + ExternalAsyncSubagentSetup original = asyncSetup(); + original.callSubmitRequest(sessionId, callId, SUCCESS_PROMPT); + assertThat(client.taskCount()).isEqualTo(tasksBefore + 1); + + // Failover: a brand-new setup instance with nothing in memory. + ExternalAsyncSubagentSetup recovered = asyncSetup(); + assertThat(recovered.callQueryStatus(sessionId, callId).getState()) + .describedAs("the recovered setup must find the original in-flight task") + .isEqualTo(BaseAsyncSubagentSetup.RunStatus.State.RUNNING); + + // The reconciler sees RUNNING and repairs nothing; no duplicate task appears. + recovered.reconcileForTest(sessionId, callId, SUCCESS_PROMPT); + assertThat(client.taskCount()).isEqualTo(tasksBefore + 1); + + // The recovered setup tracks the original task to its terminal result. + BaseAsyncSubagentSetup.RunStatus probe = waitForTerminal(recovered, sessionId, callId); + assertThat(probe.getState()).isEqualTo(BaseAsyncSubagentSetup.RunStatus.State.COMPLETED); + SubagentResult result = recovered.callFetchResult(sessionId, callId); + assertThat(result.isSuccess()).isTrue(); + assertSuccessfulEcho("true|" + result.getResult(), SUCCESS_PROMPT); + assertThat(client.taskCount()).isEqualTo(tasksBefore + 1); + LOG.info("[test] reconcile after crash done, original task resumed without duplication"); + } + + /** + * Simulates a crash where the POST never landed: the reconciler's probe sees NOT_STARTED and + * starts the run; a second reconciliation finds it RUNNING and adds no duplicate. + */ + @Test + void reconcileRepostsWhenThePostNeverLanded() throws Exception { + String sessionId = "reco-lost-" + UUID.randomUUID(); + String callId = "call-1"; + LOG.info("[test] reconcile lost POST, session={}", sessionId); + ExternalAgentClient client = new ExternalAgentClient(baseUrl); + int tasksBefore = client.taskCount(); + + ExternalAsyncSubagentSetup setup = asyncSetup(); + assertThat(setup.callQueryStatus(sessionId, callId).getState()) + .isEqualTo(BaseAsyncSubagentSetup.RunStatus.State.NOT_STARTED); + + // First reconciliation: nothing on the service, so the POST is (re)sent. + setup.reconcileForTest(sessionId, callId, SUCCESS_PROMPT); + assertThat(client.taskCount()).isEqualTo(tasksBefore + 1); + assertThat(setup.callQueryStatus(sessionId, callId).getState()) + .isEqualTo(BaseAsyncSubagentSetup.RunStatus.State.RUNNING); + + // Second reconciliation: the run exists now; nothing is repaired or duplicated. + setup.reconcileForTest(sessionId, callId, SUCCESS_PROMPT); + assertThat(client.taskCount()).isEqualTo(tasksBefore + 1); + LOG.info("[test] reconcile lost POST done, task started exactly once"); + } + + /** Probes until the run reaches a terminal state; fails after 15 seconds. */ + private static BaseAsyncSubagentSetup.RunStatus waitForTerminal( + ExternalAsyncSubagentSetup setup, String sessionId, String callId) throws Exception { + long deadline = System.currentTimeMillis() + 15_000; + while (true) { + BaseAsyncSubagentSetup.RunStatus probe = setup.callQueryStatus(sessionId, callId); + switch (probe.getState()) { + case COMPLETED: + case FAILED: + return probe; + default: + if (System.currentTimeMillis() > deadline) { + throw new AssertionError( + "task of " + + sessionId + + "#" + + callId + + " never reached a" + + " terminal state"); + } + Thread.sleep(500); + } + } + } + + // ------------------------------------------------------------------------------------------ + // Assertions and harness plumbing + // ------------------------------------------------------------------------------------------ + + /** + * Asserts a {@code true|} output. The mock backend echoes deterministically, so its + * answer is asserted exactly; a real LLM backend only needs a non-blank answer. + */ + private static void assertSuccessfulEcho(String output, String prompt) { + assertThat(output).startsWith("true|"); + String answer = output.substring("true|".length()); + if ("mock".equals(llmBackend)) { + assertThat(answer).isEqualTo("[offline-mock] echo: " + prompt); + } else { + assertThat(answer).isNotBlank(); + } + } + + @SuppressWarnings("unchecked") + private static void run( + KeyedOneInputStreamOperatorTestHarness harness, long value) + throws Exception { + harness.processElement(new StreamRecord<>(value)); + ((ActionExecutionOperator) harness.getOperator()) + .waitInFlightEventsFinished(); + } + + /** An async-mode setup against the running service, paced to the derived probe interval. */ + private static ExternalAsyncSubagentSetup asyncSetup() { + return new ExternalAsyncSubagentSetup(baseUrl, probeIntervalMillis); + } + + /** A deferred-mode setup against the running service, paced to the derived probe interval. */ + private static ExternalDeferredSubagentSetup deferredSetup() { + return new ExternalDeferredSubagentSetup(baseUrl, probeIntervalMillis); + } + + private static AgentPlan plan(String actionMethod, Object setup) throws Exception { + Agent agent = new Agent(); + agent.addResource(RESOURCE_NAME, ResourceType.AGENT, setup); + agent.addAction( + new String[] {InputEvent.EVENT_TYPE}, + ExternalAgentSubagentSetupTest.class.getMethod( + actionMethod, Event.class, RunnerContext.class)); + return new AgentPlan(agent, new AgentConfiguration()); + } + + private static KeyedOneInputStreamOperatorTestHarness harness( + AgentPlan plan) throws Exception { + return new KeyedOneInputStreamOperatorTestHarness<>( + new ActionExecutionOperatorFactory<>(plan, true), + (KeySelector) value -> value, + TypeExtractor.getForClass(Long.class)); + } +} diff --git a/runtime/src/test/java/org/apache/flink/agents/runtime/subagent/external/ExternalAsyncSubagentSetup.java b/runtime/src/test/java/org/apache/flink/agents/runtime/subagent/external/ExternalAsyncSubagentSetup.java new file mode 100644 index 000000000..29c25d7dc --- /dev/null +++ b/runtime/src/test/java/org/apache/flink/agents/runtime/subagent/external/ExternalAsyncSubagentSetup.java @@ -0,0 +1,123 @@ +/* + * 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.external; + +import org.apache.flink.agents.api.subagent.SubagentResult; +import org.apache.flink.agents.runtime.subagent.BaseAsyncSubagentSetup; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import javax.annotation.Nullable; + +import java.util.ArrayList; +import java.util.Collections; +import java.util.List; + +/** + * Integration of {@link BaseAsyncSubagentSetup} against the external async-task agent demo service: + * the pub is {@code POST /tasks}, the sub probes {@code GET /tasks/{id}} and fetches {@code GET + * /tasks/{id}/result}. The remote task id is derived deterministically from the {@code (sessionId, + * callId)} pair through {@link ExternalAgentClient#taskIdFor} and sent as the idempotency key of + * the POST, so the setup keeps no local state: after a failover the probe finds the original task + * again, and a reconciled resubmission never creates a duplicate. A probe hitting an unknown id + * reports {@link RunStatus#notStarted()}, letting the base's reconciler re-post the submission. The + * service offers no cancel endpoint, so the cancel hook stays the default no-op. + * + *

      Pacing: the await probes the remote status right after the submission, at the interval the + * setup is created with, which is chosen so that a run of the service's task delay is checked + * several times before it finishes. Every probe is logged and recorded for test assertions. + */ +public class ExternalAsyncSubagentSetup extends BaseAsyncSubagentSetup { + + private static final Logger LOG = LoggerFactory.getLogger(ExternalAsyncSubagentSetup.class); + + private final String baseUrl; + private final List probeTimestamps = Collections.synchronizedList(new ArrayList<>()); + + @Nullable private transient ExternalAgentClient client; + + public ExternalAsyncSubagentSetup(String baseUrl, long probeIntervalMillis) { + this.baseUrl = baseUrl; + this.statusPollIntervalMillis = probeIntervalMillis; + } + + private ExternalAgentClient client() { + if (client == null) { + client = new ExternalAgentClient(baseUrl); + } + return client; + } + + @Override + protected void callSubmitRequest(String sessionId, String callId, Object prompt) + throws Exception { + LOG.info("[async] submit {}#{} prompt={}", sessionId, callId, prompt); + client().submit(String.valueOf(prompt), sessionId, taskId(sessionId, callId)); + } + + @Override + protected RunStatus callQueryStatus(String sessionId, String callId) { + probeTimestamps.add(System.currentTimeMillis()); + ExternalAgentClient.TaskStatus probe; + try { + probe = client().status(taskId(sessionId, callId)); + } catch (RuntimeException e) { + throw e; + } catch (Exception e) { + // A broken probe is a system-level failure: propagate and let the job fail over. + throw new RuntimeException("status probe failed for " + sessionId + "#" + callId, e); + } + if (probe == null) { + // No task under the deterministic id: the POST never landed. + return RunStatus.notStarted(); + } + switch (probe.getStatus()) { + case ExternalAgentClient.SUCCEEDED: + return RunStatus.completed(); + case ExternalAgentClient.FAILED: + return RunStatus.failed( + probe.getError() == null ? "remote task failed" : probe.getError()); + default: + // pending or running + return RunStatus.running(); + } + } + + @Override + protected SubagentResult callFetchResult(String sessionId, String callId) throws Exception { + LOG.info("[async] fetch {}#{}", sessionId, callId); + return client().fetchResult(taskId(sessionId, callId)); + } + + private static String taskId(String sessionId, String callId) { + return ExternalAgentClient.taskIdFor(sessionId, callId); + } + + /** Test-facing entry into the crash-window reconciliation of the durable POST. */ + void reconcileForTest(String sessionId, String callId, Object prompt) throws Exception { + reconcileSubmitRequest(sessionId, callId, prompt); + } + + /** Timestamps of every status probe, for asserting the probe pacing in tests. */ + public List probeTimestamps() { + synchronized (probeTimestamps) { + return new ArrayList<>(probeTimestamps); + } + } +} diff --git a/runtime/src/test/java/org/apache/flink/agents/runtime/subagent/external/ExternalDeferredSubagentSetup.java b/runtime/src/test/java/org/apache/flink/agents/runtime/subagent/external/ExternalDeferredSubagentSetup.java new file mode 100644 index 000000000..439b86a6a --- /dev/null +++ b/runtime/src/test/java/org/apache/flink/agents/runtime/subagent/external/ExternalDeferredSubagentSetup.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.subagent.external; + +import org.apache.flink.agents.api.context.DurableCallable; +import org.apache.flink.agents.api.context.RunnerContext; +import org.apache.flink.agents.api.subagent.SubagentResult; +import org.apache.flink.agents.runtime.subagent.BaseDeferredSubagentSetup; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import javax.annotation.Nullable; + +import java.util.concurrent.atomic.AtomicInteger; + +/** + * Integration of {@link BaseDeferredSubagentSetup} against the external async-task agent demo + * service: each deferred handle, when resolved, runs one self-contained invocation — submit the + * task under the deterministic idempotency key derived from {@code (sessionId, callId)} through + * {@link ExternalAgentClient#taskIdFor}, then poll the status until a terminal state, and fetch the + * result. The poll interval is chosen so that a run of the service's task delay is checked several + * times before it finishes. The durable id is derived solely from the {@code (sessionId, callId)} + * pair, so a replay after failover re-issues the same logical invocation and the idempotent submit + * reuses the original remote task. + */ +public class ExternalDeferredSubagentSetup extends BaseDeferredSubagentSetup { + + private static final Logger LOG = LoggerFactory.getLogger(ExternalDeferredSubagentSetup.class); + + private final String baseUrl; + private final long pollIntervalMillis; + private final AtomicInteger pollCount = new AtomicInteger(); + + @Nullable private transient ExternalAgentClient client; + + public ExternalDeferredSubagentSetup(String baseUrl, long pollIntervalMillis) { + this.baseUrl = baseUrl; + this.pollIntervalMillis = pollIntervalMillis; + } + + private ExternalAgentClient client() { + if (client == null) { + client = new ExternalAgentClient(baseUrl); + } + return client; + } + + @Override + protected DurableCallable prepare( + RunnerContext ctx, Object prompt, String sessionId, String callId) { + return new DurableCallable() { + @Override + public String getId() { + return sessionId + "#" + callId; + } + + @Override + public Class getResultClass() { + return SubagentResult.class; + } + + @Override + public SubagentResult call() { + try { + LOG.info("[deferred] resolve {}#{} prompt={}", sessionId, callId, prompt); + String taskId = + client().submit( + String.valueOf(prompt), + sessionId, + ExternalAgentClient.taskIdFor(sessionId, callId)); + LOG.info( + "[deferred] {}#{} submitted as {}; polling once per second", + sessionId, + callId, + taskId); + while (true) { + int check = pollCount.incrementAndGet(); + ExternalAgentClient.TaskStatus probe = client().status(taskId); + LOG.info( + "[deferred] {}#{} check #{} -> {}", + sessionId, + callId, + check, + probe == null ? "404" : probe.getStatus()); + if (probe == null) { + return SubagentResult.error("remote task disappeared: " + taskId); + } + String status = probe.getStatus(); + if (ExternalAgentClient.SUCCEEDED.equals(status) + || ExternalAgentClient.FAILED.equals(status)) { + SubagentResult result = client().fetchResult(taskId); + LOG.info( + "[deferred] {}#{} finished: success={}, result={}, error={}", + sessionId, + callId, + result.isSuccess(), + result.getResult(), + result.getErrorMessage()); + return result; + } + Thread.sleep(pollIntervalMillis); + } + } catch (Exception e) { + return SubagentResult.error(e); + } + } + }; + } + + /** Total number of status polls across all invocations, for test pacing assertions. */ + public int pollCount() { + return pollCount.get(); + } +}