Skip to content

(feat) : implement cancellation support with CancellationToken and CancellationTokenSource - #1434

Open
rohan-naik07 wants to merge 3 commits into
google:mainfrom
rohan-naik07:feature/run-cancellation-abort-support
Open

(feat) : implement cancellation support with CancellationToken and CancellationTokenSource#1434
rohan-naik07 wants to merge 3 commits into
google:mainfrom
rohan-naik07:feature/run-cancellation-abort-support

Conversation

@rohan-naik07

@rohan-naik07 rohan-naik07 commented Aug 15, 2026

Copy link
Copy Markdown

Link to Issue or Description of Change

Related issue: #1341

Problem:

Java ADK does not provide an invocation-wide cancellation mechanism. Disposing the returned
RxJava subscription only controls that subscription and cannot be inspected consistently by nested
agents, model flows, tools, or callbacks. As a result, applications cannot implement graceful Stop
buttons or propagate request deadlines through an agent invocation.

Solution:

Add explicit, cooperative cancellation through CancellationToken and
CancellationTokenSource:

CancellationTokenSource cts = new CancellationTokenSource();
RunConfig config =
    RunConfig.builder()
        .cancellationToken(cts.token())
        .build();

Flowable<Event> events = runner.runAsync(userId, sessionId, content, config);

// From a Stop button, timeout handler, or request-disconnect handler:
cts.cancel();

The token is stored in RunConfig and shared through InvocationContext, so the same signal is
visible throughout the invocation, including nested agents. CancellationToken.none() preserves
the existing behavior when callers do not configure cancellation. CancellationTokenSource uses
atomic state, making cancel() thread-safe and idempotent.

Cancellation checkpoints cover:

  • runner and agent startup;
  • request and response processors;
  • model callback boundaries and immediately before model calls;
  • individual streamed model responses;
  • tool callback boundaries and immediately before tool execution;
  • function-response construction, agent transfers, and multi-step loop continuation;
  • both runAsync and runLive execution.

Cancellation is graceful: the event stream completes normally, without an error or a synthetic
cancelled event. Events emitted before cancellation remain available to the caller and can still be
persisted. An already-running atomic model or tool operation may finish before the following
checkpoint. When a live receive flow terminates, its send task is disposed and its model connection
is closed.

RxJava disposal remains complementary and continues to cancel disposable upstream operations. It
is not used as the invocation-wide public cancellation contract.

Testing Plan

Unit Tests:

  • Added tests for the default non-cancelled token.
  • Added tests for token propagation through RunConfig.builder(existingConfig).
  • Added tests for thread-visible, idempotent cancellation and close() behavior.
  • Added a test proving pre-cancelled execution completes without calling the model.
  • Added a test proving cancellation at the before-tool boundary preserves the emitted model
    event and prevents tool execution.
  • Focused cancellation tests pass locally.
  • Full repository test suite passes locally.

Focused test command:

./mvnw -pl core \
  -Dtest=CancellationTokenSourceTest,RunConfigTest,BaseLlmFlowTest test

Result: 42 tests run, 0 failures, 0 errors.

Compilation also succeeds:

./mvnw -pl core -DskipTests compile

The required ./mvnw test command was also run. It currently reports unrelated failures in the
checkout, including NoSuchElementException from the pre-existing
LlmAgent.maybeSaveOutputToState() implementation. The cancellation-focused tests pass in both
configured Surefire executions.

Manual End-to-End (E2E) Tests:

Not run. The behavior is covered at the model and tool execution boundaries with deterministic unit
tests.

Checklist

  • I have read the CONTRIBUTING.md document.
  • My pull request contains a single commit.
  • I have performed a self-review of my own code.
  • I have commented my code, particularly in hard-to-understand areas.
  • I have added tests that prove my fix is effective or that my feature works.
  • New and existing unit tests pass locally with my changes.
  • I have manually tested my changes end-to-end.
  • No dependent downstream changes are required.

Additional context

The API follows the explicit token/source pattern used by other structured cancellation systems
while retaining RxJava as the execution mechanism. It also aligns with TypeScript ADK's documented
graceful, cooperative, and idempotent cancellation semantics.

No cancellation event is written into conversation history. Cancellation represents invocation
lifecycle metadata rather than model-visible conversation content; applications that need to show
a cancelled status can derive it from the token they own.

@google-cla

google-cla Bot commented Aug 15, 2026

Copy link
Copy Markdown

Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

View this failed invocation of the CLA check for more information.

For the most up to date status, view the checks section at the bottom of the pull request.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant