[wellkilo]

Linked issue: #939

Closes #939

Purpose of change

Flink Agents workflows now complete every keyed input when Flink uses its batch keyed-state backend, instead of rejecting explicit batch mode or silently retaining only the last key's output. Bounded AUTOMATIC jobs receive the same behavior when Flink selects that backend.

Runtime flow

Both Java and Python entry points still converge on CompileUtils.connectToAgent. The previous graph-construction rejection is removed. During ActionExecutionOperator.open, the operator records whether Flink created a BatchExecutionKeyedStateBackend. processElement continues to enqueue action work through the mailbox; before it returns under that backend, it yields until the current input's complete action chain has removed its processing key. General-purpose keyed-state backends skip this drain and preserve cross-key concurrency.

Key decisions

• Detect the backend Flink actually instantiated rather than infer execution mode from graph-construction configuration. This covers explicit BATCH and bounded AUTOMATIC jobs.
• Reuse the existing mailbox continuation and waitInFlightEventsFinished path instead of introducing a second synchronous action executor or moving all keyed runtime state into operator-owned maps.
• Serialize inputs only under the batch backend. Its state contract already requires a key to be fully processed before moving to another key; streaming backends retain their existing concurrent behavior.

Behavioral Semantics

Interaction decisions

Runtime backend	Action form	Behavior
Batch keyed-state backend	Synchronous chain	The chain finishes before processElement returns
Batch keyed-state backend	Async continuation	Mailbox yielding continues until the action chain finishes
General-purpose keyed-state backend	Sync or async	Existing mailbox scheduling and cross-key concurrency are unchanged

Behavioral contracts

1. Every input key in a bounded, single-subtask batch job produces its workflow output.
2. Async action continuations finish without losing keyed state when the batch backend is active.
3. AUTOMATIC mode behaves correctly when Flink resolves bounded input to batch execution.
4. Streaming/general-purpose state backends retain cross-key concurrent scheduling.

Failure behavior

Mailbox continuation failures still propagate through the operator's existing action-task failure path and fail the task. The batch drain does not retry, fall back, or absorb failures. No new configuration or external-service failure path is introduced.

Tests

Contract	Tests
1	CompileUtilsTest.processesEveryKeyWithDefaultBatchStateBackend; FlinkIntegrationTest.testBatchExecutionWithMultipleKeys
2	CompileUtilsTest.processesAsyncActionsWithDefaultBatchStateBackend
3	CompileUtilsTest.processesEveryKeyWhenAutomaticModeSelectsBatchExecution
4	ActionExecutionOperatorTest.testDifferentKeyDataCanRunConcurrently; testExecuteAsyncWithMultipleKeys

Verification:

• Full Java unit-test reactor passed; runtime: 1,067 tests, 0 failures, 1 skipped.
• FlinkIntegrationTest.testBatchExecutionWithMultipleKeys passed on Flink 1.20.5 and 2.3.0 locally; CI runs the same test on every supported Flink line from 1.20 through 2.3.
• Full 36-module Spotless check passed.
• Apache RAT license check passed.

Not verified locally: the E2E test was not individually executed against Flink 2.0, 2.1, and 2.2; their dist builds succeeded, and the repository's existing Java IT matrix runs the test for each version.

Implementation invariants and supporting evidence

• The unfixed runtime was reproduced after removing the temporary guard: five distinct keys produced only [10] instead of [2, 4, 6, 8, 10].
• BatchExecutionKeyedStateBackend is present under the same fully qualified name in the supported Flink 1.20.5, 2.0.2, 2.1.3, 2.2.1, and 2.3.0 artifacts.
• No new dependency, persisted state, serializer, checkpoint format, event schema, or public method is introduced.

API

No public signature, JSON, event, or IDL changes. Existing callers using explicit batch mode no longer receive the temporary configuration-time IllegalStateException; the previously unsupported configuration now executes correctly. Java and Python bridge calls share the same runtime operator behavior.

Documentation

• doc-needed
• doc-not-needed
• doc-included

The runtime method and API contracts are recorded in runtime/Method.md and runtime/API.md.

Was this patch authored or co-authored using generative AI tooling?

• Yes
• No

Generated-by: TraeCode (GPT-5)

[wenjin272]

Thanks for the contribution! I’m currently focused on preparing for the 0.4 release, so it may take me a while to review this PR. Thanks for your patience, and other reviewers are welcome to take a look in the meantime.

[sangkyoonnam]

After rebasing onto main, the batch drain needs to move outside the parallel engine's lock. Where it sits now, it fails the task under the batch backend whenever that engine is picked. #1147 merged the day after this PR opened, so the interaction didn't exist at your base. The rebase also hits one import conflict in FlinkIntegrationTest.

What I ran, Corretto 17.0.14:

• PR head: CompileUtilsTest passes 12/12. With only ActionExecutionOperator.java reverted to its parent, the three new tests fail with [10] instead of [2, 4, 6, 8, 10] and [60] instead of [20, 40, 60].
• PR head merged with upstream/main (d6e2ecc9): the same three tests error with IllegalStateException: ParallelExecutionLock is non-reentrant; the owner cannot acquire it again. JDK 21 gives the same result in unit tests. I think that's because surefire runs from target/classes, where the META-INF/versions/21 continuation classes aren't picked up, so these plans still get the parallel engine. I haven't tested the packaged JDK 21 continuation path.

This comes from the drain meeting #1147's lock wrapper. For these Java-only test plans, the parallel engine is selected by default when continuation support isn't available. On that engine processElement runs processElementInternal inside runWithParallelExecutionLock. Both waitForCurrentInputInBatchMode() calls are in processElementInternal, and waitInFlightEventsFinished() now takes the lock itself through acquireByMain(). Making only the lock reentrant wouldn't fix this. Yielding while holding it would block the workers, which need it to run actions.

On the rebased code, this change fixes it in my runs:

@@ -468,6 +468,7 @@
         } else {
             processElementInternal(record);
         }
+        waitForCurrentInputInBatchMode();
     }
 
     private void processElementInternal(StreamRecord<IN> record) throws Exception {
@@ -501,4 +502,3 @@
                 enqueuePendingInputEvent(inputEvent);
-                waitForCurrentInputInBatchMode();
                 return;
             }
@@ -512,3 +512,2 @@
         processInputEvent(key, inputEvent);
-        waitForCurrentInputInBatchMode();
     }

With that on the merged branch, ActionExecutionOperator*Test plus CompileUtilsTest pass, 89 tests, 0 failures. I also ran the synchronous five-key batch case with parallel-execution.enabled=false as a local-only test. It passes before and after the move. I haven't run the e2e FlinkIntegrationTest or the Python path.

[sangkyoonnam]

In the rebased code, this call and the one at L311 run while processElement holds ParallelExecutionLock. waitInFlightEventsFinished() then tries to acquire it again. I'd remove both and add one call after the if/else in processElement, as in the summary. The drain still finishes before the next record can switch keys.

Could you add a serial variant of the five-key batch test with AgentExecutionOptions.PARALLEL_EXECUTION_ENABLED=false in the plan config? After rebasing, I'd also pin one batch test to the parallel engine: set the option to true and add assumeFalse(ContinuationActionExecutor.isContinuationSupported()), like the parallel tests in ActionExecutionOperatorTest. The current batch tests rely on the default engine selection.

[sangkyoonnam]

Optional cleanup: I'd remove runtime/API.md and runtime/Method.md. Nothing in the repository references them, and they sit outside docs/content, so the docs build doesn't publish them. The Javadoc on waitForCurrentInputInBatchMode already states the contract. If no user-facing doc change is planned, doc-not-needed fits better than doc-included. A note that BATCH now works would belong under docs/content.

[wenjin272]

Sorry for the delayed review—I’ve been focused on the 0.4 release work. Thanks for your patience and for working on this!

We explored the alternative of setting execution.batch-state-backend.enabled=false. BATCH jobs can still use keyed state with a regular state backend, allowing Agents to retain cross-key concurrency. For these workloads, disabling execution.sorted-inputs.enabled as well would avoid unnecessary input-sorting overhead, provided other operators don’t require it.

However, both settings apply job-wide and default to true. Supporting that default configuration correctly is valuable, so I think the approach in this PR is reasonable. Users who need cross-key concurrency can still disable the batch-specific backend.

It would be helpful to document this tradeoff and retain regression coverage for the disabled-backend configuration alongside the new default-backend tests. WDYT?

[wellkilo]

Thanks @wenjin272 — 8ba0573 covers the points from the review:

• The drain now runs once at the end of processElement, after the parallel
  engine's lock is released, so waitInFlightEventsFinished no longer re-enters
  ParallelExecutionLock.
• BatchExecutionTest covers serial and parallel execution, repeated keys on the
  batch backend, and cross-key concurrency with the regular backend;
  CompileUtilsTest is updated accordingly.
• runtime/API.md and runtime/Method.md are removed, and the backend /
  sorted-inputs tradeoff is documented in
  docs/content/docs/operations/deployment.md.

[wenjin272]

Thanks for confirming, @wellkilo, and thanks @sangkyoonnam for the detailed review!

I’ve also rebased this PR onto the latest main, which includes #1371’s fix for the unrelated CI failure in DecisionRoutingRecoveryTest. The PR changes are unchanged; the new head is 47dcdbdb4.

The focused tests passed locally after the rebase: 118 passed and 2 skipped on JDK 11. Let’s see how the updated CI run goes.