From bd1b9e194d55f0b1b6c2f08a8d8a1f0def4dac6f Mon Sep 17 00:00:00 2001 From: Kamil Tomaszek Date: Fri, 25 Sep 2026 06:03:23 -0700 Subject: [PATCH] ADK changes PiperOrigin-RevId: 988228337 --- .gemini/styleguide.md | 4 + .../analyze-releases-for-adk-docs-updates.yml | 2 +- .github/workflows/pr-commit-check.yml | 4 +- .github/workflows/pr-title-check.yml | 2 +- .github/workflows/pr-triage-adk-java.yml | 2 +- .github/workflows/release-please.yaml | 2 +- .../spam-detection-adk-java-issues.yml | 2 +- .github/workflows/stale-adk-java-issues.yml | 2 +- .github/workflows/triage-adk-java-issues.yml | 2 +- .github/workflows/validation.yml | 2 +- AGENTS.md | 4 + CLAUDE.md | 1 + .../google/adk/a2a/agent/RemoteA2AAgent.java | 2 - .../adk/sessions/FirestoreSessionService.java | 34 +- .../runner/FirestoreDatabaseRunnerTest.java | 3 +- .../sessions/FirestoreSessionServiceTest.java | 132 +- .../properties/SpringAIProperties.java | 4 +- .../models/springai/MessageConverterTest.java | 2 +- .../java/com/google/adk/agents/BaseAgent.java | 47 +- .../google/adk/agents/InvocationContext.java | 270 ++- .../java/com/google/adk/agents/LlmAgent.java | 161 +- .../java/com/google/adk/agents/LoopAgent.java | 153 +- .../com/google/adk/agents/ParallelAgent.java | 77 +- .../google/adk/agents/SequentialAgent.java | 102 +- .../adk/agents/WorkflowAgentResumption.java | 11 +- .../adk/agents/WorkflowAgentStates.java | 103 + .../main/java/com/google/adk/apps/App.java | 11 +- .../google/adk/apps/ResumabilityConfig.java | 43 +- .../java/com/google/adk/events/Event.java | 8 +- .../com/google/adk/events/EventActions.java | 41 +- .../adk/flows/llmflows/AgentTransfer.java | 36 +- .../adk/flows/llmflows/BaseLlmFlow.java | 159 +- .../google/adk/flows/llmflows/Contents.java | 2 - .../google/adk/flows/llmflows/Functions.java | 10 +- .../google/adk/flows/llmflows/StepResume.java | 241 +++ .../agentanalytics/BigQuerySchema.java | 1 - .../java/com/google/adk/runner/Runner.java | 653 +++++- .../com/google/adk/sessions/ApiClient.java | 1 - .../adk/sessions/BaseSessionService.java | 7 + .../adk/sessions/InMemorySessionService.java | 60 +- .../google/adk/sessions/SessionException.java | 3 + .../adk/sessions/SessionJsonConverter.java | 40 +- .../google/adk/sessions/VertexAiClient.java | 4 +- .../adk/tools/GoogleSearchAgentTool.java | 3 +- .../google/adk/tools/GoogleSearchTool.java | 25 +- .../com/google/adk/tools/ToolContext.java | 2 - .../com/google/adk/tools/UrlContextTool.java | 15 +- .../tools/mcp/DefaultMcpTransportBuilder.java | 2 +- .../google/adk/tools/mcp/McpAsyncTool.java | 2 - .../adk/tools/mcp/McpSessionManager.java | 1 - .../com/google/adk/tools/mcp/McpTool.java | 1 - .../adk/agents/InvocationContextTest.java | 510 ++++- .../com/google/adk/agents/LlmAgentTest.java | 303 +++ .../com/google/adk/agents/LoopAgentTest.java | 141 ++ .../google/adk/agents/ParallelAgentTest.java | 202 ++ .../adk/agents/SequentialAgentTest.java | 166 ++ .../google/adk/events/EventActionsTest.java | 53 + .../google/adk/examples/ExampleUtilsTest.java | 1 - .../adk/flows/llmflows/AgentTransferTest.java | 81 +- .../adk/flows/llmflows/FunctionsTest.java | 33 + .../chat/ChatCompletionsResponseTest.java | 2 +- .../BigQueryAgentAnalyticsPluginTest.java | 5 +- .../runner/RunnerLegacyResumabilityTest.java | 565 +++-- .../adk/runner/RunnerResumabilityTest.java | 1864 +++++++++++++++++ .../sessions/InMemorySessionServiceTest.java | 87 + .../sessions/SessionJsonConverterTest.java | 101 + .../adk/testing/ResumabilityTestUtils.java | 294 +++ .../com/google/adk/testing/TestUtils.java | 42 +- .../com/google/adk/tools/BaseToolTest.java | 79 +- dev/INTENRAL_TODOS.md | 22 - dev/README.md | 33 + .../adk/plugins/recordings/LlmRecording.java | 2 +- .../adk/plugins/recordings/Recording.java | 2 +- .../adk/plugins/recordings/ToolRecording.java | 2 +- .../java/com/google/adk/web/AdkWebServer.java | 36 +- .../com/google/adk/web/config/BackendUrl.java | 118 ++ .../google/adk/web/config/DevUiAssets.java | 13 +- .../controller/RuntimeConfigController.java | 104 + .../adk/web/controller/VersionController.java | 49 + .../google/adk/web/dto/AgentRunRequest.java | 8 +- .../web/AdkWebServerProxyRedirectTest.java | 12 +- .../adk/web/BackendUrlRedirectTest.java | 169 ++ .../google/adk/web/config/BackendUrlTest.java | 203 ++ .../adk/web/config/DevUiAssetsTest.java | 6 + .../RuntimeConfigControllerMergeTest.java | 123 ++ .../RuntimeConfigControllerTest.java | 140 ++ .../web/controller/VersionControllerTest.java | 58 + .../kotlin/com/google/adk/tokt/JavaAdkToKt.kt | 56 +- .../com/google/adk/tokt/KotlinAdkToJava.kt | 18 +- .../com/google/adk/tokt/KtRunnerToJava.kt | 36 +- .../adk/tokt/adapters/ActionReconciler.kt | 6 +- .../com/google/adk/tokt/codecs/EventCodec.kt | 18 +- .../google/adk/tokt/codecs/KtToJavaCodecs.kt | 39 +- .../google/adk/tokt/codecs/SessionCodec.kt | 4 +- .../tokt/services/KtSessionServiceToJava.kt | 59 +- .../google/adk/tokt/KtRunnerInteropTest.kt | 844 +++++++- 96 files changed, 8543 insertions(+), 667 deletions(-) create mode 100644 .gemini/styleguide.md create mode 100644 AGENTS.md create mode 100644 CLAUDE.md create mode 100644 core/src/main/java/com/google/adk/agents/WorkflowAgentStates.java create mode 100644 core/src/main/java/com/google/adk/flows/llmflows/StepResume.java create mode 100644 core/src/test/java/com/google/adk/runner/RunnerResumabilityTest.java create mode 100644 core/src/test/java/com/google/adk/testing/ResumabilityTestUtils.java delete mode 100644 dev/INTENRAL_TODOS.md create mode 100644 dev/src/main/java/com/google/adk/web/config/BackendUrl.java create mode 100644 dev/src/main/java/com/google/adk/web/controller/RuntimeConfigController.java create mode 100644 dev/src/main/java/com/google/adk/web/controller/VersionController.java create mode 100644 dev/src/test/java/com/google/adk/web/BackendUrlRedirectTest.java create mode 100644 dev/src/test/java/com/google/adk/web/config/BackendUrlTest.java create mode 100644 dev/src/test/java/com/google/adk/web/controller/RuntimeConfigControllerMergeTest.java create mode 100644 dev/src/test/java/com/google/adk/web/controller/RuntimeConfigControllerTest.java create mode 100644 dev/src/test/java/com/google/adk/web/controller/VersionControllerTest.java diff --git a/.gemini/styleguide.md b/.gemini/styleguide.md new file mode 100644 index 000000000..51a585547 --- /dev/null +++ b/.gemini/styleguide.md @@ -0,0 +1,4 @@ +# Gemini Code Assist Customization + +_Placeholder - customization for Gemini Code Assist in this repository is coming +soon._ diff --git a/.github/workflows/analyze-releases-for-adk-docs-updates.yml b/.github/workflows/analyze-releases-for-adk-docs-updates.yml index 4a964e92c..050d1cf30 100644 --- a/.github/workflows/analyze-releases-for-adk-docs-updates.yml +++ b/.github/workflows/analyze-releases-for-adk-docs-updates.yml @@ -23,7 +23,7 @@ on: jobs: analyze-new-release-for-adk-docs-updates: - runs-on: ubuntu-latest + runs-on: ubuntu-24.04 # Dry-run reads only (this repo + the public docs repo) and skips writes, so # the built-in GITHUB_TOKEN suffices. For --no-dry-run, use a PAT with write # access to the docs repo. diff --git a/.github/workflows/pr-commit-check.yml b/.github/workflows/pr-commit-check.yml index 1e31e42f3..8caaef527 100644 --- a/.github/workflows/pr-commit-check.yml +++ b/.github/workflows/pr-commit-check.yml @@ -13,8 +13,8 @@ on: # Defines the jobs that will run as part of the workflow. jobs: check-commit-count: - # The type of runner that the job will run on. 'ubuntu-latest' is a good default. - runs-on: ubuntu-latest + # The type of runner that the job will run on, pinned to a fixed Ubuntu release. + runs-on: ubuntu-24.04 # The steps that will be executed as part of the job. steps: diff --git a/.github/workflows/pr-title-check.yml b/.github/workflows/pr-title-check.yml index 45ec0b8b2..0e33b0edd 100644 --- a/.github/workflows/pr-title-check.yml +++ b/.github/workflows/pr-title-check.yml @@ -13,7 +13,7 @@ permissions: jobs: check-pr-title: - runs-on: ubuntu-latest + runs-on: ubuntu-24.04 steps: - name: Validate Conventional Commit title env: diff --git a/.github/workflows/pr-triage-adk-java.yml b/.github/workflows/pr-triage-adk-java.yml index 177ef1698..ab990fae7 100644 --- a/.github/workflows/pr-triage-adk-java.yml +++ b/.github/workflows/pr-triage-adk-java.yml @@ -44,7 +44,7 @@ concurrency: jobs: agent-triage-pull-request: - runs-on: ubuntu-latest + runs-on: ubuntu-24.04 # Only run on the upstream repo, for newly-opened/reopened/edited PRs or a # manual dispatch. if: >- diff --git a/.github/workflows/release-please.yaml b/.github/workflows/release-please.yaml index 258cf90db..9879cd6db 100644 --- a/.github/workflows/release-please.yaml +++ b/.github/workflows/release-please.yaml @@ -10,7 +10,7 @@ permissions: name: release-please jobs: release-please: - runs-on: ubuntu-latest + runs-on: ubuntu-24.04 steps: - uses: googleapis/release-please-action@v4 with: diff --git a/.github/workflows/spam-detection-adk-java-issues.yml b/.github/workflows/spam-detection-adk-java-issues.yml index 3312841c1..e3f7d2fac 100644 --- a/.github/workflows/spam-detection-adk-java-issues.yml +++ b/.github/workflows/spam-detection-adk-java-issues.yml @@ -41,7 +41,7 @@ concurrency: jobs: agent-scan-issues: - runs-on: ubuntu-latest + runs-on: ubuntu-24.04 # Only run on the upstream repo, for newly-opened issues, the scheduled # sweep, or a manual dispatch. if: >- diff --git a/.github/workflows/stale-adk-java-issues.yml b/.github/workflows/stale-adk-java-issues.yml index 983422339..0bfc4dbb4 100644 --- a/.github/workflows/stale-adk-java-issues.yml +++ b/.github/workflows/stale-adk-java-issues.yml @@ -34,7 +34,7 @@ concurrency: jobs: agent-audit-stale-issues: - runs-on: ubuntu-latest + runs-on: ubuntu-24.04 # Only run on the upstream repo, on the daily schedule or a manual dispatch. if: >- github.repository == 'google/adk-java' && ( diff --git a/.github/workflows/triage-adk-java-issues.yml b/.github/workflows/triage-adk-java-issues.yml index 7ae4b6ef5..4d1b78fa3 100644 --- a/.github/workflows/triage-adk-java-issues.yml +++ b/.github/workflows/triage-adk-java-issues.yml @@ -32,7 +32,7 @@ concurrency: jobs: agent-triage-issues: - runs-on: ubuntu-latest + runs-on: ubuntu-24.04 # Only run on the upstream repo, for newly-opened issues, the scheduled # batch sweep, or a manual dispatch. if: >- diff --git a/.github/workflows/validation.yml b/.github/workflows/validation.yml index 339ad9f3e..7a3fe8366 100644 --- a/.github/workflows/validation.yml +++ b/.github/workflows/validation.yml @@ -12,7 +12,7 @@ concurrency: jobs: build-modules: - runs-on: ubuntu-latest + runs-on: ubuntu-24.04 timeout-minutes: 30 strategy: matrix: diff --git a/AGENTS.md b/AGENTS.md new file mode 100644 index 000000000..6d86f1b09 --- /dev/null +++ b/AGENTS.md @@ -0,0 +1,4 @@ +# AGENTS.md + +_Placeholder - guidance for AI coding agents working in this repository is coming +soon._ diff --git a/CLAUDE.md b/CLAUDE.md new file mode 100644 index 000000000..543a9c15f --- /dev/null +++ b/CLAUDE.md @@ -0,0 +1 @@ +See [AGENTS.md](./AGENTS.md) for project context, commands, and contribution guidelines for AI coding agents. diff --git a/a2a/src/main/java/com/google/adk/a2a/agent/RemoteA2AAgent.java b/a2a/src/main/java/com/google/adk/a2a/agent/RemoteA2AAgent.java index f134ee43b..571493b8d 100644 --- a/a2a/src/main/java/com/google/adk/a2a/agent/RemoteA2AAgent.java +++ b/a2a/src/main/java/com/google/adk/a2a/agent/RemoteA2AAgent.java @@ -317,8 +317,6 @@ synchronized void handleError(String message, Throwable e) { emitter.tryOnError(new A2AClientError(message, e)); } - // TODO: b/483038527 - The synchronized block might block the thread, we should optimize for - // performance in the future. synchronized void handleEvent(ClientEvent clientEvent, AgentCard unused) { // Mark the flow as done if it is already cancelled. if (!done) { diff --git a/contrib/firestore-session-service/src/main/java/com/google/adk/sessions/FirestoreSessionService.java b/contrib/firestore-session-service/src/main/java/com/google/adk/sessions/FirestoreSessionService.java index f4e68e3ca..340c24d45 100644 --- a/contrib/firestore-session-service/src/main/java/com/google/adk/sessions/FirestoreSessionService.java +++ b/contrib/firestore-session-service/src/main/java/com/google/adk/sessions/FirestoreSessionService.java @@ -21,7 +21,9 @@ import com.google.adk.utils.Constants; import com.google.api.core.ApiFuture; import com.google.api.core.ApiFutures; +import com.google.api.gax.rpc.AlreadyExistsException; import com.google.cloud.firestore.CollectionReference; +import com.google.cloud.firestore.DocumentReference; import com.google.cloud.firestore.DocumentSnapshot; import com.google.cloud.firestore.Firestore; import com.google.cloud.firestore.Query; @@ -48,9 +50,10 @@ import java.util.UUID; import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.ConcurrentMap; +import java.util.concurrent.ExecutionException; import java.util.concurrent.atomic.AtomicBoolean; import java.util.regex.Matcher; -import javax.annotation.Nullable; +import org.jspecify.annotations.Nullable; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -73,6 +76,9 @@ public class FirestoreSessionService implements BaseSessionService { private static final String UPDATE_TIME_KEY = Constants.KEY_UPDATE_TIME; private static final String TIMESTAMP_KEY = Constants.KEY_TIMESTAMP; + /** Random token each create stores, so a retried create can detect its own write. */ + static final String CREATE_TOKEN_KEY = "createToken"; + /** Constructor for FirestoreSessionService. */ public FirestoreSessionService(Firestore firestore) { this.firestore = firestore; @@ -96,7 +102,10 @@ public Single createSession( return createSession(appName, userId, (Map) state, sessionId); } - /** Creates a new session in Firestore. */ + /** + * Creates a new session in Firestore. Session IDs are unique per user across apps, so creating + * one the user already has under any app fails with {@link SessionException}. + */ @Override public Single createSession( String appName, @@ -139,11 +148,22 @@ public Single createSession( sessionData.put(USER_ID_KEY, newSession.userId()); sessionData.put(UPDATE_TIME_KEY, newSession.lastUpdateTime().toString()); sessionData.put(STATE_KEY, newSession.state()); - - // Asynchronously write to Firestore and wait for the result - ApiFuture future = - getSessionsCollection(userId).document(resolvedSessionId).set(sessionData); - future.get(); // Block until the write is complete + String createToken = UUID.randomUUID().toString(); + sessionData.put(CREATE_TOKEN_KEY, createToken); + + // Unlike set(), create() fails if the session already exists instead of replacing it. + DocumentReference sessionDoc = getSessionsCollection(userId).document(resolvedSessionId); + try { + sessionDoc.create(sessionData).get(); + } catch (ExecutionException e) { + if (!(e.getCause() instanceof AlreadyExistsException)) { + throw e; + } + // A retry after a lost reply fails on this call's own write; its token means success. + if (!createToken.equals(sessionDoc.get().get().getString(CREATE_TOKEN_KEY))) { + throw new SessionException(SessionException.SESSION_ALREADY_EXISTS, e.getCause()); + } + } return newSession; }); diff --git a/contrib/firestore-session-service/src/test/java/com/google/adk/runner/FirestoreDatabaseRunnerTest.java b/contrib/firestore-session-service/src/test/java/com/google/adk/runner/FirestoreDatabaseRunnerTest.java index 05c78efe8..eae6de3b4 100644 --- a/contrib/firestore-session-service/src/test/java/com/google/adk/runner/FirestoreDatabaseRunnerTest.java +++ b/contrib/firestore-session-service/src/test/java/com/google/adk/runner/FirestoreDatabaseRunnerTest.java @@ -175,7 +175,8 @@ void run_withUserInput_createsSessionAndExecutesAgent() throws Exception { when(mockUserDocRef.collection(anyString())).thenReturn(mockSessionsCollection); when(mockSessionsCollection.document(anyString())).thenReturn(mockSessionDocRef); - when(mockSessionDocRef.set(anyMap())).thenReturn(ApiFutures.immediateFuture(mockWriteResult)); + when(mockSessionDocRef.create(anyMap())) + .thenReturn(ApiFutures.immediateFuture(mockWriteResult)); when(mockSessionDocRef.update(anyMap())) .thenReturn(ApiFutures.immediateFuture(mockWriteResult)); // Mock the event sub-collection chain diff --git a/contrib/firestore-session-service/src/test/java/com/google/adk/sessions/FirestoreSessionServiceTest.java b/contrib/firestore-session-service/src/test/java/com/google/adk/sessions/FirestoreSessionServiceTest.java index 71fa8a7bd..9b7d1a522 100644 --- a/contrib/firestore-session-service/src/test/java/com/google/adk/sessions/FirestoreSessionServiceTest.java +++ b/contrib/firestore-session-service/src/test/java/com/google/adk/sessions/FirestoreSessionServiceTest.java @@ -31,6 +31,10 @@ import com.google.adk.events.EventActions; import com.google.adk.utils.Constants; import com.google.api.core.ApiFutures; +import com.google.api.gax.grpc.GrpcStatusCode; +import com.google.api.gax.rpc.AlreadyExistsException; +import com.google.api.gax.rpc.PermissionDeniedException; +import com.google.api.gax.rpc.UnavailableException; import com.google.cloud.firestore.CollectionReference; import com.google.cloud.firestore.DocumentReference; import com.google.cloud.firestore.DocumentSnapshot; @@ -45,6 +49,7 @@ import com.google.common.collect.ImmutableMap; import com.google.genai.types.Content; import com.google.genai.types.Part; +import io.grpc.Status; import io.reactivex.rxjava3.observers.TestObserver; import java.time.Instant; import java.util.Collections; @@ -52,10 +57,12 @@ import java.util.Map; import java.util.Optional; import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.ExecutionException; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; import org.junit.jupiter.api.extension.ExtendWith; import org.mockito.ArgumentCaptor; +import org.mockito.Captor; import org.mockito.Mock; import org.mockito.junit.jupiter.MockitoExtension; @@ -89,6 +96,7 @@ public class FirestoreSessionServiceTest { @Mock private QuerySnapshot mockQuerySnapshot; @Mock private WriteResult mockWriteResult; @Mock private WriteBatch mockWriteBatch; + @Captor private ArgumentCaptor> sessionDataCaptor; private FirestoreSessionService sessionService; @@ -130,7 +138,7 @@ public void setup() { // Default mock for writes lenient() - .when(mockSessionDocRef.set(anyMap())) + .when(mockSessionDocRef.create(anyMap())) .thenReturn(ApiFutures.immediateFuture(mockWriteResult)); lenient() .when(mockSessionDocRef.update(anyMap())) @@ -307,7 +315,7 @@ void createSession_withSessionId_returnsNewSession() { assertThat(session.id()).isEqualTo(SESSION_ID); return true; }); - verify(mockSessionDocRef).set(anyMap()); + verify(mockSessionDocRef).create(anyMap()); } /** Tests that createSession creates a new session with a generated session ID. */ @@ -334,7 +342,7 @@ void createSession_withNullSessionId_generatesNewId() { assertThat(session.id()).isNotEmpty(); return true; }); - verify(mockSessionDocRef).set(anyMap()); + verify(mockSessionDocRef).create(anyMap()); } /** Tests that createSession creates a new session with an empty session ID. */ @@ -358,7 +366,7 @@ void createSession_withEmptySessionId_generatesNewId() { assertThat(session.id()).isNotEqualTo(" "); return true; }); - verify(mockSessionDocRef).set(anyMap()); + verify(mockSessionDocRef).create(anyMap()); } /** Tests that createSession creates a new session with an empty state when null state is */ @@ -391,6 +399,114 @@ void createSession_withNullAppName_throwsNullPointerException() { .assertError(NullPointerException.class); } + /** Tests that createSession rejects a session ID that is already taken. */ + @Test + void createSession_withSessionIdAlreadyTaken_failsWithSessionException() { + // Arrange + when(mockSessionsCollection.document(SESSION_ID)).thenReturn(mockSessionDocRef); + when(mockSessionDocRef.create(anyMap())) + .thenReturn(ApiFutures.immediateFailedFuture(alreadyExists())); + when(mockSessionDocRef.get()).thenReturn(ApiFutures.immediateFuture(mockSessionSnapshot)); + // Another caller's session, or one written before sessions carried a token. + when(mockSessionSnapshot.getString(FirestoreSessionService.CREATE_TOKEN_KEY)).thenReturn(null); + + // Act + TestObserver testObserver = + sessionService.createSession(APP_NAME, USER_ID, null, SESSION_ID).test(); + + // Assert + testObserver.assertError( + e -> { + assertThat(e).isInstanceOf(SessionException.class); + assertThat(e).hasMessageThat().isEqualTo(SessionException.SESSION_ALREADY_EXISTS); + assertThat(e).hasCauseThat().isInstanceOf(AlreadyExistsException.class); + return true; + }); + } + + /** Tests that createSession succeeds when a client retry fails on the session's own write. */ + @Test + void createSession_whenRetryHitsItsOwnWrite_returnsSession() { + // Arrange + when(mockSessionsCollection.document(SESSION_ID)).thenReturn(mockSessionDocRef); + when(mockSessionDocRef.create(sessionDataCaptor.capture())) + .thenReturn(ApiFutures.immediateFailedFuture(alreadyExists())); + when(mockSessionDocRef.get()).thenReturn(ApiFutures.immediateFuture(mockSessionSnapshot)); + when(mockSessionSnapshot.getString(FirestoreSessionService.CREATE_TOKEN_KEY)) + .thenAnswer( + unused -> sessionDataCaptor.getValue().get(FirestoreSessionService.CREATE_TOKEN_KEY)); + + // Act + TestObserver testObserver = + sessionService.createSession(APP_NAME, USER_ID, null, SESSION_ID).test(); + + // Assert + testObserver.assertComplete(); + testObserver.assertValue(session -> session.id().equals(SESSION_ID)); + } + + /** + * Tests that a failed read-back after a duplicate is propagated, since it cannot tell whose + * session exists. + */ + @Test + void createSession_whenReadBackFails_propagatesFailure() { + // Arrange + when(mockSessionsCollection.document(SESSION_ID)).thenReturn(mockSessionDocRef); + when(mockSessionDocRef.create(anyMap())) + .thenReturn(ApiFutures.immediateFailedFuture(alreadyExists())); + when(mockSessionDocRef.get()) + .thenReturn( + ApiFutures.immediateFailedFuture( + new UnavailableException( + "Service unavailable", + /* cause= */ null, + GrpcStatusCode.of(Status.Code.UNAVAILABLE), + /* retryable= */ true))); + + // Act + TestObserver testObserver = + sessionService.createSession(APP_NAME, USER_ID, null, SESSION_ID).test(); + + // Assert + testObserver.assertError( + e -> { + assertThat(e).isInstanceOf(ExecutionException.class); + assertThat(e).hasCauseThat().isInstanceOf(UnavailableException.class); + return true; + }); + } + + /** + * Tests that a write failure unrelated to a duplicate session ID is propagated rather than + * reported as one. + */ + @Test + void createSession_whenWriteFailsForAnotherReason_propagatesFailure() { + // Arrange + when(mockSessionsCollection.document(SESSION_ID)).thenReturn(mockSessionDocRef); + when(mockSessionDocRef.create(anyMap())) + .thenReturn( + ApiFutures.immediateFailedFuture( + new PermissionDeniedException( + "Missing or insufficient permissions", + /* cause= */ null, + GrpcStatusCode.of(Status.Code.PERMISSION_DENIED), + /* retryable= */ false))); + + // Act + TestObserver testObserver = + sessionService.createSession(APP_NAME, USER_ID, null, SESSION_ID).test(); + + // Assert + testObserver.assertError( + e -> { + assertThat(e).isInstanceOf(ExecutionException.class); + assertThat(e).hasCauseThat().isInstanceOf(PermissionDeniedException.class); + return true; + }); + } + // --- appendEvent Tests --- /** Tests that appendEvent persists the event and updates the session's updateTime. */ @Test @@ -978,4 +1094,12 @@ void deleteSession_appNameMismatch_doesNotDelete() { verify(mockSessionDocRef, never()).delete(); verify(mockWriteBatch, never()).commit(); } + + private static AlreadyExistsException alreadyExists() { + return new AlreadyExistsException( + "Document already exists", + /* cause= */ null, + GrpcStatusCode.of(Status.Code.ALREADY_EXISTS), + /* retryable= */ false); + } } diff --git a/contrib/spring-ai/src/main/java/com/google/adk/models/springai/properties/SpringAIProperties.java b/contrib/spring-ai/src/main/java/com/google/adk/models/springai/properties/SpringAIProperties.java index 5e972ebcd..b3f356a78 100644 --- a/contrib/spring-ai/src/main/java/com/google/adk/models/springai/properties/SpringAIProperties.java +++ b/contrib/spring-ai/src/main/java/com/google/adk/models/springai/properties/SpringAIProperties.java @@ -15,10 +15,10 @@ */ package com.google.adk.models.springai.properties; -import jakarta.annotation.Nullable; import jakarta.validation.constraints.DecimalMax; import jakarta.validation.constraints.DecimalMin; import jakarta.validation.constraints.Min; +import org.jspecify.annotations.Nullable; import org.springframework.boot.context.properties.ConfigurationProperties; import org.springframework.validation.annotation.Validated; @@ -41,7 +41,7 @@ @Validated public class SpringAIProperties { - @Nullable private String model; + private @Nullable String model; /** Default temperature for controlling randomness in responses. Must be between 0.0 and 2.0. */ @DecimalMin(value = "0.0", message = "Temperature must be at least 0.0") diff --git a/contrib/spring-ai/src/test/java/com/google/adk/models/springai/MessageConverterTest.java b/contrib/spring-ai/src/test/java/com/google/adk/models/springai/MessageConverterTest.java index 4abb7ad1d..310bf3d07 100644 --- a/contrib/spring-ai/src/test/java/com/google/adk/models/springai/MessageConverterTest.java +++ b/contrib/spring-ai/src/test/java/com/google/adk/models/springai/MessageConverterTest.java @@ -818,7 +818,7 @@ public Optional declaration() { @Test void testToolOptionsPreserveProviderSpecificTypeToAvoidClassCastException() { - // Regression test for b/527041291 (GitHub adk-java #1295): Spring AI OpenAI 2.0.0 casts + // Regression test for GitHub adk-java #1295: Spring AI OpenAI 2.0.0 casts // Prompt.getOptions() directly to OpenAiChatOptions in createRequest(). When ADK passed a // provider-neutral DefaultToolCallingChatOptions, that cast threw a ClassCastException. Basing // the prompt options on the model's own options must keep the concrete provider type. diff --git a/core/src/main/java/com/google/adk/agents/BaseAgent.java b/core/src/main/java/com/google/adk/agents/BaseAgent.java index eeb04bdf1..c4bb0fbb7 100644 --- a/core/src/main/java/com/google/adk/agents/BaseAgent.java +++ b/core/src/main/java/com/google/adk/agents/BaseAgent.java @@ -16,6 +16,7 @@ package com.google.adk.agents; +import static com.google.common.base.Preconditions.checkArgument; import static com.google.common.base.Strings.isNullOrEmpty; import static com.google.common.collect.ImmutableList.toImmutableList; import static java.lang.String.format; @@ -23,6 +24,7 @@ import com.google.adk.agents.Callbacks.AfterAgentCallback; import com.google.adk.agents.Callbacks.BeforeAgentCallback; import com.google.adk.events.Event; +import com.google.adk.events.EventActions; import com.google.adk.plugins.Plugin; import com.google.adk.telemetry.Instrumentation; import com.google.adk.telemetry.Instrumentation.AgentInvocation; @@ -39,6 +41,7 @@ import java.util.ArrayList; import java.util.HashSet; import java.util.List; +import java.util.Map; import java.util.Optional; import java.util.function.Function; import java.util.regex.Pattern; @@ -136,10 +139,8 @@ private static void validateAgentName(String name) { throw new IllegalArgumentException( format("Agent name '%s' does not match regex '%s'.", name, IDENTIFIER_REGEX)); } - if (name.equals(Role.USER)) { - throw new IllegalArgumentException( - "Agent name cannot be 'user'; reserved for end-user input."); - } + checkArgument( + !name.equals(Role.USER), "Agent name cannot be 'user'; reserved for end-user input."); } /** @@ -468,6 +469,44 @@ public Flowable runLive(InvocationContext parentContext) { return run(parentContext, this::runLiveImpl); } + /** + * Records this agent's end-of-agent checkpoint and returns it as a single-event stream. The + * recorded state is cleared and the agent marked finished, so a later run skips it. + * + * @param context Current invocation context. + * @return stream of the single {@code endOfAgent = true} checkpoint event. + */ + final Flowable endOfAgentAndRecord(InvocationContext context) { + context.setAgentState(name(), /* agentState= */ null, /* endOfAgent= */ true); + return Flowable.just(checkpointEvent(context, EventActions.builder().endOfAgent(true).build())); + } + + /** + * Records {@code agentState} for this agent and returns the matching checkpoint event as a + * single-event stream. The agent is left unfinished, so a later run resumes from this checkpoint. + * + * @param context Current invocation context. + * @param agentState The serialized agent state to persist. + * @return stream of the single checkpoint event carrying {@code agentState}. + */ + final Flowable checkpointAndRecord( + InvocationContext context, Map agentState) { + context.setAgentState(name(), agentState, /* endOfAgent= */ false); + return Flowable.just( + checkpointEvent(context, EventActions.builder().agentState(agentState).build())); + } + + /** Builds a resumability checkpoint event authored by this agent carrying {@code actions}. */ + private Event checkpointEvent(InvocationContext context, EventActions actions) { + return Event.builder() + .id(Event.generateEventId()) + .invocationId(context.invocationId()) + .author(name()) + .branch(context.branch().orElse(null)) + .actions(actions) + .build(); + } + /** * Agent-specific asynchronous logic. * diff --git a/core/src/main/java/com/google/adk/agents/InvocationContext.java b/core/src/main/java/com/google/adk/agents/InvocationContext.java index 384fd8372..6c874c3e1 100644 --- a/core/src/main/java/com/google/adk/agents/InvocationContext.java +++ b/core/src/main/java/com/google/adk/agents/InvocationContext.java @@ -20,9 +20,11 @@ import static com.google.common.collect.ImmutableList.toImmutableList; import static com.google.common.collect.ImmutableSet.toImmutableSet; +import com.google.adk.annotations.Experimental; import com.google.adk.apps.ResumabilityConfig; import com.google.adk.artifacts.BaseArtifactService; import com.google.adk.events.Event; +import com.google.adk.flows.llmflows.Functions; import com.google.adk.memory.BaseMemoryService; import com.google.adk.models.LlmCallsLimitExceededException; import com.google.adk.plugins.Plugin; @@ -31,22 +33,27 @@ import com.google.adk.sessions.Session; import com.google.adk.summarizer.EventsCompactionConfig; import com.google.common.collect.ImmutableList; +import com.google.common.collect.ImmutableMap; import com.google.common.collect.ImmutableSet; import com.google.errorprone.annotations.CanIgnoreReturnValue; import com.google.genai.types.Content; import com.google.genai.types.FunctionCall; import com.google.genai.types.FunctionResponse; +import java.util.ArrayList; import java.util.Collections; +import java.util.HashSet; +import java.util.LinkedHashMap; +import java.util.List; import java.util.Map; import java.util.Objects; import java.util.Optional; +import java.util.Set; import java.util.UUID; import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.atomic.AtomicInteger; import org.jspecify.annotations.Nullable; /** The context for an agent invocation. */ -@SuppressWarnings("deprecation") // Plumbs the deprecated ResumabilityConfig. public class InvocationContext { private final BaseSessionService sessionService; @@ -64,6 +71,9 @@ public class InvocationContext { private final @Nullable ResumabilityConfig resumabilityConfig; private final InvocationCostManager invocationCostManager; private final Map callbackContextData; + // Shared by reference so a sub-agent's checkpoint is visible to its parent and the runner. + private final Map> agentStates; + private final Map endOfAgents; @Nullable private String branch; private BaseAgent agent; @@ -91,6 +101,8 @@ protected InvocationContext(Builder builder) { // invocation invocation so that Plugins can access the same data it during the invocation // across all types of callbacks. this.callbackContextData = builder.callbackContextData; + this.agentStates = builder.agentStates; + this.endOfAgents = builder.endOfAgents; } /** Returns a new {@link Builder} for creating {@link InvocationContext} instances. */ @@ -304,16 +316,253 @@ public Optional contextCacheConfig() { return Optional.ofNullable(contextCacheConfig); } + /** Returns whether the current invocation is resumable. */ + @Experimental + public boolean isResumable() { + return resumabilityConfig != null && resumabilityConfig.isResumable(); + } + /** - * Returns whether the current invocation is resumable. Mirrors Python ADK v1's {@code - * InvocationContext.is_resumable}. The deprecated plain-text continuation shim selects the same - * resumption behavior, so it reports resumable too; the two are mutually exclusive. + * Returns whether the invocation runs the legacy resumption flow: the behavior {@link + * #isResumable()} had before durable checkpoints existed. It is a separate flow, not a weaker + * {@link #isResumable()}, so a caller wanting either has to ask for both. + * + * @deprecated Reports the deprecated plain-text continuation shim and goes away with it; use + * {@link #isResumable()}. */ + @Deprecated @SuppressWarnings("deprecation") // The shim it reads is deprecated by design. - public boolean isResumable() { - return resumabilityConfig != null - && (resumabilityConfig.isResumable() - || resumabilityConfig.isPlainTextContinuationAutoResume()); + public boolean isLegacyResumability() { + return resumabilityConfig != null && resumabilityConfig.isPlainTextContinuationAutoResume(); + } + + /** + * Returns an unmodifiable view of the per-agent resumability checkpoint states for this + * invocation, keyed by agent name. The backing map is shared by reference across derived contexts + * within the invocation; mutate it only through {@link #setAgentState}. + */ + @Experimental + public Map> agentStates() { + return Collections.unmodifiableMap(agentStates); + } + + /** + * Returns an unmodifiable view of the per-agent end-of-agent flags for this invocation, keyed by + * agent name. + */ + @Experimental + public Map endOfAgents() { + return Collections.unmodifiableMap(endOfAgents); + } + + /** + * Sets the checkpoint state of an agent explicitly. Does not implicitly initialize. + * + * @param agentName the agent whose state to set. + * @param agentState the serialized agent state to store; ignored when {@code endOfAgent} is true. + * @param endOfAgent when true, marks the agent finished and drops any stored state. + */ + void setAgentState( + String agentName, @Nullable Map agentState, boolean endOfAgent) { + if (endOfAgent) { + endOfAgents.put(agentName, true); + agentStates.remove(agentName); + } else if (agentState != null) { + // LinkedHashMap, not ImmutableMap.copyOf: a deserialized agentState may carry a null value. + agentStates.put(agentName, Collections.unmodifiableMap(new LinkedHashMap<>(agentState))); + endOfAgents.put(agentName, false); + } else { + endOfAgents.remove(agentName); + agentStates.remove(agentName); + } + } + + /** Recursively resets the checkpoint state of all sub-agents of the given agent. */ + void resetSubAgentStates(String agentName) { + Optional target = agent.findAgent(agentName); + if (target.isEmpty()) { + return; + } + for (BaseAgent subAgent : target.get().subAgents()) { + setAgentState(subAgent.name(), /* agentState= */ null, /* endOfAgent= */ false); + resetSubAgentStates(subAgent.name()); + } + } + + /** + * Rehydrates {@link #agentStates()} and {@link #endOfAgents()} from the current invocation's + * history when this invocation is resumable. For each event carrying agent-state information, + * sets the authoring agent's checkpoint; for a non-workflow author that already produced content, + * seeds an empty state so it is treated as mid-run. + */ + @Experimental + public void populateInvocationAgentStates() { + if (!isResumable()) { + return; + } + for (Event event : events(/* currentInvocation= */ true, /* currentBranch= */ false)) { + String author = event.author(); + if (author == null) { + continue; + } + Optional> agentState = event.actions().agentState(); + // The deprecated setEndInvocation aliases endOfAgent, so only a content-less marker counts. + if (event.actions().endOfAgent() && event.content().isEmpty()) { + endOfAgents.put(author, true); + agentStates.remove(author); + } else if (agentState.isPresent()) { + setAgentState(author, agentState.get(), /* endOfAgent= */ false); + } else if (!author.equals(Role.USER) + && event.content().isPresent() + && !agentStates.containsKey(author)) { + // Content after an end-of-agent marker reopens the agent, as in Python. + agentStates.put(author, ImmutableMap.of()); + endOfAgents.put(author, false); + } + } + } + + /** + * Returns the current session's events, optionally filtered to the current invocation and/or the + * current branch. Reads the in-memory {@link Session#events()} list, which {@link + * BaseSessionService#appendEvent} keeps in sync. A {@code null}-branch event is visible on any + * branch. + * + * @param currentInvocation whether to filter to events from this invocation. + * @param currentBranch whether to filter to events on this branch (or with no branch). + */ + @Experimental + public ImmutableList events(boolean currentInvocation, boolean currentBranch) { + // eventsOnCurrentBranch owns branch filtering: author-asymmetric rule and cross-branch guard. + List results; + if (currentBranch) { + results = new ArrayList<>(eventsOnCurrentBranch()); + } else { + // session.events() is a synchronized list; copy it under its own monitor. + List sessionEvents = session.events(); + synchronized (sessionEvents) { + results = new ArrayList<>(sessionEvents); + } + } + if (currentInvocation) { + results.removeIf(event -> !invocationId.equals(event.invocationId())); + } + return ImmutableList.copyOf(results); + } + + /** + * Returns whether to pause the invocation right after this event: it is resumable and the event + * carries a long-running function call, including a synthetic {@code adk_request_confirmation} + * HITL request. Pausing, unlike ending, leaves the invocation resumable. + */ + @Experimental + public boolean shouldPauseInvocation(Event event) { + return isResumable() && carriesLongRunningCall(event); + } + + /** + * Returns whether the event carries a long-running function call, independent of the mode. Shared + * with the legacy flow's check so the two cannot drift apart. + */ + private static boolean carriesLongRunningCall(Event event) { + return Functions.hasPendingLongRunningCall(event); + } + + /** + * Returns whether either of the last two events on this branch pauses the invocation, the + * condition Python uses to decide whether to withhold an agent's end-of-agent checkpoint. + */ + boolean lastEventsPauseInvocation() { + ImmutableList events = events(/* currentInvocation= */ true, /* currentBranch= */ true); + return events.subList(Math.max(0, events.size() - 2), events.size()).stream() + .anyMatch(this::shouldPauseInvocation); + } + + /** + * Returns whether a long-running call made anywhere inside {@code agent}'s subtree is still + * unanswered. A workflow agent needs this because a sub-agent can pause without emitting an + * event, and because Java gives every sub-agent of a {@link ParallelAgent} that agent's own + * branch, which no branch filter can tell apart. Scoping by subtree keeps a paused parallel + * sibling from stalling an unrelated branch, as Python does. + */ + boolean hasUnansweredLongRunningCallIn(BaseAgent agent) { + Set scope = new HashSet<>(); + collectAgentNames(agent, scope); + return hasUnansweredLongRunningCall( + events(/* currentInvocation= */ true, /* currentBranch= */ false), scope); + } + + private static void collectAgentNames(BaseAgent agent, Set names) { + if (!names.add(agent.name())) { + return; + } + for (BaseAgent subAgent : agent.subAgents()) { + collectAgentNames(subAgent, names); + } + } + + /** + * Whether {@code events} hold a long-running call with no response. {@code scope} restricts which + * authors' calls count; responses always count whoever authored them, since a resumed answer is + * authored by the user. + */ + private static boolean hasUnansweredLongRunningCall( + ImmutableList events, Set scope) { + if (events.isEmpty()) { + return false; + } + Set awaited = new HashSet<>(); + for (Event event : events) { + if (!carriesLongRunningCall(event)) { + continue; + } + if (!scope.contains(event.author())) { + continue; + } + // Only long-running calls count, as in Python: an ordinary call in the event runs inline. + Set longRunningIds = event.longRunningToolIds().orElse(ImmutableSet.of()); + for (FunctionCall call : event.functionCalls()) { + call.id().filter(longRunningIds::contains).ifPresent(awaited::add); + } + } + if (awaited.isEmpty()) { + return false; + } + Set answered = new HashSet<>(); + for (Event event : events) { + for (FunctionResponse response : event.functionResponses()) { + response.id().ifPresent(answered::add); + } + } + return !answered.containsAll(awaited); + } + + /** + * Finds the current-invocation event whose function call matches any function response id in + * {@code functionResponseEvent}, searching newest-first. Matching any id (not just the first) + * keeps parallel function responses resolvable when their calls interleave. + */ + Optional findMatchingFunctionCall(Event functionResponseEvent) { + Set targetIds = new HashSet<>(); + for (FunctionResponse response : functionResponseEvent.functionResponses()) { + response.id().ifPresent(targetIds::add); + } + if (targetIds.isEmpty()) { + return Optional.empty(); + } + ImmutableList events = events(/* currentInvocation= */ true, /* currentBranch= */ false); + for (int i = events.size() - 1; i >= 0; i--) { + // The response event never answers itself, as in Python. + if (events.get(i).id().equals(functionResponseEvent.id())) { + continue; + } + for (FunctionCall call : events.get(i).functionCalls()) { + if (call.id().filter(targetIds::contains).isPresent()) { + return Optional.of(events.get(i)); + } + } + } + return Optional.empty(); } private static class InvocationCostManager { @@ -375,6 +624,9 @@ private Builder(InvocationContext context) { // invocation invocation so that Plugins can access the same data it during the invocation // across all types of callbacks. this.callbackContextData = context.callbackContextData; + // Shared by reference, not copied: the checkpoints belong to the invocation, not a context. + this.agentStates = context.agentStates; + this.endOfAgents = context.endOfAgents; } private BaseSessionService sessionService; @@ -395,6 +647,8 @@ private Builder(InvocationContext context) { private @Nullable ResumabilityConfig resumabilityConfig; private InvocationCostManager invocationCostManager = new InvocationCostManager(); private Map callbackContextData = new ConcurrentHashMap<>(); + private Map> agentStates = new ConcurrentHashMap<>(); + private Map endOfAgents = new ConcurrentHashMap<>(); /** * Sets the session service for managing session state. diff --git a/core/src/main/java/com/google/adk/agents/LlmAgent.java b/core/src/main/java/com/google/adk/agents/LlmAgent.java index fa754e0c0..d36a2b2f6 100644 --- a/core/src/main/java/com/google/adk/agents/LlmAgent.java +++ b/core/src/main/java/com/google/adk/agents/LlmAgent.java @@ -55,6 +55,7 @@ import com.google.adk.tools.BaseToolset; import com.google.common.base.Preconditions; import com.google.common.collect.ImmutableList; +import com.google.common.collect.Iterables; import com.google.errorprone.annotations.CanIgnoreReturnValue; import com.google.genai.types.Content; import com.google.genai.types.GenerateContentConfig; @@ -70,6 +71,7 @@ import java.util.Objects; import java.util.Optional; import java.util.concurrent.Executor; +import java.util.concurrent.atomic.AtomicBoolean; import java.util.function.Function; import org.jspecify.annotations.Nullable; import org.slf4j.Logger; @@ -80,6 +82,9 @@ public class LlmAgent extends BaseAgent { private static final Logger logger = LoggerFactory.getLogger(LlmAgent.class); + /** Name of the built-in tool whose response records an agent transfer. */ + private static final String TRANSFER_TO_AGENT_TOOL = "transfer_to_agent"; + /** * Enum to define if contents of previous events should be included in requests to the underlying * LLM. @@ -277,7 +282,7 @@ public Builder disallowTransferToPeers(boolean disallowTransferToPeers) { return this; } - // (b/476510024): Temporary workaround for ces + /** Removes any before-model callbacks set on this builder. */ @CanIgnoreReturnValue public Builder clearBeforeModelCallbacks() { this.beforeModelCallback = null; @@ -660,7 +665,135 @@ private static boolean isThought(Part part) { @Override protected Flowable runAsyncImpl(InvocationContext invocationContext) { - return llmFlow.run(invocationContext).doOnNext(this::maybeSaveOutputToState); + if (!invocationContext.isResumable()) { + return llmFlow.run(invocationContext).doOnNext(this::maybeSaveOutputToState); + } + return Flowable.defer( + () -> { + // Resumed after a transfer: continue the sub-agent instead of re-invoking the model. + if (invocationContext.agentStates().containsKey(name())) { + Optional resumeTarget = findSubAgentToResume(invocationContext); + if (resumeTarget.isPresent()) { + AtomicBoolean resumePaused = new AtomicBoolean(false); + return resumeTarget + .get() + .runAsync(invocationContext) + .doOnNext( + event -> { + if (invocationContext.shouldPauseInvocation(event)) { + resumePaused.set(true); + } + }) + .concatWith( + Flowable.defer( + () -> { + if (resumePaused.get()) { + return Flowable.empty(); + } + return endOfAgentAndRecord(invocationContext); + })); + } + } + // Normal path: emit an end-of-agent checkpoint so a later run can skip this agent. + Flowable events = + llmFlow.run(invocationContext).doOnNext(this::maybeSaveOutputToState); + AtomicBoolean paused = new AtomicBoolean(false); + return events + .doOnNext( + event -> { + if (invocationContext.shouldPauseInvocation(event)) { + paused.set(true); + } + }) + .concatWith( + Flowable.defer( + () -> { + // Only the checkpoint is withheld, as in Python; the flow still runs. A + // transferring agent stays open so a resume re-enters it and reaches the + // sub-agent it handed off to. + if (paused.get() || invocationContext.lastEventsPauseInvocation()) { + return Flowable.empty(); + } + return endOfAgentAndRecord(invocationContext); + })); + }); + } + + /** + * Returns the name of the agent this agent transferred to in {@code event}, or empty when {@code + * event} is not such a transfer. A {@code transfer_to_agent} function response is required, as in + * Python ADK, so a callback that sets the action on an unrelated event does not read as one. + */ + private Optional transferTargetNameFrom(Event event) { + if (!name().equals(event.author()) + || event.functionResponses().stream() + .noneMatch( + response -> response.name().filter(TRANSFER_TO_AGENT_TOOL::equals).isPresent())) { + return Optional.empty(); + } + return event.actions().transferToAgent().filter(target -> !target.equals(name())); + } + + /** + * Resolves a transfer target recorded in history back to its agent. A live transfer fails in + * {@code BaseLlmFlow} when the target is gone; a resumed one has no such check, and returning + * empty here would silently re-run this agent instead of the sub-agent it handed off to. + */ + private BaseAgent resolveTransferTarget(String targetName) { + // Declared targets first, as in Python, so a same-named agent on an unrelated branch is not + // resumed into; then the same tree-wide search the live transfer path uses. + for (BaseAgent target : transferTargets()) { + if (target.name().equals(targetName)) { + return target; + } + } + return rootAgent() + .findAgent(targetName) + .orElseThrow( + () -> + new IllegalStateException( + "Cannot resume agent " + + name() + + ": the agent it transferred to is no longer in the agent tree.")); + } + + /** + * When this agent is being resumed, returns the sub-agent it had transferred to (so the resume + * continues that sub-agent), or empty when this agent should continue itself. + */ + private Optional findSubAgentToResume(InvocationContext context) { + ImmutableList events = + context.events(/* currentInvocation= */ true, /* currentBranch= */ true); + if (events.isEmpty()) { + return Optional.empty(); + } + Event lastEvent = Iterables.getLast(events); + if (name().equals(lastEvent.author())) { + return transferTargetNameFrom(lastEvent).map(this::resolveTransferTarget); + } + if (Objects.equals(lastEvent.author(), Role.USER) && !lastEvent.functionResponses().isEmpty()) { + // IAE (not ISE): an unresolvable resume surfaces through Runner.runAsync's IAE contract. + Event functionCallEvent = + context + .findMatchingFunctionCall(lastEvent) + .orElseThrow( + () -> + new IllegalArgumentException( + "No matching function call to resume agent " + + name() + + " from a function response.")); + if (name().equals(functionCallEvent.author())) { + return Optional.empty(); + } + } + // A user event answering no call falls through; Python fails the resume instead. + for (int i = events.size() - 2; i >= 0; i--) { + Optional targetName = transferTargetNameFrom(events.get(i)); + if (targetName.isPresent()) { + return Optional.of(resolveTransferTarget(targetName.get())); + } + } + return Optional.empty(); } @Override @@ -789,6 +922,30 @@ public boolean disallowTransferToPeers() { return disallowTransferToPeers; } + /** + * Returns the agents this one may transfer to: its sub-agents, plus its parent and peers unless + * disallowed. Transfer to a parent or peer needs an {@link LlmAgent} parent. Package-private + * because {@code AgentTransfer} cannot reach it and keeps its own copy of the same rule. + */ + ImmutableList transferTargets() { + List targets = new ArrayList<>(subAgents()); + BaseAgent parent = parentAgent(); + if (!(parent instanceof LlmAgent)) { + return ImmutableList.copyOf(targets); + } + if (!disallowTransferToParent()) { + targets.add(parent); + } + if (!disallowTransferToPeers()) { + for (BaseAgent peer : parent.subAgents()) { + if (!peer.name().equals(name())) { + targets.add(peer); + } + } + } + return ImmutableList.copyOf(targets); + } + public List beforeModelCallback() { return beforeModelCallback; } diff --git a/core/src/main/java/com/google/adk/agents/LoopAgent.java b/core/src/main/java/com/google/adk/agents/LoopAgent.java index 19fd4c497..927e8ee9a 100644 --- a/core/src/main/java/com/google/adk/agents/LoopAgent.java +++ b/core/src/main/java/com/google/adk/agents/LoopAgent.java @@ -18,11 +18,14 @@ import com.google.adk.agents.ConfigAgentUtils.ConfigurationException; import com.google.adk.events.Event; +import com.google.common.collect.ImmutableMap; import com.google.errorprone.annotations.CanIgnoreReturnValue; import io.reactivex.rxjava3.core.Flowable; import java.util.List; +import java.util.Map; import java.util.concurrent.atomic.AtomicBoolean; import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.atomic.AtomicReference; import org.jspecify.annotations.Nullable; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -96,7 +99,6 @@ public Builder maxIterations(@Nullable Integer maxIterations) { @Override public LoopAgent build() { - // TODO(b/410859954): Add validation for required fields like name. return new LoopAgent( name, description, subAgents, maxIterations, beforeAgentCallback, afterAgentCallback); } @@ -136,12 +138,17 @@ public static LoopAgent fromConfig(LoopAgentConfig config, String configAbsPath) } @Override + @SuppressWarnings("deprecation") // The shim it dispatches on is deprecated by design. protected Flowable runAsyncImpl(InvocationContext invocationContext) { List subAgents = subAgents(); if (subAgents == null || subAgents.isEmpty()) { return Flowable.empty(); } + if (invocationContext.isLegacyResumability()) { + return runAsyncLegacyResumption(invocationContext, subAgents); + } + if (!invocationContext.isResumable()) { return Flowable.fromIterable(subAgents) .concatMap(subAgent -> subAgent.runAsync(invocationContext)) @@ -149,16 +156,33 @@ protected Flowable runAsyncImpl(InvocationContext invocationContext) { .takeUntil(LoopAgent::hasEscalateAction); } - // Resumable: stop looping once a sub-agent emits a pending long-running call (e.g. HITL), - // matching Python ADK v1 and avoiding a runaway loop. The current sub-agent still finishes; - // resuming into the paused iteration needs persisted state (future work). + // Resumable: checkpoint before each sub-agent, pause (not end) on a long-running call. + return Flowable.defer( + () -> { + Map state = invocationContext.agentStates().get(name()); + int startTimesLooped = WorkflowAgentStates.timesLooped(state); + // Unlike a sequence, a loop has no finished state to restore into, so it restarts at 0. + int resumeIndex = WorkflowAgentStates.resumeIndex(state, subAgents, logger); + int startIndex = resumeIndex == WorkflowAgentStates.NO_SUB_AGENT_NAMED ? 0 : resumeIndex; + LoopState loopState = new LoopState(state != null, startTimesLooped); + return runLoopIteration(invocationContext, subAgents, startIndex, loopState); + }); + } + + /** + * Runs the loop under the deprecated legacy resumption flow. Frozen copy of the behavior + * resumability had before durable checkpoints: stop looping once a sub-agent emits a pending + * long-running call, with no state read or written, so a resume cannot re-enter that iteration. + */ + private Flowable runAsyncLegacyResumption( + InvocationContext invocationContext, List subAgents) { AtomicBoolean paused = new AtomicBoolean(false); AtomicInteger timesLooped = new AtomicInteger(0); return Flowable.fromIterable(subAgents) .concatMap( subAgent -> paused.get() - ? Flowable.empty() + ? Flowable.empty() : subAgent .runAsync(invocationContext) .doOnNext( @@ -174,6 +198,125 @@ protected Flowable runAsyncImpl(InvocationContext invocationContext) { .takeUntil(LoopAgent::hasEscalateAction); } + /** Mutable state shared across the iterations of one resumable {@link LoopAgent} run. */ + private static final class LoopState { + /** True until the sub-agent being resumed into has run; that sub-agent skips its checkpoint. */ + final AtomicBoolean resuming; + + final AtomicInteger timesLooped; + final AtomicBoolean shouldExit = new AtomicBoolean(false); + final AtomicBoolean paused = new AtomicBoolean(false); + + /** Set once the loop must stop repeating, whether it closed, escalated or paused. */ + final AtomicBoolean finished = new AtomicBoolean(false); + + LoopState(boolean resuming, int timesLooped) { + this.resuming = new AtomicBoolean(resuming); + this.timesLooped = new AtomicInteger(timesLooped); + } + } + + /** + * Runs one loop iteration over the sub-agents from {@code startIndex}, then either recurses for + * the next iteration or terminates (emitting end-of-agent unless paused). + * + *

Only the previous sub-agent is checked for a silent pause, which under a {@link + * ParallelAgent} sits on a branch this agent cannot see: a call still unanswered inside the + * sub-agent about to run is what the resume is for, and must not stall it. The first sub-agent + * this run enters is re-entered unchecked, so one that pauses without emitting an event must be + * idempotent on re-entry. + */ + private Flowable runLoopIteration( + InvocationContext context, + List subAgents, + int startIndex, + LoopState state) { + AtomicInteger nextStart = new AtomicInteger(startIndex); + // Repeated rather than recursive: recursion nested one operator layer per iteration, which is + // unbounded when maxIterations is. + return Flowable.defer(() -> oneIteration(context, subAgents, nextStart.getAndSet(0), state)) + .repeatUntil(state.finished::get); + } + + /** + * Runs the sub-agents once from {@code startIndex}, marking {@code state} finished at the end. + */ + private Flowable oneIteration( + InvocationContext context, + List subAgents, + int startIndex, + LoopState state) { + return Flowable.defer( + () -> { + // Checked here so a resume already at the cap stops, not just a completed iteration. + if (maxIterations != null && state.timesLooped.get() >= maxIterations) { + state.finished.set(true); + return endOfAgentAndRecord(context); + } + AtomicReference ranLast = new AtomicReference<>(null); + Flowable iteration = + Flowable.fromIterable(subAgents.subList(startIndex, subAgents.size())) + .concatMap( + subAgent -> + Flowable.defer( + () -> { + if (state.shouldExit.get() || state.paused.get()) { + return Flowable.empty(); + } + // Catches a silent pause by the previous sub-agent. + BaseAgent previous = ranLast.get(); + if (previous != null + && context.hasUnansweredLongRunningCallIn(previous)) { + state.paused.set(true); + return Flowable.empty(); + } + ranLast.set(subAgent); + Flowable checkpoint = Flowable.empty(); + if (!state.resuming.getAndSet(false)) { + ImmutableMap subState = + ImmutableMap.of( + WorkflowAgentStates.CURRENT_SUB_AGENT, subAgent.name(), + WorkflowAgentStates.TIMES_LOOPED, + state.timesLooped.get()); + checkpoint = checkpointAndRecord(context, subState); + } + Flowable run = + subAgent + .runAsync(context) + .doOnNext( + event -> { + if (hasEscalateAction(event)) { + state.shouldExit.set(true); + } + if (context.shouldPauseInvocation(event)) { + state.paused.set(true); + } + }); + return checkpoint.concatWith(run); + })); + return iteration.concatWith( + Flowable.defer( + () -> { + // An escalation exit still closes the loop, as in Python. + if (state.paused.get() + || (!state.shouldExit.get() + && context.hasUnansweredLongRunningCallIn(this))) { + state.finished.set(true); + return Flowable.empty(); + } + // Anything but a pause completes the iteration, so count and clear state. + state.timesLooped.incrementAndGet(); + context.resetSubAgentStates(name()); + if (state.shouldExit.get()) { + state.finished.set(true); + return endOfAgentAndRecord(context); + } + // resuming is already false here, so the next iteration checkpoints normally. + return Flowable.empty(); + })); + }); + } + @Override protected Flowable runLiveImpl(InvocationContext invocationContext) { return Flowable.error( diff --git a/core/src/main/java/com/google/adk/agents/ParallelAgent.java b/core/src/main/java/com/google/adk/agents/ParallelAgent.java index d4066ed98..e78d47796 100644 --- a/core/src/main/java/com/google/adk/agents/ParallelAgent.java +++ b/core/src/main/java/com/google/adk/agents/ParallelAgent.java @@ -19,6 +19,7 @@ import com.google.adk.agents.ConfigAgentUtils.ConfigurationException; import com.google.adk.events.Event; +import com.google.common.collect.ImmutableMap; import com.google.common.collect.ImmutableSet; import com.google.errorprone.annotations.CanIgnoreReturnValue; import io.reactivex.rxjava3.core.Flowable; @@ -28,6 +29,7 @@ import java.util.List; import java.util.Objects; import java.util.Set; +import java.util.concurrent.atomic.AtomicBoolean; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -184,13 +186,76 @@ protected Flowable runAsyncImpl(InvocationContext invocationContext) { .filter(Objects::nonNull) .collect(ImmutableSet.toImmutableSet()); - var updatedInvocationContext = setBranchForCurrentAgent(this, invocationContext); - List> agentFlowables = new ArrayList<>(); - for (BaseAgent subAgent : currentSubAgents) { - agentFlowables.add(subAgent.runAsync(updatedInvocationContext).subscribeOn(scheduler)); + if (!invocationContext.isResumable()) { + var updatedInvocationContext = setBranchForCurrentAgent(this, invocationContext); + List> agentFlowables = new ArrayList<>(); + for (BaseAgent subAgent : currentSubAgents) { + agentFlowables.add(subAgent.runAsync(updatedInvocationContext).subscribeOn(scheduler)); + } + return Flowable.merge(agentFlowables) + .takeUntil((Event event) -> asksThisAgentToExit(event, directSubAgentNames)); } - return Flowable.merge(agentFlowables) - .takeUntil((Event event) -> asksThisAgentToExit(event, directSubAgentNames)); + + // Resumable: skip finished branches, checkpoint the start, pause instead of ending. + return Flowable.defer( + () -> { + List activeSubAgents = new ArrayList<>(); + for (BaseAgent subAgent : currentSubAgents) { + if (!invocationContext.endOfAgents().getOrDefault(subAgent.name(), false)) { + activeSubAgents.add(subAgent); + } + } + + Flowable initialCheckpoint = Flowable.empty(); + if (!invocationContext.agentStates().containsKey(name())) { + initialCheckpoint = checkpointAndRecord(invocationContext, ImmutableMap.of()); + } + + var updatedInvocationContext = setBranchForCurrentAgent(this, invocationContext); + AtomicBoolean paused = new AtomicBoolean(false); + AtomicBoolean escalated = new AtomicBoolean(false); + List> agentFlowables = new ArrayList<>(); + for (BaseAgent subAgent : activeSubAgents) { + agentFlowables.add( + subAgent + .runAsync(updatedInvocationContext) + .subscribeOn(scheduler) + .doOnNext( + event -> { + if (invocationContext.shouldPauseInvocation(event)) { + paused.set(true); + } + if (asksThisAgentToExit(event, directSubAgentNames)) { + escalated.set(true); + } + })); + } + Flowable merged = + Flowable.merge(agentFlowables) + .takeUntil((Event event) -> asksThisAgentToExit(event, directSubAgentNames)); + + return initialCheckpoint + .concatWith(merged) + .concatWith( + Flowable.defer( + () -> { + if (paused.get()) { + return Flowable.empty(); + } + // A sub-agent that never records endOfAgent keeps this agent from ending. + boolean allEnded = + activeSubAgents.stream() + .allMatch( + a -> + invocationContext + .endOfAgents() + .getOrDefault(a.name(), false)); + if (escalated.get() || allEnded) { + return endOfAgentAndRecord(invocationContext); + } + return Flowable.empty(); + })); + }); } /** diff --git a/core/src/main/java/com/google/adk/agents/SequentialAgent.java b/core/src/main/java/com/google/adk/agents/SequentialAgent.java index 963c3d109..7ed0d90a7 100644 --- a/core/src/main/java/com/google/adk/agents/SequentialAgent.java +++ b/core/src/main/java/com/google/adk/agents/SequentialAgent.java @@ -17,9 +17,12 @@ import com.google.adk.agents.ConfigAgentUtils.ConfigurationException; import com.google.adk.events.Event; +import com.google.common.collect.ImmutableMap; import io.reactivex.rxjava3.core.Flowable; import java.util.List; +import java.util.Map; import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicReference; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -77,7 +80,6 @@ public static class Builder extends BaseAgent.Builder { @Override public SequentialAgent build() { - // TODO(b/410859954): Add validation for required fields like name. return new SequentialAgent( name, description, subAgents, beforeAgentCallback, afterAgentCallback); } @@ -90,24 +92,38 @@ public static Builder builder() { /** * Runs sub-agents sequentially. * - *

When resumability is enabled, on resume execution fast-forwards to the sub-agent being - * resumed (completed ones are not re-run) and pauses on a pending long-running call; when - * disabled, sub-agents simply run in order (matches Python ADK v1 with resumability off). - * Temporary, event-based. + *

Three modes: with resumability on, a resume fast-forwards to the checkpointed sub-agent and + * pauses on a long-running call; with the deprecated shim on, the same happens but the resume + * point is reconstructed from history and nothing is checkpointed; with neither, sub-agents just + * run in order. * * @param invocationContext Invocation context. * @return Flowable emitting events from sub-agents. */ @Override + @SuppressWarnings("deprecation") // The shim it dispatches on is deprecated by design. protected Flowable runAsyncImpl(InvocationContext invocationContext) { List subAgents = subAgents(); if (subAgents.isEmpty()) { return Flowable.empty(); } - if (!invocationContext.isResumable()) { - return Flowable.fromIterable(subAgents) - .concatMap(subAgent -> subAgent.runAsync(invocationContext)); + if (invocationContext.isResumable()) { + return runAsyncResumable(invocationContext, subAgents); } + if (invocationContext.isLegacyResumability()) { + return runAsyncLegacyResumption(invocationContext, subAgents); + } + return Flowable.fromIterable(subAgents) + .concatMap(subAgent -> subAgent.runAsync(invocationContext)); + } + + /** + * Runs sub-agents under the deprecated legacy resumption flow, reconstructing the resume point + * from session events. Frozen copy of the behavior resumability had before durable checkpoints: + * no state is read or written, so only history decides where the sequence restarts. + */ + private Flowable runAsyncLegacyResumption( + InvocationContext invocationContext, List subAgents) { int startIndex = WorkflowAgentResumption.resumeSubAgentIndex(invocationContext, subAgents).orElse(0); AtomicBoolean paused = new AtomicBoolean(false); @@ -126,6 +142,76 @@ protected Flowable runAsyncImpl(InvocationContext invocationContext) { })); } + /** + * Runs sub-agents under durable resumability, matching Python ADK: checkpoint each sub-agent + * before it runs, fast-forward to the checkpoint on resume, and pause (without ending) on a + * long-running call. + * + *

Only the previous sub-agent is checked for a silent pause: a call still unanswered inside + * the sub-agent about to run is what the resume is for, and must not stall it. The first + * sub-agent this run enters is re-entered unchecked, so one that pauses without emitting an event + * must be idempotent on re-entry. + */ + private Flowable runAsyncResumable( + InvocationContext invocationContext, List subAgents) { + // Deferred so each subscription reads the checkpoint and builds its own mutable state. + return Flowable.defer( + () -> { + Map state = invocationContext.agentStates().get(name()); + int resumeIndex = WorkflowAgentStates.resumeIndex(state, subAgents, logger); + // A checkpoint naming no sub-agent means the sequence already finished. + int startIndex = + resumeIndex == WorkflowAgentStates.NO_SUB_AGENT_NAMED + ? subAgents.size() + : resumeIndex; + AtomicBoolean paused = new AtomicBoolean(false); + AtomicBoolean resuming = new AtomicBoolean(state != null); + AtomicReference ranLast = new AtomicReference<>(null); + return Flowable.fromIterable(subAgents.subList(startIndex, subAgents.size())) + .concatMap( + subAgent -> + Flowable.defer( + () -> { + // Catches a previous sub-agent that paused without emitting an event. + BaseAgent previous = ranLast.get(); + if (paused.get() + || (previous != null + && invocationContext.hasUnansweredLongRunningCallIn( + previous))) { + paused.set(true); + return Flowable.empty(); + } + ranLast.set(subAgent); + Flowable checkpoint = Flowable.empty(); + if (!resuming.getAndSet(false)) { + ImmutableMap subState = + ImmutableMap.of( + WorkflowAgentStates.CURRENT_SUB_AGENT, subAgent.name()); + checkpoint = checkpointAndRecord(invocationContext, subState); + } + Flowable run = + subAgent + .runAsync(invocationContext) + .doOnNext( + event -> { + if (invocationContext.shouldPauseInvocation(event)) { + paused.set(true); + } + }); + return checkpoint.concatWith(run); + })) + .concatWith( + Flowable.defer( + () -> { + if (paused.get() + || invocationContext.hasUnansweredLongRunningCallIn(this)) { + return Flowable.empty(); + } + return endOfAgentAndRecord(invocationContext); + })); + }); + } + /** * Runs sub-agents sequentially in live mode. * diff --git a/core/src/main/java/com/google/adk/agents/WorkflowAgentResumption.java b/core/src/main/java/com/google/adk/agents/WorkflowAgentResumption.java index 2bff47803..811be0c89 100644 --- a/core/src/main/java/com/google/adk/agents/WorkflowAgentResumption.java +++ b/core/src/main/java/com/google/adk/agents/WorkflowAgentResumption.java @@ -22,14 +22,15 @@ import java.util.Optional; /** - * Helpers for resuming workflow agents from session events. Temporary until session resumption - * (persisted agent state) is available. + * Legacy-resumption helper: reconstructs a workflow agent's resume point from session events. Used + * only by the legacy flow, which writes no durable agent state; the resumable flow resumes from its + * checkpoints instead, as Python ADK does. */ final class WorkflowAgentResumption { /** - * Index of the direct sub-agent whose subtree authored the call the latest event resumes, or - * empty when not resuming into this workflow. + * Index of the direct sub-agent whose subtree authored the call being resumed, or empty when not + * resuming into this workflow. */ static Optional resumeSubAgentIndex( InvocationContext invocationContext, List subAgents) { @@ -50,6 +51,8 @@ static Optional resumeSubAgentIndex( /** * Whether the event emits a long-running call still awaiting a response (e.g. a HITL request). + * The legacy flow pauses on this rather than on {@link InvocationContext#shouldPauseInvocation}, + * which the resumable flow uses. */ static boolean hasPendingLongRunningCall(Event event) { return Functions.hasPendingLongRunningCall(event); diff --git a/core/src/main/java/com/google/adk/agents/WorkflowAgentStates.java b/core/src/main/java/com/google/adk/agents/WorkflowAgentStates.java new file mode 100644 index 000000000..60eb936ec --- /dev/null +++ b/core/src/main/java/com/google/adk/agents/WorkflowAgentStates.java @@ -0,0 +1,103 @@ +/* + * Copyright 2025 Google LLC + * + * Licensed 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 com.google.adk.agents; + +import java.util.List; +import java.util.Map; +import org.jspecify.annotations.Nullable; +import org.slf4j.Logger; + +/** + * Wire-format keys and helpers for workflow-agent resumability checkpoints. The keys match Python + * and Kotlin ADK so persisted state is portable across languages. + */ +final class WorkflowAgentStates { + + /** Key holding the name of the current/next sub-agent in a Sequential or Loop checkpoint. */ + static final String CURRENT_SUB_AGENT = "current_sub_agent"; + + /** Key holding the completed-iteration count in a Loop checkpoint. */ + static final String TIMES_LOOPED = "times_looped"; + + /** + * Returns the index of the sub-agent to resume from by name, or 0 when the name is null or (with + * a warning) when it is no longer present in the sub-agents list. + */ + static int findIndexForResumption( + List subAgents, @Nullable String agentName, Logger logger) { + if (agentName == null) { + return 0; + } + for (int i = 0; i < subAgents.size(); i++) { + if (agentName.equals(subAgents.get(i).name())) { + return i; + } + } + // Agent names are developer-assigned identifiers, not user data, so log the missing name. + logger.warn("Restored sub-agent '{}' not found; resuming from index 0.", agentName); + return 0; + } + + /** Returned by {@link #resumeIndex} when the restored checkpoint names no sub-agent. */ + static final int NO_SUB_AGENT_NAMED = -1; + + /** + * Returns the index of the sub-agent to resume from, or {@link #NO_SUB_AGENT_NAMED} when the + * checkpoint names none -- an absent or empty name, which Python treats as "already finished". + * Any other unreadable value throws rather than being read as finished, so a session written by + * another runtime cannot silently skip the whole workflow. + */ + static int resumeIndex( + @Nullable Map state, List subAgents, Logger logger) { + if (state == null) { + return 0; + } + if (!state.containsKey(CURRENT_SUB_AGENT)) { + return NO_SUB_AGENT_NAMED; + } + Object current = state.get(CURRENT_SUB_AGENT); + // Restarting the workflow would re-run side-effecting sub-agents, so fail the resume instead. + if (!(current instanceof String name)) { + throw new IllegalStateException(rejection(CURRENT_SUB_AGENT, current, "a string")); + } + return name.isEmpty() ? NO_SUB_AGENT_NAMED : findIndexForResumption(subAgents, name, logger); + } + + /** + * Returns the completed-iteration count from a Loop checkpoint, or 0 when it holds none. As with + * {@link #resumeIndex}, an unreadable value throws instead of silently restarting the count. + */ + static int timesLooped(@Nullable Map state) { + if (state == null || !state.containsKey(TIMES_LOOPED)) { + return 0; + } + Object value = state.get(TIMES_LOOPED); + if (!(value instanceof Number number)) { + throw new IllegalStateException(rejection(TIMES_LOOPED, value, "a number")); + } + return number.intValue(); + } + + /** Message for a checkpoint field that cannot be read; names the type, never the value. */ + private static String rejection(String field, @Nullable Object value, String expected) { + return String.format( + "Cannot resume: checkpoint field %s is %s, not %s.", + field, value == null ? "null" : "a " + value.getClass().getSimpleName(), expected); + } + + private WorkflowAgentStates() {} +} diff --git a/core/src/main/java/com/google/adk/apps/App.java b/core/src/main/java/com/google/adk/apps/App.java index 9120954b5..4d2d1b3d9 100644 --- a/core/src/main/java/com/google/adk/apps/App.java +++ b/core/src/main/java/com/google/adk/apps/App.java @@ -19,6 +19,7 @@ import com.google.adk.agents.BaseAgent; import com.google.adk.agents.ContextCacheConfig; import com.google.adk.agents.Role; +import com.google.adk.annotations.Experimental; import com.google.adk.plugins.Plugin; import com.google.adk.summarizer.EventsCompactionConfig; import com.google.common.collect.ImmutableList; @@ -35,7 +36,6 @@ * and communication across all agents in the hierarchy. The {@code plugins} are application-wide * components that provide shared capabilities and services to the entire system. */ -@SuppressWarnings("deprecation") // Plumbs the deprecated ResumabilityConfig. public class App { private static final Pattern IDENTIFIER_PATTERN = Pattern.compile("[a-zA-Z_][a-zA-Z0-9_]*"); @@ -83,6 +83,7 @@ public ContextCacheConfig contextCacheConfig() { return contextCacheConfig; } + @Experimental public @Nullable ResumabilityConfig resumabilityConfig() { return resumabilityConfig; } @@ -132,14 +133,8 @@ public Builder contextCacheConfig(ContextCacheConfig contextCacheConfig) { return this; } - /** - * Sets the app resumability config. - * - * @deprecated See {@link ResumabilityConfig}: partial feature, full resumability not yet - * available. - */ @CanIgnoreReturnValue - @Deprecated + @Experimental public Builder resumabilityConfig(ResumabilityConfig resumabilityConfig) { this.resumabilityConfig = resumabilityConfig; return this; diff --git a/core/src/main/java/com/google/adk/apps/ResumabilityConfig.java b/core/src/main/java/com/google/adk/apps/ResumabilityConfig.java index 590d9177a..b19e04b82 100644 --- a/core/src/main/java/com/google/adk/apps/ResumabilityConfig.java +++ b/core/src/main/java/com/google/adk/apps/ResumabilityConfig.java @@ -18,40 +18,40 @@ import static com.google.common.base.Preconditions.checkArgument; +import com.google.adk.annotations.Experimental; import com.google.auto.value.AutoValue; import com.google.errorprone.annotations.CanIgnoreReturnValue; /** - * App resumability config, mirroring Python ADK v1's {@code ResumabilityConfig}: pause on a - * long-running call and resume from the last event. Applies to all agents in the app. + * Resumability config for every agent in an app: pause on a long-running call and resume from the + * last event. * - *

The two flags select the same resumption behavior and are mutually exclusive: set {@link - * #isResumable()}, or the deprecated shim, but not both. + *

The two flags select different resumption flows and are mutually exclusive: {@link + * #isResumable()} checkpoints agent state and resumes from it, {@link + * #isPlainTextContinuationAutoResume()} selects the deprecated legacy flow, and with neither set + * the app never pauses. * - * @deprecated Partial feature: only event-reconstruction-based pause/resume for {@code - * SequentialAgent} is implemented. Full session resumability (persisted agent state, durable - * resume, other workflow agents) is not yet available. Forward-compatible: the same config will - * drive full resumability once it lands. + *

Experimental: resume is best-effort and at-least-once, so a resuming tool must be idempotent + * and in-memory state is lost on resumption. */ -@Deprecated +@Experimental @AutoValue public abstract class ResumabilityConfig { - /** Whether the app supports agent resumption. */ + /** Whether the app supports agent resumption, checkpointing agent state as it runs. */ public abstract boolean isResumable(); /** * Whether a plain-text {@code runAsync} continuation -- a user message that is not a function - * response -- resumes the last unfinished invocation instead of starting a new one. Off by - * default, matching Python ADK, where a plain-text {@code runAsync} always starts a new - * invocation and a paused invocation is resumed explicitly. + * response -- resumes the last unfinished invocation instead of starting a new one; off by + * default. * - *

Selects the same resumption behavior as {@link #isResumable()}, with which it is mutually - * exclusive; it differs only in also resuming on a plain-text continuation. + *

Selects the legacy resumption flow, which reconstructs the resume point from session events + * and never persists agent state; mutually exclusive with {@link #isResumable()}. * - * @deprecated Back-compat shim for callers that deliver a resume as a plain-text turn. Migrate to - * {@code Runner.runAsync(userId, sessionId, invocationId, message, runConfig, stateDelta)} - * (or send a function response to the paused call) and set {@link #isResumable()} instead; + * @deprecated Back-compat shim for callers that deliver a resume as a plain-text turn; migrate to + * {@code Runner.runAsync(userId, sessionId, invocationId, message, runConfig, stateDelta)} or + * answer the paused call with a function response, and set {@link #isResumable()} instead; * this flag will be removed. */ @Deprecated @@ -83,16 +83,15 @@ public abstract static class Builder { /** * Builds the config, rejecting a combination of flags that has no defined behavior. * - * @throws IllegalArgumentException if both resumability and the deprecated shim are set; they - * select the same behavior, so exactly one may be enabled. + * @throws IllegalArgumentException if both resumability and the legacy shim are enabled; they + * select different resumption flows, so exactly one may be set. */ - @SuppressWarnings("deprecation") // Validating the deprecated shim against the supported flag. public ResumabilityConfig build() { ResumabilityConfig config = autoBuild(); checkArgument( !(config.isResumable() && config.isPlainTextContinuationAutoResume()), "resumable and plainTextContinuationAutoResume are mutually exclusive: set resumable for" - + " the supported flag, or the deprecated shim, but not both."); + + " the supported flow, or the deprecated shim for the legacy one."); return config; } } diff --git a/core/src/main/java/com/google/adk/events/Event.java b/core/src/main/java/com/google/adk/events/Event.java index c62df0985..40556a53f 100644 --- a/core/src/main/java/com/google/adk/events/Event.java +++ b/core/src/main/java/com/google/adk/events/Event.java @@ -42,8 +42,6 @@ import java.util.UUID; import org.jspecify.annotations.Nullable; -// TODO - b/413761119 update Agent.java when resolved. - /** Represents an event in a session. */ @JsonDeserialize(builder = Event.Builder.class) public class Event extends JsonBaseModel { @@ -174,6 +172,12 @@ public void setErrorCode(@Nullable FinishReason errorCode) { this.errorCode = errorCode; } + /** + * Sets the finish reason from an {@link Optional}. + * + * @deprecated Use {@link #setFinishReason(FinishReason)} with {@code finishReason.orElse(null)} + * instead. + */ @Deprecated @SuppressWarnings("checkstyle:IllegalType") public void setFinishReason(Optional finishReason) { diff --git a/core/src/main/java/com/google/adk/events/EventActions.java b/core/src/main/java/com/google/adk/events/EventActions.java index 743e42c4e..964733705 100644 --- a/core/src/main/java/com/google/adk/events/EventActions.java +++ b/core/src/main/java/com/google/adk/events/EventActions.java @@ -31,7 +31,6 @@ import org.jspecify.annotations.Nullable; /** Represents the actions attached to an event. */ -// TODO - b/414081262 make json wire camelCase @JsonDeserialize(builder = EventActions.Builder.class) public class EventActions extends JsonBaseModel { @@ -47,6 +46,7 @@ public class EventActions extends JsonBaseModel { private @Nullable Map agentState; private @Nullable EventCompaction compaction; private @Nullable Object setModelResponse; + private @Nullable String rewindBeforeInvocationId; /** Default constructor for Jackson. */ public EventActions() { @@ -71,6 +71,7 @@ private EventActions(Builder builder) { this.agentState = builder.agentState; this.compaction = builder.compaction; this.setModelResponse = builder.setModelResponse; + this.rewindBeforeInvocationId = builder.rewindBeforeInvocationId; } @JsonProperty("skipSummarization") @@ -91,7 +92,13 @@ public Map stateDelta() { return stateDelta; } - @Deprecated // Use stateDelta() and removeStateByKey() instead. + /** + * Replaces the state delta. + * + * @deprecated Update the map returned by {@link #stateDelta()} instead, and use {@link + * #removeStateByKey(String)} to remove a key. + */ + @Deprecated public void setStateDelta(ConcurrentMap stateDelta) { this.stateDelta = stateDelta; } @@ -231,6 +238,20 @@ public void setSetModelResponse(@Nullable Object setModelResponse) { this.setModelResponse = setModelResponse; } + /** + * On a rewind event, the ID of the invocation to rewind before: the history from that invocation + * up to the rewind event is dropped. ADK Java does not apply rewinds itself; the field exists for + * runtimes that do, such as ADK Kotlin. + */ + @JsonProperty("rewindBeforeInvocationId") + public Optional rewindBeforeInvocationId() { + return Optional.ofNullable(rewindBeforeInvocationId); + } + + public void setRewindBeforeInvocationId(@Nullable String rewindBeforeInvocationId) { + this.rewindBeforeInvocationId = rewindBeforeInvocationId; + } + public static Builder builder() { return new Builder(); } @@ -258,7 +279,8 @@ public boolean equals(Object o) { && (endOfAgent == that.endOfAgent) && Objects.equals(agentState, that.agentState) && Objects.equals(compaction, that.compaction) - && Objects.equals(setModelResponse, that.setModelResponse); + && Objects.equals(setModelResponse, that.setModelResponse) + && Objects.equals(rewindBeforeInvocationId, that.rewindBeforeInvocationId); } @Override @@ -275,7 +297,8 @@ public int hashCode() { endOfAgent, agentState, compaction, - setModelResponse); + setModelResponse, + rewindBeforeInvocationId); } /** Builder for {@link EventActions}. */ @@ -292,6 +315,7 @@ public static class Builder { private @Nullable Map agentState; private @Nullable EventCompaction compaction; private @Nullable Object setModelResponse; + private @Nullable String rewindBeforeInvocationId; public Builder() { this.stateDelta = new ConcurrentHashMap<>(); @@ -315,6 +339,7 @@ private Builder(EventActions eventActions) { this.agentState = eventActions.agentState; this.compaction = eventActions.compaction; this.setModelResponse = eventActions.setModelResponse; + this.rewindBeforeInvocationId = eventActions.rewindBeforeInvocationId; } @CanIgnoreReturnValue @@ -435,6 +460,13 @@ public Builder setModelResponse(@Nullable Object value) { return this; } + @CanIgnoreReturnValue + @JsonProperty("rewindBeforeInvocationId") + public Builder rewindBeforeInvocationId(@Nullable String value) { + this.rewindBeforeInvocationId = value; + return this; + } + @CanIgnoreReturnValue public Builder merge(EventActions other) { other.skipSummarization().ifPresent(this::skipSummarization); @@ -449,6 +481,7 @@ public Builder merge(EventActions other) { other.agentState().ifPresent(this::agentState); other.compaction().ifPresent(this::compaction); other.setModelResponse().ifPresent(this::setModelResponse); + other.rewindBeforeInvocationId().ifPresent(this::rewindBeforeInvocationId); return this; } diff --git a/core/src/main/java/com/google/adk/flows/llmflows/AgentTransfer.java b/core/src/main/java/com/google/adk/flows/llmflows/AgentTransfer.java index 0a0da8761..62ff664b5 100644 --- a/core/src/main/java/com/google/adk/flows/llmflows/AgentTransfer.java +++ b/core/src/main/java/com/google/adk/flows/llmflows/AgentTransfer.java @@ -44,16 +44,15 @@ public Single processRequest( "Base agent in InvocationContext is not an instance of Agent."); } - List transferTargets = getTransferTargets(agent); - if (transferTargets.isEmpty()) { + List targets = transferTargets(agent); + if (targets.isEmpty()) { return Single.just( RequestProcessor.RequestProcessingResult.create(request, ImmutableList.of())); } LlmRequest.Builder builder = request.toBuilder() - .appendInstructions( - ImmutableList.of(buildTargetAgentsInstructions(agent, transferTargets))); + .appendInstructions(ImmutableList.of(buildTargetAgentsInstructions(agent, targets))); FunctionTool agentTransferTool = createTransferToAgentTool(); agentTransferTool.processLlmRequest(builder, ToolContext.builder(context).build()); @@ -119,30 +118,29 @@ private String buildTargetAgentsInstructions(LlmAgent agent, List tra return sb.toString(); } - /** Returns valid transfer targets: sub-agents, parent, and peers (if allowed). */ - private List getTransferTargets(LlmAgent agent) { - List transferTargets = new ArrayList<>(); - transferTargets.addAll(agent.subAgents()); // Add all sub-agents - + /** + * Returns the agents {@code agent} may transfer to: its sub-agents, plus its parent and peers + * unless disallowed. Transfer to a parent or peer needs an {@link LlmAgent} parent. Order is + * load-bearing: it is the order the target list is rendered to the model in. {@code LlmAgent} + * keeps the same rule for its own use, since it cannot reach this one. + */ + static ImmutableList transferTargets(LlmAgent agent) { + List targets = new ArrayList<>(agent.subAgents()); BaseAgent parent = agent.parentAgent(); - // Agents eligible to transfer must have an LLM-based agent parent. if (!(parent instanceof LlmAgent)) { - return transferTargets; + return ImmutableList.copyOf(targets); } - if (!agent.disallowTransferToParent()) { - transferTargets.add(parent); + targets.add(parent); } - if (!agent.disallowTransferToPeers()) { - for (BaseAgent peerAgent : parent.subAgents()) { - if (!peerAgent.name().equals(agent.name())) { - transferTargets.add(peerAgent); + for (BaseAgent peer : parent.subAgents()) { + if (!peer.name().equals(agent.name())) { + targets.add(peer); } } } - - return transferTargets; + return ImmutableList.copyOf(targets); } @Schema( diff --git a/core/src/main/java/com/google/adk/flows/llmflows/BaseLlmFlow.java b/core/src/main/java/com/google/adk/flows/llmflows/BaseLlmFlow.java index 91cc225f2..f255bbe03 100644 --- a/core/src/main/java/com/google/adk/flows/llmflows/BaseLlmFlow.java +++ b/core/src/main/java/com/google/adk/flows/llmflows/BaseLlmFlow.java @@ -60,6 +60,7 @@ import java.util.List; import java.util.Optional; import java.util.Set; +import java.util.concurrent.atomic.AtomicBoolean; import java.util.concurrent.atomic.AtomicReference; import java.util.function.BiFunction; import org.slf4j.Logger; @@ -432,7 +433,14 @@ private Flowable runOneStep(Context spanContext, InvocationContext contex return Flowable.defer( () -> { + AtomicBoolean preprocessAnswered = new AtomicBoolean(false); return preprocess(context, llmRequestRef) + .doOnNext( + event -> { + if (!event.functionResponses().isEmpty() && event.finalResponse()) { + preprocessAnswered.set(true); + } + }) .concatWith( Flowable.defer( () -> { @@ -441,6 +449,30 @@ private Flowable runOneStep(Context spanContext, InvocationContext contex logger.debug("End invocation requested during preprocessing."); return Flowable.empty(); } + // A tool answered during preprocessing would be replayed below. + if (context.isResumable() && preprocessAnswered.get()) { + logger.debug("Preprocessing produced a final response; ending the step."); + return Flowable.empty(); + } + + // As in Python, decide before the model call: the branch may owe work. + StepResume.Decision resume = + StepResume.decide(context, llmRequestAfterPreprocess.tools()); + if (resume.action == StepResume.Action.PAUSE) { + logger.debug("Pausing the flow: a call is still unanswered."); + return Flowable.empty(); + } + if (resume.action == StepResume.Action.REPLAY_CALLS) { + logger.debug("Replaying function calls a previous run did not execute."); + // Re-id as Python's replay does; PersistBarrier skips an id-less event. + return runFunctionCalls( + context, + resume.replayEvent(), + llmRequestAfterPreprocess, + spanContext) + .map(event -> event.toBuilder().id(Event.generateEventId()).build()) + .concatMap(event -> followTransfer(event, context, spanContext)); + } try { context.incrementLlmCallsCount(); @@ -477,27 +509,7 @@ private Flowable runOneStep(Context spanContext, InvocationContext contex String newId = Event.generateEventId(); logger.debug("Resetting event ID from {} to {}", oldId, newId); event = event.toBuilder().id(newId).build(); - Flowable postProcessedEvents = Flowable.just(event); - if (event.actions().transferToAgent().isPresent()) { - String agentToTransfer = - event.actions().transferToAgent().get(); - BaseAgent rootAgent = context.agent().rootAgent(); - Optional nextAgent = - rootAgent.findAgent(agentToTransfer); - if (nextAgent.isEmpty()) { - logger.error("Agent not found: {}", agentToTransfer); - return postProcessedEvents.concatWith( - Flowable.error( - new IllegalStateException( - "Agent not found: " + agentToTransfer))); - } - return postProcessedEvents.concatWith( - nextAgent - .get() - .runAsync(context) - .compose(Tracing.withContext(spanContext))); - } - return postProcessedEvents; + return followTransfer(event, context, spanContext); }); })); }); @@ -521,6 +533,8 @@ private Flowable run( logger.debug("Ending flow execution because max steps reached."); return currentStepEvents; } + @SuppressWarnings("deprecation") // The shim it reports is deprecated by design. + boolean legacyResumption = invocationContext.isLegacyResumability(); return currentStepEvents.concatWith( currentStepEvents @@ -534,11 +548,8 @@ private Flowable run( "Ending flow execution based on final response, endInvocation action or" + " empty event list."); return Flowable.empty(); - } else if (invocationContext.isResumable() - && Functions.hasPendingLongRunningCall(eventList)) { - // When resumable, a pending long-running call (e.g. HITL) pauses the flow - // instead of calling the model again, matching Python ADK v1 and avoiding a - // runaway re-issue loop. The disabled path is unchanged. + } else if (legacyResumption && Functions.hasPendingLongRunningCall(eventList)) { + // The resumable flow decides before the model call instead, in StepResume. logger.debug("Pausing flow execution on a pending long-running call."); return Flowable.empty(); } else { @@ -758,30 +769,86 @@ private Flowable buildPostprocessingEvents( return processorEvents.concatWith(Flowable.just(modelResponseEvent)); } - Flowable functionEvents; + Flowable functionEvents = + runFunctionCalls(context, modelResponseEvent, llmRequest, parentContext); + + return processorEvents.concatWith(Flowable.just(modelResponseEvent)).concatWith(functionEvents); + } + + /** + * Emits {@code event}, followed by the transferred-to agent's run when it carries a transfer. + * Shared so a replayed transfer continues into its target the way a fresh one does. + */ + private Flowable followTransfer( + Event event, InvocationContext context, Context spanContext) { + Flowable self = Flowable.just(event); + if (event.actions().transferToAgent().isEmpty()) { + return self; + } + String agentToTransfer = event.actions().transferToAgent().get(); + Optional nextAgent = context.agent().rootAgent().findAgent(agentToTransfer); + if (nextAgent.isEmpty()) { + logger.error("Agent not found: {}", agentToTransfer); + return self.concatWith( + Flowable.error(new IllegalStateException("Agent not found: " + agentToTransfer))); + } + BaseAgent target = nextAgent.get(); + // Resumable only: restricting the default flow too would move behavior for existing callers. + if (context.isResumable() && context.agent() instanceof LlmAgent llmAgent) { + Optional declared = declaredTransferTarget(llmAgent, agentToTransfer); + if (declared.isEmpty()) { + return self.concatWith( + Flowable.error( + new IllegalStateException( + "Agent " + + llmAgent.name() + + " is not allowed to transfer to the requested agent."))); + } + target = declared.get(); + } + return self.concatWith(target.runAsync(context).compose(Tracing.withContext(spanContext))); + } + + /** + * Returns the agent {@code from} may transfer to under {@code name}, or empty when the transfer + * is not allowed. Resolves through the declared targets rather than the tree, as Python does, so + * a same-named agent on an unrelated branch cannot stand in for the declared one. + */ + private static Optional declaredTransferTarget(LlmAgent from, String name) { + if (name.equals(from.name())) { + return Optional.of(from); + } + return AgentTransfer.transferTargets(from).stream() + .filter(candidate -> candidate.name().equals(name)) + .findFirst(); + } + + /** + * Runs {@code callEvent}'s function calls and emits what follows from them: the tool-confirmation + * request, the function response, and any structured final response. Shared by the normal + * postprocessing path and by replaying calls a previous run never executed. + */ + private Flowable runFunctionCalls( + InvocationContext context, Event callEvent, LlmRequest llmRequest, Context parentContext) { try (Scope scope = parentContext.makeCurrent()) { Maybe maybeFunctionResponseEvent = context.runConfig().streamingMode() == StreamingMode.BIDI - ? Functions.handleFunctionCallsLive(context, modelResponseEvent, llmRequest.tools()) - : Functions.handleFunctionCalls(context, modelResponseEvent, llmRequest.tools()); - functionEvents = - maybeFunctionResponseEvent.flatMapPublisher( - functionResponseEvent -> { - Optional toolConfirmationEvent = - Functions.generateRequestConfirmationEvent( - context, modelResponseEvent, functionResponseEvent); - List events = new ArrayList<>(); - toolConfirmationEvent.ifPresent(events::add); - events.add(functionResponseEvent); - OutputSchema.getStructuredModelResponse(functionResponseEvent) - .ifPresent( - json -> - events.add(OutputSchema.createFinalModelResponseEvent(context, json))); - return Flowable.fromIterable(events); - }); + ? Functions.handleFunctionCallsLive(context, callEvent, llmRequest.tools()) + : Functions.handleFunctionCalls(context, callEvent, llmRequest.tools()); + return maybeFunctionResponseEvent.flatMapPublisher( + functionResponseEvent -> { + Optional toolConfirmationEvent = + Functions.generateRequestConfirmationEvent( + context, callEvent, functionResponseEvent); + List events = new ArrayList<>(); + toolConfirmationEvent.ifPresent(events::add); + events.add(functionResponseEvent); + OutputSchema.getStructuredModelResponse(functionResponseEvent) + .ifPresent( + json -> events.add(OutputSchema.createFinalModelResponseEvent(context, json))); + return Flowable.fromIterable(events); + }); } - - return processorEvents.concatWith(Flowable.just(modelResponseEvent)).concatWith(functionEvents); } /** diff --git a/core/src/main/java/com/google/adk/flows/llmflows/Contents.java b/core/src/main/java/com/google/adk/flows/llmflows/Contents.java index a80dfa09d..4778ebe27 100644 --- a/core/src/main/java/com/google/adk/flows/llmflows/Contents.java +++ b/core/src/main/java/com/google/adk/flows/llmflows/Contents.java @@ -521,8 +521,6 @@ private static List rearrangeEventsForLatestFunctionResponse(List return events; } - // TODO: b/412663475 - Handle parallel function calls within the same event. Currently, this - // throws an error. if (events.isEmpty() || Iterables.getLast(events).functionResponses().isEmpty()) { // No need to process if the list is empty or the last event is not a function response return events; diff --git a/core/src/main/java/com/google/adk/flows/llmflows/Functions.java b/core/src/main/java/com/google/adk/flows/llmflows/Functions.java index 3a1d45c2f..9e2accdc3 100644 --- a/core/src/main/java/com/google/adk/flows/llmflows/Functions.java +++ b/core/src/main/java/com/google/adk/flows/llmflows/Functions.java @@ -130,8 +130,6 @@ public static void populateClientFunctionCallId(Event modelResponseEvent) { } } - // TODO - b/413761119 add the remaining methods for function call id. - /** Handles standard, non-streaming function calls. */ public static Maybe handleFunctionCalls( InvocationContext invocationContext, Event functionCallEvent, Map tools) { @@ -457,9 +455,11 @@ public static boolean hasPendingLongRunningCall(Event event) { } /** - * Returns whether the last one or two events hold a pending long-running call, meaning a - * resumable flow should pause instead of calling the model again. Mirrors Python ADK v1's - * flow-level pause check on {@code events[-1]} and {@code events[-2]}. + * Returns whether either of the last two events emits a long-running call, meaning the legacy + * resumption flow should pause instead of calling the model again. Responses are not matched + * against calls: a long-running tool that returns a value in the same turn still pauses, because + * its call event is inside the window. The resumable flow does not use this -- it decides in + * {@code StepResume}, which does match responses. */ static boolean hasPendingLongRunningCall(List events) { int from = Math.max(0, events.size() - 2); diff --git a/core/src/main/java/com/google/adk/flows/llmflows/StepResume.java b/core/src/main/java/com/google/adk/flows/llmflows/StepResume.java new file mode 100644 index 000000000..35acd4ee8 --- /dev/null +++ b/core/src/main/java/com/google/adk/flows/llmflows/StepResume.java @@ -0,0 +1,241 @@ +/* + * Copyright 2026 Google LLC + * + * Licensed 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 com.google.adk.flows.llmflows; + +import com.google.adk.agents.InvocationContext; +import com.google.adk.events.Event; +import com.google.adk.tools.BaseTool; +import com.google.common.collect.ImmutableList; +import com.google.common.collect.ImmutableSet; +import com.google.genai.types.FunctionCall; +import com.google.genai.types.FunctionResponse; +import java.util.HashSet; +import java.util.List; +import java.util.Map; +import java.util.Set; +import org.jspecify.annotations.Nullable; + +/** + * Decides how a resumable flow's next step continues: carry on to the model, stop because a call is + * still unanswered, or run calls a previous run never executed. + * + *

Ports Python ADK's resume decision. Java has no per-call branch run ids, so the sub-branch + * case Python recognizes (a HITL answer returning against a branch the call opened) has no + * counterpart here and is treated as absent. + */ +final class StepResume { + + /** What the flow should do with the events it resumed from. */ + enum Action { + /** Nothing outstanding; proceed to the model call. */ + CONTINUE, + /** A call is still unanswered; stop without emitting anything. */ + PAUSE, + /** A call was never executed; run the calls on {@link Decision#event}. */ + REPLAY_CALLS + } + + /** The action to take, and the event it applies to. */ + static final class Decision { + final Action action; + private final @Nullable Event event; + + private Decision(Action action, @Nullable Event event) { + this.action = action; + this.event = event; + } + + /** The event whose calls to run; only a {@code REPLAY_CALLS} decision carries one. */ + Event replayEvent() { + if (event == null) { + throw new IllegalStateException(action + " decision carries no event to replay"); + } + return event; + } + } + + private static final Decision CONTINUE = new Decision(Action.CONTINUE, null); + private static final Decision PAUSE = new Decision(Action.PAUSE, null); + + /** + * Decides how the next step resumes. The branch's last event decides the case: content means the + * normal flow carries on, while a call that was never executed has to run first. + */ + static Decision decide(InvocationContext context, Map tools) { + if (!context.isResumable()) { + return CONTINUE; + } + ImmutableList events = + context.events(/* currentInvocation= */ true, /* currentBranch= */ true); + if (events.isEmpty()) { + return CONTINUE; + } + if (events.size() > 1) { + Decision decision = decideFromBranch(context, events, tools); + if (decision.action != Action.CONTINUE) { + return decision; + } + } + // Calls on the last event are unanswered by construction -- nothing follows them to answer. + Event last = events.get(events.size() - 1); + if (!last.partial().orElse(false) + && !last.functionCalls().isEmpty() + // A caller-supplied event would otherwise run a tool with no model turn behind it. + && context.agent().name().equals(last.author())) { + return new Decision(Action.REPLAY_CALLS, last); + } + return CONTINUE; + } + + /** The multi-event core: pause on an unanswered call, else replay an unexecuted one. */ + private static Decision decideFromBranch( + InvocationContext context, List events, Map tools) { + Event last = events.get(events.size() - 1); + boolean pausedByLast = context.shouldPauseInvocation(last); + if (!pausedByLast && pauseLeftCallsUnanswered(context, events)) { + return PAUSE; + } + + int callIdx = findTargetCallEventIndex(context, events, tools); + if (callIdx < 0) { + return pausedByLast ? PAUSE : CONTINUE; + } + Event callEvent = events.get(callIdx); + Set callNames = new HashSet<>(); + Set callIds = new HashSet<>(); + // An id-less call can never be answered, so it pauses; Python keeps a None in its id set. + boolean hasIdlessCall = false; + for (FunctionCall call : callEvent.functionCalls()) { + call.name().ifPresent(callNames::add); + call.id().ifPresent(callIds::add); + hasIdlessCall |= call.id().isEmpty(); + } + Set longRunningIds = new HashSet<>(); + for (int i = callIdx; i < events.size(); i++) { + longRunningIds.addAll(events.get(i).longRunningToolIds().orElse(ImmutableSet.of())); + } + callIds.addAll(longRunningIds); + + Set answeredIds = new HashSet<>(); + for (int i = callIdx + 1; i < events.size(); i++) { + for (FunctionResponse response : events.get(i).functionResponses()) { + response.id().ifPresent(answeredIds::add); + } + } + ImmutableList answers = answerEvent(events, callIdx, callIds, callNames); + + boolean lroUnanswered = !longRunningIds.isEmpty() && disjoint(longRunningIds, answeredIds); + boolean callUnanswered = + (!callIds.isEmpty() || hasIdlessCall) + && disjoint(callIds, answeredIds) + && answers.stream() + .noneMatch(response -> response.name().map(callNames::contains).orElse(false)); + if (lroUnanswered || callUnanswered) { + return PAUSE; + } + if (needsCallReplay(callNames, answers)) { + return new Decision(Action.REPLAY_CALLS, callEvent); + } + return pausedByLast ? PAUSE : CONTINUE; + } + + /** + * Whether a pause earlier in {@code events} is still waiting. Every event before the last counts: + * a long-running call followed by several responses leaves the pausing call further back than a + * two-event window can see. + */ + private static boolean pauseLeftCallsUnanswered(InvocationContext context, List events) { + Set awaited = new HashSet<>(); + for (int i = 0; i < events.size() - 1; i++) { + Event event = events.get(i); + if (!context.shouldPauseInvocation(event)) { + continue; + } + // Union, as in Python: every call on a pausing event plus its long-running ids. A different + // rule from shouldPauseInvocation's, which intersects the two. + for (FunctionCall call : event.functionCalls()) { + call.id().ifPresent(awaited::add); + } + awaited.addAll(event.longRunningToolIds().orElse(ImmutableSet.of())); + } + if (awaited.isEmpty()) { + return false; + } + Set answered = new HashSet<>(); + for (Event event : events) { + for (FunctionResponse response : event.functionResponses()) { + response.id().ifPresent(answered::add); + } + } + // A partially answered pause keeps waiting. + return !answered.containsAll(awaited); + } + + /** + * Index of the most recent event before the last that calls a tool this flow owns and the current + * agent authored, or -1 when there is none. + */ + private static int findTargetCallEventIndex( + InvocationContext context, List events, Map tools) { + String agentName = context.agent().name(); + for (int i = events.size() - 2; i >= 0; i--) { + Event event = events.get(i); + if (!agentName.equals(event.author())) { + continue; + } + boolean callsOwnedTool = + event.functionCalls().stream() + .anyMatch(call -> call.name().map(tools::containsKey).orElse(false)); + if (callsOwnedTool) { + return i; + } + } + return -1; + } + + /** The responses of the event answering the call, or of the last event when none does. */ + private static ImmutableList answerEvent( + List events, int callIdx, Set callIds, Set callNames) { + for (int i = events.size() - 1; i > callIdx; i--) { + for (FunctionResponse response : events.get(i).functionResponses()) { + boolean matchesId = response.id().map(callIds::contains).orElse(false); + boolean matchesName = + response.id().isEmpty() && response.name().map(callNames::contains).orElse(false); + if (matchesId || matchesName) { + return events.get(i).functionResponses(); + } + } + } + return events.get(events.size() - 1).functionResponses(); + } + + /** Whether the calls still need running: nothing answered them, or something else did. */ + private static boolean needsCallReplay(Set callNames, List answers) { + if (callNames.isEmpty()) { + return false; + } + return answers.isEmpty() + || answers.stream() + .anyMatch(response -> !response.name().map(callNames::contains).orElse(false)); + } + + private static boolean disjoint(Set left, Set right) { + return left.stream().noneMatch(right::contains); + } + + private StepResume() {} +} diff --git a/core/src/main/java/com/google/adk/plugins/agentanalytics/BigQuerySchema.java b/core/src/main/java/com/google/adk/plugins/agentanalytics/BigQuerySchema.java index bf50a26f1..2d03fe6dd 100644 --- a/core/src/main/java/com/google/adk/plugins/agentanalytics/BigQuerySchema.java +++ b/core/src/main/java/com/google/adk/plugins/agentanalytics/BigQuerySchema.java @@ -68,7 +68,6 @@ private BigQuerySchema() {} "{\"encoding\": \"WKT\"}")); /** Returns the BigQuery schema for the events table. */ - // TODO(b/491848381): Rely on the same schema defined for python plugin. public static Schema getEventsSchema() { return Schema.of( Field.newBuilder("timestamp", StandardSQLTypeName.TIMESTAMP) diff --git a/core/src/main/java/com/google/adk/runner/Runner.java b/core/src/main/java/com/google/adk/runner/Runner.java index 2642033a3..7ec2e9265 100644 --- a/core/src/main/java/com/google/adk/runner/Runner.java +++ b/core/src/main/java/com/google/adk/runner/Runner.java @@ -17,6 +17,9 @@ package com.google.adk.runner; import static com.google.common.base.Preconditions.checkArgument; +import static com.google.common.base.Preconditions.checkState; +import static com.google.common.collect.ImmutableList.toImmutableList; +import static java.util.stream.Collectors.toCollection; import com.google.adk.agents.ActiveStreamingTool; import com.google.adk.agents.BaseAgent; @@ -24,9 +27,11 @@ import com.google.adk.agents.InvocationContext; import com.google.adk.agents.LiveRequestQueue; import com.google.adk.agents.LlmAgent; +import com.google.adk.agents.ParallelAgent; import com.google.adk.agents.Role; import com.google.adk.agents.RunConfig; import com.google.adk.agents.SequentialAgent; +import com.google.adk.annotations.Experimental; import com.google.adk.apps.App; import com.google.adk.apps.ResumabilityConfig; import com.google.adk.artifacts.BaseArtifactService; @@ -56,6 +61,8 @@ import com.google.errorprone.annotations.CanIgnoreReturnValue; import com.google.genai.types.AudioTranscriptionConfig; import com.google.genai.types.Content; +import com.google.genai.types.FunctionCall; +import com.google.genai.types.FunctionResponse; import com.google.genai.types.Modality; import com.google.genai.types.Part; import io.opentelemetry.api.trace.Span; @@ -69,15 +76,17 @@ import java.util.ArrayList; import java.util.Arrays; import java.util.Collections; +import java.util.HashSet; import java.util.List; import java.util.Map; +import java.util.Objects; import java.util.Optional; +import java.util.Set; import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.ConcurrentMap; import org.jspecify.annotations.Nullable; /** The main class for the GenAI Agents runner. */ -@SuppressWarnings("deprecation") // Plumbs the deprecated ResumabilityConfig. public class Runner { private final BaseAgent agent; private final String appName; @@ -360,6 +369,22 @@ private Single createUserMessageEvent( InvocationContext invocationContext, boolean saveInputBlobsAsArtifacts, @Nullable Map stateDelta) { + return appendNewMessageToSession( + session, + newMessage, + invocationContext, + saveInputBlobsAsArtifacts, + stateDelta, + /* branch= */ null); + } + + private Single appendNewMessageToSession( + Session session, + Content newMessage, + InvocationContext invocationContext, + boolean saveInputBlobsAsArtifacts, + @Nullable Map stateDelta, + @Nullable String branch) { checkArgument(newMessage.parts().isPresent(), "No parts in the new_message."); Content messageToAppend = newMessage; @@ -393,6 +418,7 @@ private Single createUserMessageEvent( .id(Event.generateEventId()) .invocationId(invocationContext.invocationId()) .author(Role.USER) + .branch(branch) .content(messageToAppend); // Add state delta if provided @@ -450,6 +476,104 @@ public Flowable runAsync( this.runAsyncImpl(session, newMessage, runConfig, stateDelta))) .compose(Tracing.trace("invocation")); + return serializePerSession(sessionId, result); + } + + /** See {@link #runAsync(String, String, Content, RunConfig, Map)}. */ + public Flowable runAsync( + SessionKey sessionKey, + Content newMessage, + RunConfig runConfig, + @Nullable Map stateDelta) { + return runAsync(sessionKey.userId(), sessionKey.id(), newMessage, runConfig, stateDelta); + } + + /** See {@link #runAsync(String, String, Content, RunConfig, Map)}. */ + public Flowable runAsync(SessionKey sessionKey, Content newMessage, RunConfig runConfig) { + return runAsync(sessionKey, newMessage, runConfig, /* stateDelta= */ null); + } + + /** See {@link #runAsync(String, String, Content, RunConfig, Map)}. */ + public Flowable runAsync(SessionKey sessionKey, Content newMessage) { + return runAsync(sessionKey, newMessage, RunConfig.builder().build()); + } + + /** See {@link #runAsync(String, String, Content, RunConfig, Map)}. */ + public Flowable runAsync(String userId, String sessionId, Content newMessage) { + return runAsync(userId, sessionId, newMessage, RunConfig.builder().build()); + } + + /** + * Runs the agent, resuming an existing invocation instead of starting a new one. The invocation + * is resolved from {@code invocationId}, or from a function response carried by {@code + * newMessage}. Agent checkpoints are rehydrated from history and an invocation whose active agent + * already finished resolves to a no-op. + * + * @param userId the user id of the session. + * @param sessionId the session id. + * @param invocationId the invocation to resume; may be {@code null} when it can be inferred from + * {@code newMessage}. + * @param newMessage an optional message (typically a function response) to append before running. + * @param runConfig the run configuration. + * @param stateDelta optional state updates to merge into the session for this run. + * @return the events generated while resuming, or an empty stream when there is nothing to + * resume. + * @throws IllegalStateException if the app is not resumable. + * @throws IllegalArgumentException if the invocation cannot be resolved. + */ + @Experimental + public Flowable runAsync( + String userId, + String sessionId, + @Nullable String invocationId, + @Nullable Content newMessage, + RunConfig runConfig, + @Nullable Map stateDelta) { + checkState( + isResumable(), + "Resuming an invocation requires an App configured with a resumable ResumabilityConfig."); + Flowable result = + Flowable.defer( + () -> + this.sessionService + .getSession(appName, userId, sessionId, Optional.empty()) + .switchIfEmpty( + Single.defer( + () -> { + // Same get-or-create as a new run, as in Python. + if (runConfig.autoCreateSession()) { + return this.sessionService.createSession( + appName, userId, (Map) null, sessionId); + } + return Single.error( + new IllegalArgumentException( + String.format( + "Session not found: %s for user %s", + sessionId, userId))); + })) + .flatMapPublisher( + session -> + runResumableFromSession( + session, invocationId, newMessage, runConfig, stateDelta))) + .compose(Tracing.trace("invocation")); + + // Resuming reads session events around an append, so it shares the new run's per-session turn. + return serializePerSession(sessionId, result); + } + + /** Returns {@code content} with the user role applied when it carries none, as Python does. */ + private static @Nullable Content withUserRole(@Nullable Content content) { + if (content == null || content.role().filter(role -> !role.isEmpty()).isPresent()) { + return content; + } + return content.toBuilder().role(Role.USER).build(); + } + + /** + * Chains {@code result} behind any run already in flight for {@code sessionId}, so two runs never + * read and append the session's events concurrently. A null session id runs unsequenced. + */ + private Flowable serializePerSession(@Nullable String sessionId, Flowable result) { return Flowable.defer( () -> { if (sessionId == null) { @@ -480,30 +604,6 @@ public Flowable runAsync( }); } - /** See {@link #runAsync(String, String, Content, RunConfig, Map)}. */ - public Flowable runAsync( - SessionKey sessionKey, - Content newMessage, - RunConfig runConfig, - @Nullable Map stateDelta) { - return runAsync(sessionKey.userId(), sessionKey.id(), newMessage, runConfig, stateDelta); - } - - /** See {@link #runAsync(String, String, Content, RunConfig, Map)}. */ - public Flowable runAsync(SessionKey sessionKey, Content newMessage, RunConfig runConfig) { - return runAsync(sessionKey, newMessage, runConfig, /* stateDelta= */ null); - } - - /** See {@link #runAsync(String, String, Content, RunConfig, Map)}. */ - public Flowable runAsync(SessionKey sessionKey, Content newMessage) { - return runAsync(sessionKey, newMessage, RunConfig.builder().build()); - } - - /** See {@link #runAsync(String, String, Content, RunConfig, Map)}. */ - public Flowable runAsync(String userId, String sessionId, Content newMessage) { - return runAsync(userId, sessionId, newMessage, RunConfig.builder().build()); - } - /** * Runs the agent asynchronously using a provided Session object. * @@ -521,6 +621,19 @@ protected Flowable runAsyncImpl( Preconditions.checkNotNull(session, "session cannot be null"); Preconditions.checkNotNull(newMessage, "newMessage cannot be null"); Preconditions.checkNotNull(runConfig, "runConfig cannot be null"); + if (isResumable()) { + return runResumableFromSession( + session, /* providedInvocationId= */ null, newMessage, runConfig, stateDelta); + } + return runNewInvocation(session, newMessage, runConfig, stateDelta); + } + + /** Starts a brand-new invocation for {@code newMessage} (the default, non-resume flow). */ + private Flowable runNewInvocation( + Session session, + Content newMessage, + RunConfig runConfig, + @Nullable Map stateDelta) { return Flowable.defer( () -> { Context capturedContext = Context.current(); @@ -558,12 +671,7 @@ protected Flowable runAsyncImpl( userEvent -> runAgentForUserEvent(initialContext, session, userEvent, rootAgent) .compose(Tracing.withContext(capturedContext))) - .doOnError( - throwable -> - this.pluginManager - .runOnRunErrorCallback(initialContext, throwable) - .onErrorComplete() - .subscribe()); + .doOnError(throwable -> runOnRunError(initialContext, throwable)); }) .doOnError( throwable -> { @@ -621,27 +729,35 @@ private Flowable runAgentWithUpdatedSession( .userContent(event.content().orElseGet(Content::fromParts)) .build(); - // Call beforeRunCallback with updated session - Maybe beforeRunEvent = - this.pluginManager - .beforeRunCallback(contextWithUpdatedSession) - .map( - content -> - Event.builder() - .id(Event.generateEventId()) - .invocationId(contextWithUpdatedSession.invocationId()) - .author("model") - .content(content) - .build()); + // If beforeRunCallback returns content, emit it and skip agent. + Maybe beforeRunEvent = beforeRunEventFor(contextWithUpdatedSession); + Context capturedContext = Context.current(); + return executeAgentPipeline( + contextWithUpdatedSession, + updatedSession, + beforeRunEvent, + // TODO: remove this hack after deprecating runAsync with Session. + () -> copySessionStates(updatedSession, initialContext.session())) + .compose(Tracing.withContext(capturedContext)); + } + /** + * Runs {@code context.agent()} and drives the persist / plugin-callback pipeline shared by both + * invocation paths, through the after-run and compaction brackets. A {@code beforeRunEvent} + * short-circuits the agent run; {@code onEachPersisted} runs after each persisted event. + */ + private Flowable executeAgentPipeline( + InvocationContext context, + Session sessionToPersist, + Maybe beforeRunEvent, + @Nullable Runnable onEachPersisted) { // Let BaseLlmFlow block each step until this Runner has persisted the prior step's events. - PersistBarrier.enable(contextWithUpdatedSession); + PersistBarrier.enable(context); - // Agent execution Flowable agentEvents = - contextWithUpdatedSession + context .agent() - .runAsync(contextWithUpdatedSession) + .runAsync(context) .concatMap( agentEvent -> { // Mirror ADK Python (runners.py): partial events are streamed to the caller but @@ -651,39 +767,32 @@ private Flowable runAgentWithUpdatedSession( Single persistStep = agentEvent.partial().orElse(false) ? Single.just(agentEvent) - : this.sessionService.appendEvent(updatedSession, agentEvent); + : this.sessionService.appendEvent(sessionToPersist, agentEvent); return persistStep // Release (or fail) BaseLlmFlow's wait for this step; the Runner stays the // sole appendEvent caller (see PersistBarrier). .doOnSuccess( - unusedEvent -> - PersistBarrier.markPersisted( - contextWithUpdatedSession, agentEvent.id())) + unusedEvent -> PersistBarrier.markPersisted(context, agentEvent.id())) .doOnError( - error -> - PersistBarrier.markFailed( - contextWithUpdatedSession, agentEvent.id(), error)) + error -> PersistBarrier.markFailed(context, agentEvent.id(), error)) .flatMap( registeredEvent -> { - // TODO: remove this hack after deprecating runAsync with Session. - copySessionStates(updatedSession, initialContext.session()); - return contextWithUpdatedSession + if (onEachPersisted != null) { + onEachPersisted.run(); + } + return context .pluginManager() - .onEventCallback(contextWithUpdatedSession, registeredEvent) + .onEventCallback(context, registeredEvent) .defaultIfEmpty(registeredEvent); }) .toFlowable(); }); - // If beforeRunCallback returns content, emit it and skip agent - Context capturedContext = Context.current(); return beforeRunEvent .toFlowable() .switchIfEmpty(agentEvents) - .concatWith( - Completable.defer(() -> pluginManager.afterRunCallback(contextWithUpdatedSession))) - .concatWith(Completable.defer(() -> compactEvents(updatedSession))) - .compose(Tracing.withContext(capturedContext)); + .concatWith(Completable.defer(() -> pluginManager.afterRunCallback(context))) + .concatWith(Completable.defer(() -> compactEvents(sessionToPersist))); } private Completable compactEvents(Session session) { @@ -694,6 +803,367 @@ private Completable compactEvents(Session session) { .orElseGet(Completable::complete); } + /** Wraps before-run callback content as a model event that short-circuits the agent run. */ + private Maybe beforeRunEventFor(InvocationContext context) { + return this.pluginManager + .beforeRunCallback(context) + .map( + content -> + Event.builder() + .id(Event.generateEventId()) + .invocationId(context.invocationId()) + .author("model") + .content(content) + .build()); + } + + /** Fires the on-run-error plugin callback; the run paths share this error handler. */ + private void runOnRunError(InvocationContext context, Throwable throwable) { + this.pluginManager.runOnRunErrorCallback(context, throwable).onErrorComplete().subscribe(); + } + + /** + * Resumes an existing invocation when one resolves from {@code providedInvocationId} or a + * function response in {@code rawMessage}; otherwise starts a new invocation. Requires + * resumability. + */ + private Flowable runResumableFromSession( + Session session, + @Nullable String providedInvocationId, + @Nullable Content rawMessage, + RunConfig runConfig, + @Nullable Map stateDelta) { + // As in Python, a role-less message is a user turn; otherwise the request drops it. + Content newMessage = withUserRole(rawMessage); + return Flowable.defer( + () -> { + // Rejected before the invocation-id fallback; defer turns the throw into an error signal. + if (newMessage != null) { + validateResumeMessage(session, newMessage); + } + String resolvedInvocationId = + resolveInvocationId(session, newMessage, providedInvocationId); + if (resolvedInvocationId == null) { + if (newMessage == null) { + return Flowable.error( + new IllegalArgumentException( + "No new message provided and no resumable invocation to resume.")); + } + return runNewInvocation(session, newMessage, runConfig, stateDelta); + } + if (session.events().isEmpty()) { + // Only a wholly empty session is rejected; Python does not reject an unmatched id. + return Flowable.error( + new IllegalArgumentException("Session has no events to resume.")); + } + return resumeCore(session, resolvedInvocationId, newMessage, runConfig, stateDelta); + }); + } + + /** + * Runs an existing invocation on the given session: optionally appends {@code newMessage}, + * rehydrates agent checkpoints, skips a completed invocation, and runs the resolved agent under + * the resumed invocation id. + */ + private Flowable resumeCore( + Session session, + String resolvedInvocationId, + @Nullable Content newMessage, + RunConfig runConfig, + @Nullable Map stateDelta) { + return Flowable.defer( + () -> { + Context capturedContext = Context.current(); + + // Not new input: Python drops it and its state delta rather than adding a second prompt. + boolean supersededByOriginalMessage = + newMessage != null + && !carriesFunctionResponse(newMessage) + && findOriginalUserMessage(session, resolvedInvocationId).isPresent(); + + // Merged eagerly so the on-user-message callback sees it. + if (!supersededByOriginalMessage && stateDelta != null && !stateDelta.isEmpty()) { + stateDelta.forEach((key, value) -> session.state().put(key, value)); + } + + // For the pre-run plugin callbacks; the agent is resolved later, in runResumedAgent. + InvocationContext initialContext = + newInvocationContextBuilder(session) + .invocationId(resolvedInvocationId) + .runConfig(runConfig) + .userContent( + newMessage != null + ? newMessage + : originalUserMessage(session, resolvedInvocationId)) + .build(); + + Flowable events; + if (supersededByOriginalMessage) { + events = + runResumedAgent(session, resolvedInvocationId, /* userContent= */ null, runConfig); + } else if (newMessage != null) { + // Inherit the branch of the call it answers, so routing and rehydration see it. + String branch = + matchingFunctionCallEvent(session, newMessage).flatMap(Event::branch).orElse(null); + events = + this.pluginManager + .onUserMessageCallback(initialContext, newMessage) + .compose(Tracing.withContext(capturedContext)) + .defaultIfEmpty(newMessage) + .flatMap( + content -> + appendNewMessageToSession( + session, + content, + initialContext, + runConfig.saveInputBlobsAsArtifacts(), + stateDelta, + branch)) + // Persist first: an unpersisted response leaves the call looking unanswered. + .flatMap(userEvent -> this.sessionService.appendEvent(session, userEvent)) + .flatMapPublisher( + userEvent -> + runResumedAgent( + session, + resolvedInvocationId, + userEvent.content().orElse(null), + runConfig)); + } else if (stateDelta != null && !stateDelta.isEmpty()) { + // No message to carry the delta, so persist a content-less event, as Python does. + Event stateDeltaEvent = + Event.builder() + .id(Event.generateEventId()) + .invocationId(resolvedInvocationId) + .author(Role.USER) + .actions( + EventActions.builder() + .stateDelta(new ConcurrentHashMap<>(stateDelta)) + .build()) + .build(); + events = + this.sessionService + .appendEvent(session, stateDeltaEvent) + .ignoreElement() + .andThen( + runResumedAgent( + session, resolvedInvocationId, /* userContent= */ null, runConfig)); + } else { + events = + runResumedAgent(session, resolvedInvocationId, /* userContent= */ null, runConfig); + } + + return events + .doOnError(throwable -> runOnRunError(initialContext, throwable)) + .compose(Tracing.withContext(capturedContext)); + }); + } + + /** + * Runs the resolved agent for a resumed invocation with the same before-run / after-run plugin + * bracket as the new-invocation path. Rehydrates checkpoints, skips a completed invocation, and + * persists each event. + */ + private Flowable runResumedAgent( + Session session, + String resolvedInvocationId, + @Nullable Content userContent, + RunConfig runConfig) { + return Flowable.defer( + () -> { + BaseAgent rootAgent = this.agent; + InvocationContext rootContext = + newInvocationContextBuilder(session, rootAgent) + .invocationId(resolvedInvocationId) + .runConfig(runConfig) + .userContent( + userContent != null + ? userContent + : originalUserMessage(session, resolvedInvocationId)) + .build(); + // Rehydrate before routing, as Python does, so a root that checkpointed re-enters itself. + rootContext.populateInvocationAgentStates(); + BaseAgent resumeAgent = + rootContext.endOfAgents().containsKey(rootAgent.name()) + ? rootAgent + : findAgentToRun(session, rootAgent); + InvocationContext context = + resumeAgent.equals(rootAgent) + ? rootContext + : rootContext.toBuilder() + .agent(resumeAgent) + .branch(resumeParentBranch(session, resolvedInvocationId, resumeAgent)) + .build(); + + // No-op guard: a completed invocation (its active agent already finished) is not re-run. + if (context.endOfAgents().getOrDefault(context.agent().name(), false)) { + return Flowable.empty(); + } + + // before_run may short-circuit the run with a model event, as on the new-invocation path. + Maybe beforeRunEvent = beforeRunEventFor(context); + + return executeAgentPipeline( + context, session, beforeRunEvent, /* onEachPersisted= */ null); + }); + } + + /** + * Branch to seed a resumed context with so {@code resumeAgent} runs under the same branch it + * originally did. Returns the parent branch (the resolved agent's most recent event branch minus + * its own trailing name segment, which {@link BaseAgent#runAsync} re-appends), or {@code null} + * for the root branch. Non-null only for an agent nested under a {@link ParallelAgent}. + */ + private static @Nullable String resumeParentBranch( + Session session, String invocationId, BaseAgent resumeAgent) { + List events = session.events(); + for (int i = events.size() - 1; i >= 0; i--) { + Event event = events.get(i); + if (invocationId.equals(event.invocationId()) + && resumeAgent.name().equals(event.author()) + && event.branch().isPresent()) { + String branch = event.branch().get(); + String ownSegment = "." + resumeAgent.name(); + if (branch.endsWith(ownSegment)) { + String parent = branch.substring(0, branch.length() - ownSegment.length()); + return parent.isEmpty() ? null : parent; + } + return branch.equals(resumeAgent.name()) ? null : branch; + } + } + return null; + } + + /** + * Resolves which invocation a request targets: the invocation that issued the function call + * matching {@code newMessage}'s function response, else the caller-supplied {@code invocationId}. + * Returns {@code null} when neither applies (a fresh message starts a new invocation). + */ + private static @Nullable String resolveInvocationId( + Session session, @Nullable Content newMessage, @Nullable String invocationId) { + if (newMessage != null) { + return matchingFunctionCallEvent(session, newMessage) + .map(Event::invocationId) + .orElse(invocationId); + } + return invocationId; + } + + /** + * Returns the session event whose function call matches a function response id carried by {@code + * newMessage}, searching newest-first. Both the resumed invocation id and the branch of the + * appended function-response event are derived from it. + */ + private static Optional matchingFunctionCallEvent(Session session, Content newMessage) { + Set responseIds = functionResponseIds(newMessage); + if (responseIds.isEmpty()) { + return Optional.empty(); + } + List events = session.events(); + for (int i = events.size() - 1; i >= 0; i--) { + Event event = events.get(i); + for (FunctionCall call : event.functionCalls()) { + if (call.id().filter(responseIds::contains).isPresent()) { + return Optional.of(event); + } + } + } + return Optional.empty(); + } + + /** + * Rejects a resume message no invocation can be resumed from: one whose first function response + * carries no id, one whose response ids match no call in session history (it would feed the model + * an orphan response), and one whose responses span more than one invocation. Counts rather than + * ids are reported, since the ids come from the caller. Text alongside a function response is + * accepted, as on Python's agent path. + * + * @throws IllegalArgumentException if {@code newMessage} is not resumable. + */ + private static void validateResumeMessage(Session session, Content newMessage) { + ImmutableList responses = + newMessage.parts().stream() + .flatMap(List::stream) + .map(Part::functionResponse) + .flatMap(Optional::stream) + .collect(toImmutableList()); + if (responses.isEmpty()) { + return; + } + // Python checks the first response only, and raises rather than opening a new invocation that + // would carry an orphan response. + checkArgument( + responses.get(0).id().isPresent(), + "A function response id is required to resume an invocation."); + Set unmatched = functionResponseIds(newMessage); + Set invocationIds = new HashSet<>(); + List events = session.events(); + for (int i = events.size() - 1; i >= 0 && !unmatched.isEmpty(); i--) { + Event event = events.get(i); + for (FunctionCall call : event.functionCalls()) { + if (call.id().filter(unmatched::remove).isPresent()) { + invocationIds.add(event.invocationId()); + } + } + } + checkArgument( + unmatched.isEmpty(), + "No matching function call for %s of the resume message's function responses.", + unmatched.size()); + checkArgument( + invocationIds.size() <= 1, + "The resume message's function responses span %s invocations; all of them must answer the" + + " same one.", + invocationIds.size()); + } + + private static Set functionResponseIds(Content message) { + return message.parts().stream() + .flatMap(List::stream) + .map(Part::functionResponse) + .flatMap(Optional::stream) + .map(FunctionResponse::id) + .flatMap(Optional::stream) + .collect(toCollection(HashSet::new)); + } + + /** + * Returns the user message that started {@code invocationId}, so a resume that carries no message + * of its own still shows plugins, callbacks and the model the original prompt instead of empty + * content. A function-response message does not qualify: it answers an invocation rather than + * starting one. + * + * @throws IllegalArgumentException if the invocation has no such message, as in Python ADK. + */ + private static Content originalUserMessage(Session session, String invocationId) { + return findOriginalUserMessage(session, invocationId) + .orElseThrow( + () -> + new IllegalArgumentException( + "No user message to resume the requested invocation from.")); + } + + private static boolean carriesFunctionResponse(Content message) { + return message.parts().stream() + .flatMap(List::stream) + .anyMatch(part -> part.functionResponse().isPresent()); + } + + /** The user message that started {@code invocationId}, or empty when it has none. */ + private static Optional findOriginalUserMessage(Session session, String invocationId) { + return session.events().stream() + .filter(event -> invocationId.equals(event.invocationId())) + .filter(event -> Objects.equals(event.author(), Role.USER)) + .map(event -> event.content().orElse(null)) + .filter(Objects::nonNull) + .filter(content -> content.parts().stream().flatMap(List::stream).findAny().isPresent()) + .filter( + content -> + content.parts().stream() + .flatMap(List::stream) + .noneMatch(part -> part.functionResponse().isPresent())) + .findFirst(); + } + private void copySessionStates(Session source, Session target) { // TODO: remove this hack when deprecating all runAsync with Session. target.state().putAll(source.state()); @@ -725,7 +1195,7 @@ private InvocationContext newInvocationContextForLive( } } InvocationContext.Builder builder = - newInvocationContextBuilder(session) + newInvocationContextBuilder(session, findAgentToRun(session, this.agent)) .runConfig(runConfigBuilder.build()) .userContent(Content.fromParts()) .liveRequestQueue(liveRequestQueue); @@ -734,18 +1204,21 @@ private InvocationContext newInvocationContextForLive( } private InvocationContext.Builder newInvocationContextBuilder(Session session) { - BaseAgent rootAgent = this.agent; + return newInvocationContextBuilder(session, this.findAgentToRun(session, this.agent)); + } + + /** Overload for callers that already resolved the agent, so the walk is not repeated. */ + private InvocationContext.Builder newInvocationContextBuilder(Session session, BaseAgent agent) { return InvocationContext.builder() .sessionService(this.sessionService) .artifactService(this.artifactService) .memoryService(this.memoryService) .pluginManager(this.pluginManager) - .agent(rootAgent) .session(session) .eventsCompactionConfig(this.eventsCompactionConfig) .contextCacheConfig(this.contextCacheConfig) .resumabilityConfig(this.resumabilityConfig) - .agent(this.findAgentToRun(session, rootAgent)); + .agent(agent); } public Flowable runLive( @@ -832,10 +1305,7 @@ protected Flowable runLiveImpl( Span span = Span.current(); span.setStatus(StatusCode.ERROR, "Error in runLive Flowable execution"); span.recordException(throwable); - this.pluginManager - .runOnRunErrorCallback(invocationContext, throwable) - .onErrorComplete() - .subscribe(); + runOnRunError(invocationContext, throwable); }) .compose(Tracing.withContext(capturedContext)); }); @@ -863,30 +1333,34 @@ private boolean isTransferableAcrossAgentTree(BaseAgent agentToRun) { return true; } + /** Returns whether resumability is enabled for this runner's app. */ + private boolean isResumable() { + return resumabilityConfig != null && resumabilityConfig.isResumable(); + } + /** - * Returns whether resumability is enabled for this runner's app, by either the supported flag or - * the deprecated plain-text continuation shim, which selects the same behavior. + * Returns whether this runner's app runs the legacy resumption flow. A separate flow from {@link + * #isResumable()}, so a caller wanting either has to ask for both. */ @SuppressWarnings("deprecation") // The shim it reads is deprecated by design. - private boolean isResumable() { - return resumabilityConfig != null - && (resumabilityConfig.isResumable() - || resumabilityConfig.isPlainTextContinuationAutoResume()); + private boolean isLegacyResumability() { + return resumabilityConfig != null && resumabilityConfig.isPlainTextContinuationAutoResume(); } /** Returns the agent that should handle the next request based on session history. */ private BaseAgent findAgentToRun(Session session, BaseAgent rootAgent) { - // Route a function response to its call's author; when resumable, re-enter via the author's - // top-most SequentialAgent ancestor so the sequence can advance past it (else route straight to - // it, matching Python ADK v1 with resumability off). Temporary, event-based. + // Route a function response to its call's author, as Python's _agent_router does. An enclosing + // workflow advances by re-entering the root, which holds its own checkpoint, not from here. Optional functionCallAuthor = Functions.findMatchingFunctionCallEvent(session.events()) .filter(event -> event.author() != null) .flatMap(event -> rootAgent.findAgent(event.author())); if (functionCallAuthor.isPresent()) { - return isResumable() - ? topmostSequentialAncestor(functionCallAuthor.get()) - : functionCallAuthor.get(); + BaseAgent author = functionCallAuthor.get(); + if (isLegacyResumability()) { + return topmostSequentialAncestor(author); + } + return author; } List events = new ArrayList<>(session.events()); @@ -901,6 +1375,12 @@ private BaseAgent findAgentToRun(Session session, BaseAgent rootAgent) { continue; } + // Markers carry no model turn; resumable-only, since endOfAgent shares endInvocation's field. + if (isResumable() + && (event.actions().endOfAgent() || event.actions().agentState().isPresent())) { + continue; + } + if (author.equals(rootAgent.name())) { return rootAgent; } @@ -921,9 +1401,8 @@ private BaseAgent findAgentToRun(Session session, BaseAgent rootAgent) { /** * Returns the top-most ancestor reachable from {@code agent} through {@link SequentialAgent} - * parents, or {@code agent} itself otherwise. Only SequentialAgent is resume-aware; other - * workflow agents are left to resume their paused sub-agent directly (via the function-call - * author). + * parents, or {@code agent} itself otherwise. Used by the deprecated legacy flow, whose only + * resume-aware workflow agent is SequentialAgent. */ private static BaseAgent topmostSequentialAncestor(BaseAgent agent) { BaseAgent result = agent; diff --git a/core/src/main/java/com/google/adk/sessions/ApiClient.java b/core/src/main/java/com/google/adk/sessions/ApiClient.java index 0c630e460..30fc35537 100644 --- a/core/src/main/java/com/google/adk/sessions/ApiClient.java +++ b/core/src/main/java/com/google/adk/sessions/ApiClient.java @@ -154,7 +154,6 @@ OkHttpClient httpClient() { private Optional> getTimeoutHeader(HttpOptions httpOptionsToApply) { if (httpOptionsToApply.timeout().isPresent()) { int timeoutInSeconds = (int) Math.ceil((double) httpOptionsToApply.timeout().get() / 1000.0); - // TODO(b/329147724): Document the usage of X-Server-Timeout header. return Optional.of(ImmutableMap.of("X-Server-Timeout", Integer.toString(timeoutInSeconds))); } return Optional.empty(); diff --git a/core/src/main/java/com/google/adk/sessions/BaseSessionService.java b/core/src/main/java/com/google/adk/sessions/BaseSessionService.java index 8596f8eb6..655289014 100644 --- a/core/src/main/java/com/google/adk/sessions/BaseSessionService.java +++ b/core/src/main/java/com/google/adk/sessions/BaseSessionService.java @@ -35,6 +35,10 @@ * methods for creating, retrieving, listing, and deleting sessions, as well as listing and * appending events to a session. Implementations of this interface handle the underlying storage * and retrieval logic. + * + *

Implementations differ on whether they reject a duplicate session id; those that do signal it + * with a {@link SessionException} whose message is {@link SessionException#SESSION_ALREADY_EXISTS}. + * This applies to every {@code createSession} overload that accepts a session id. */ public interface BaseSessionService { @@ -68,6 +72,9 @@ Single createSession( * @param sessionId An optional client-provided identifier for the session. If empty or null, the * service should generate a unique ID. * @return The newly created {@link Session} instance. + * @throws SessionException with {@link SessionException#SESSION_ALREADY_EXISTS} if {@code + * sessionId} is already in use for this app and user; only implementations that reject + * duplicates throw it. * @throws SessionException if creation fails. */ default Single createSession( diff --git a/core/src/main/java/com/google/adk/sessions/InMemorySessionService.java b/core/src/main/java/com/google/adk/sessions/InMemorySessionService.java index e54289b76..e9c925cc1 100644 --- a/core/src/main/java/com/google/adk/sessions/InMemorySessionService.java +++ b/core/src/main/java/com/google/adk/sessions/InMemorySessionService.java @@ -36,6 +36,8 @@ import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.ConcurrentMap; import org.jspecify.annotations.Nullable; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; /** * An in-memory implementation of {@link BaseSessionService} assuming {@link Session} objects are @@ -49,6 +51,24 @@ * during retrieval operations ({@code getSession}, {@code createSession}). */ public final class InMemorySessionService implements BaseSessionService { + + private static final Logger logger = LoggerFactory.getLogger(InMemorySessionService.class); + + /** What {@code createSession} does when {@code sessionId} is already in use. */ + public enum DuplicateSessionIdBehavior { + /** + * Replaces the stored session, discarding its events and state. + * + * @deprecated Construct the service with {@link #REJECT} instead, which becomes the default in + * a future release. + */ + @Deprecated + OVERWRITE, + + /** Rejects the call with {@link SessionException}. Becomes the default in a future release. */ + REJECT; + } + // Structure: appName -> userId -> sessionId -> Session private final ConcurrentMap>> sessions; @@ -58,8 +78,26 @@ public final class InMemorySessionService implements BaseSessionService { // Structure: appName -> stateKey -> stateValue private final ConcurrentMap> appState; - /** Creates a new instance of the in-memory session service with empty storage. */ + private final DuplicateSessionIdBehavior duplicateSessionIdBehavior; + + /** + * Creates a new instance of the in-memory session service with empty storage. A duplicate session + * ID overwrites the stored session, which stays the default until callers have migrated. + */ public InMemorySessionService() { + this(DuplicateSessionIdBehavior.OVERWRITE); + } + + /** + * Creates a new instance of the in-memory session service with empty storage. + * + * @param duplicateSessionIdBehavior What {@code createSession} does when the session ID is + * already in use. + */ + public InMemorySessionService(DuplicateSessionIdBehavior duplicateSessionIdBehavior) { + this.duplicateSessionIdBehavior = + Objects.requireNonNull( + duplicateSessionIdBehavior, "duplicateSessionIdBehavior cannot be null"); this.sessions = new ConcurrentHashMap<>(); this.userState = new ConcurrentHashMap<>(); this.appState = new ConcurrentHashMap<>(); @@ -102,10 +140,22 @@ public Single createSession( .lastUpdateTime(Instant.now()) .build(); - sessions - .computeIfAbsent(appName, unused -> new ConcurrentHashMap<>()) - .computeIfAbsent(userId, unused -> new ConcurrentHashMap<>()) - .put(resolvedSessionId, newSession); + ConcurrentMap userSessions = + sessions + .computeIfAbsent(appName, unused -> new ConcurrentHashMap<>()) + .computeIfAbsent(userId, unused -> new ConcurrentHashMap<>()); + if (duplicateSessionIdBehavior == DuplicateSessionIdBehavior.REJECT) { + if (userSessions.putIfAbsent(resolvedSessionId, newSession) != null) { + return Single.error(new SessionException(SessionException.SESSION_ALREADY_EXISTS)); + } + } else { + Session previous = userSessions.put(resolvedSessionId, newSession); + if (previous != null) { + logger.info( + "Replaced an existing session because its ID was reused; its events and state are" + + " discarded."); + } + } // Create a mutable copy for the return value Session returnCopy = copySession(newSession); diff --git a/core/src/main/java/com/google/adk/sessions/SessionException.java b/core/src/main/java/com/google/adk/sessions/SessionException.java index 514811187..bf0b2bda5 100644 --- a/core/src/main/java/com/google/adk/sessions/SessionException.java +++ b/core/src/main/java/com/google/adk/sessions/SessionException.java @@ -19,6 +19,9 @@ /** Represents a general error that occurred during session management operations. */ public class SessionException extends RuntimeException { + /** Message used when {@code createSession} is given a session id that is already taken. */ + public static final String SESSION_ALREADY_EXISTS = "Session already exists"; + public SessionException(String message) { super(message); } diff --git a/core/src/main/java/com/google/adk/sessions/SessionJsonConverter.java b/core/src/main/java/com/google/adk/sessions/SessionJsonConverter.java index adc84fbb6..1e6c5327a 100644 --- a/core/src/main/java/com/google/adk/sessions/SessionJsonConverter.java +++ b/core/src/main/java/com/google/adk/sessions/SessionJsonConverter.java @@ -114,6 +114,7 @@ static String convertEventToJson(Event event, boolean useIsoString) { putIfNotEmpty(actionsJson, "requestedAuthConfigs", actions.requestedAuthConfigs()); putIfNotEmpty( actionsJson, "requestedToolConfirmations", actions.requestedToolConfirmations()); + actions.agentState().ifPresent(v -> actionsJson.put("agentState", v)); eventJson.put("actions", actionsJson); } event.content().ifPresent(c -> eventJson.put("content", SessionUtils.encodeContent(c))); @@ -151,6 +152,15 @@ private static Content convertMapToContent(Object rawContentValue) { } } + /** Returns {@code key} from {@code actions}, falling back to the copy under {@code rawEvent}. */ + private static @Nullable Object checkpointValue( + @Nullable Map actions, @Nullable Map rawActions, String key) { + if (actions != null && actions.get(key) != null) { + return actions.get(key); + } + return rawActions == null ? null : rawActions.get(key); + } + /** * Converts raw API event data into an {@link Event} object. * @@ -180,10 +190,6 @@ static Event fromApiEvent(Map apiEvent) { if (escalate != null) { eventActionsBuilder.escalate(escalate); } - Boolean endOfAgent = (Boolean) actionsMap.get("endOfAgent"); - if (endOfAgent != null) { - eventActionsBuilder.endOfAgent(endOfAgent); - } eventActionsBuilder.requestedAuthConfigs( Optional.ofNullable(actionsMap.get("requestedAuthConfigs")) .map(SessionJsonConverter::asConcurrentMapOfConcurrentMaps) @@ -194,6 +200,32 @@ static Event fromApiEvent(Map apiEvent) { .orElse(new ConcurrentHashMap<>())); } + // Python's Vertex session service leaves agentState and endOfAgent out of "actions" and keeps + // them only under "rawEvent", so a session it wrote would otherwise load with no checkpoints. + Map rawActions = null; + Object rawEvent = apiEvent.get("rawEvent"); + if (rawEvent instanceof Map) { + Object actions = ((Map) rawEvent).get("actions"); + if (actions instanceof Map) { + rawActions = (Map) actions; + } + } + // Resolved per key, not as a whole map: a rawEvent carrying only one of the two must not + // shadow the other where "actions" does carry it. + Object endOfAgent = checkpointValue(actionsMap, rawActions, "endOfAgent"); + if (endOfAgent instanceof Boolean value) { + eventActionsBuilder.endOfAgent(value); + } else if (endOfAgent != null) { + logger.warn("Ignoring 'endOfAgent' of unexpected type {}", endOfAgent.getClass().getName()); + } + Object agentState = checkpointValue(actionsMap, rawActions, "agentState"); + if (agentState instanceof Map) { + eventActionsBuilder.agentState((Map) agentState); + } else if (agentState != null) { + // Drop a foreign agentState rather than fail the whole session load on a bad cast. + logger.warn("Ignoring 'agentState' of unexpected type {}", agentState.getClass().getName()); + } + Event event = Event.builder() .id((String) Iterables.getLast(Splitter.on('/').split(apiEvent.get("name").toString()))) diff --git a/core/src/main/java/com/google/adk/sessions/VertexAiClient.java b/core/src/main/java/com/google/adk/sessions/VertexAiClient.java index 310c8f93f..6eef2e89f 100644 --- a/core/src/main/java/com/google/adk/sessions/VertexAiClient.java +++ b/core/src/main/java/com/google/adk/sessions/VertexAiClient.java @@ -147,8 +147,8 @@ Maybe listSessions(String reasoningEngineId, String userId) { } /** - * Wraps a value in an AIP-160 double-quoted string literal. Per go/aip/160, only backslashes and - * double quotes need escaping inside the quotes. + * Wraps a value in an AIP-160 double-quoted string literal. Per https://google.aip.dev/160, only + * backslashes and double quotes need escaping inside the quotes. */ private static String quoteFilterLiteral(String value) { return "\"" + value.replace("\\", "\\\\").replace("\"", "\\\"") + "\""; diff --git a/core/src/main/java/com/google/adk/tools/GoogleSearchAgentTool.java b/core/src/main/java/com/google/adk/tools/GoogleSearchAgentTool.java index 29c09028a..865cbeaa1 100644 --- a/core/src/main/java/com/google/adk/tools/GoogleSearchAgentTool.java +++ b/core/src/main/java/com/google/adk/tools/GoogleSearchAgentTool.java @@ -23,8 +23,7 @@ /** * A tool that wraps a sub-agent that only uses google_search tool. * - *

This is a workaround to support using google_search tool with other tools. TODO(b/448114567): - * Remove once the workaround is no longer needed. + *

This is a workaround to support using google_search tool with other tools. */ public class GoogleSearchAgentTool extends AgentTool { diff --git a/core/src/main/java/com/google/adk/tools/GoogleSearchTool.java b/core/src/main/java/com/google/adk/tools/GoogleSearchTool.java index d5ff9492c..3ce5803b4 100644 --- a/core/src/main/java/com/google/adk/tools/GoogleSearchTool.java +++ b/core/src/main/java/com/google/adk/tools/GoogleSearchTool.java @@ -25,8 +25,8 @@ import java.util.List; /** - * A built-in tool that is automatically invoked by Gemini 2 and 3 models to retrieve search results - * from Google Search. + * A built-in tool that is automatically invoked by Gemini models to retrieve search results from + * Google Search. * *

This tool operates internally within the model and does not require or perform local code * execution. @@ -60,27 +60,10 @@ public Completable processLlmRequest( List existingTools = configBuilder.build().tools().orElse(ImmutableList.of()); ImmutableList.Builder updatedToolsBuilder = ImmutableList.builder(); updatedToolsBuilder.addAll(existingTools); - - String model = llmRequestBuilder.build().model().orElse(null); - if (isSupportedModel(model)) { - - updatedToolsBuilder.add(Tool.builder().googleSearch(GoogleSearch.builder().build()).build()); - configBuilder.tools(updatedToolsBuilder.build()); - } else { - return Completable.error( - new IllegalArgumentException("Google search tool is not supported for model " + model)); - } + updatedToolsBuilder.add(Tool.builder().googleSearch(GoogleSearch.builder().build()).build()); + configBuilder.tools(updatedToolsBuilder.build()); llmRequestBuilder.config(configBuilder.build()); return Completable.complete(); } - - private boolean isSupportedModel(String model) { - if (model == null || !model.startsWith("gemini-")) { - return false; - } - return model.startsWith("gemini-2") - || model.startsWith("gemini-3") - || model.endsWith("-latest"); - } } diff --git a/core/src/main/java/com/google/adk/tools/ToolContext.java b/core/src/main/java/com/google/adk/tools/ToolContext.java index 974d1f017..9e0465145 100644 --- a/core/src/main/java/com/google/adk/tools/ToolContext.java +++ b/core/src/main/java/com/google/adk/tools/ToolContext.java @@ -68,13 +68,11 @@ public void toolConfirmation(ToolConfirmation toolConfirmation) { @SuppressWarnings("unused") private void requestCredential() { - // TODO: b/414678311 - Implement credential request logic. Make this public. throw new UnsupportedOperationException("Credential request not implemented yet."); } @SuppressWarnings("unused") private void getAuthResponse() { - // TODO: b/414678311 - Implement auth response retrieval logic. Make this public. throw new UnsupportedOperationException("Auth response retrieval not implemented yet."); } diff --git a/core/src/main/java/com/google/adk/tools/UrlContextTool.java b/core/src/main/java/com/google/adk/tools/UrlContextTool.java index fe7f9c77e..acf66dd7d 100644 --- a/core/src/main/java/com/google/adk/tools/UrlContextTool.java +++ b/core/src/main/java/com/google/adk/tools/UrlContextTool.java @@ -25,8 +25,8 @@ import java.util.List; /** - * A built-in tool that is automatically invoked by Gemini 2 and 3 models to retrieve information - * from the given URLs. + * A built-in tool that is automatically invoked by Gemini models to retrieve information from the + * given URLs. * *

This tool operates internally within the model and does not require or perform local code * execution. @@ -60,15 +60,8 @@ public Completable processLlmRequest( List existingTools = configBuilder.build().tools().orElse(ImmutableList.of()); ImmutableList.Builder updatedToolsBuilder = ImmutableList.builder(); updatedToolsBuilder.addAll(existingTools); - - String model = llmRequestBuilder.build().model().get(); - if (model != null && (model.startsWith("gemini-2") || model.startsWith("gemini-3"))) { - updatedToolsBuilder.add(Tool.builder().urlContext(UrlContext.builder().build()).build()); - configBuilder.tools(updatedToolsBuilder.build()); - } else { - return Completable.error( - new IllegalArgumentException("Url context tool is not supported for model " + model)); - } + updatedToolsBuilder.add(Tool.builder().urlContext(UrlContext.builder().build()).build()); + configBuilder.tools(updatedToolsBuilder.build()); llmRequestBuilder.config(configBuilder.build()); return Completable.complete(); diff --git a/core/src/main/java/com/google/adk/tools/mcp/DefaultMcpTransportBuilder.java b/core/src/main/java/com/google/adk/tools/mcp/DefaultMcpTransportBuilder.java index 086e0867d..e2fdfa7b8 100644 --- a/core/src/main/java/com/google/adk/tools/mcp/DefaultMcpTransportBuilder.java +++ b/core/src/main/java/com/google/adk/tools/mcp/DefaultMcpTransportBuilder.java @@ -64,7 +64,7 @@ public McpClientTransport build(Object connectionParams) { .orElse("")))) .build(); } else if (connectionParams instanceof StreamableHttpServerParameters streamableParams) { - // Split the URL so the transport's URI.resolve does not drop a custom path (b/513186321). + // Split the URL so the transport's URI.resolve does not drop a custom path. SplitUri split = splitBaseAndEndpoint(streamableParams.url()); HttpClientStreamableHttpTransport.Builder builder = HttpClientStreamableHttpTransport.builder(split.baseUri()) diff --git a/core/src/main/java/com/google/adk/tools/mcp/McpAsyncTool.java b/core/src/main/java/com/google/adk/tools/mcp/McpAsyncTool.java index 0f13bc9de..5d165cbbe 100644 --- a/core/src/main/java/com/google/adk/tools/mcp/McpAsyncTool.java +++ b/core/src/main/java/com/google/adk/tools/mcp/McpAsyncTool.java @@ -32,8 +32,6 @@ import org.slf4j.Logger; import org.slf4j.LoggerFactory; -// TODO(b/413489523): Add support for auth. This is a TODO for Python as well. - /** * Initializes a MCP tool. * diff --git a/core/src/main/java/com/google/adk/tools/mcp/McpSessionManager.java b/core/src/main/java/com/google/adk/tools/mcp/McpSessionManager.java index 65eb59fe9..5228e1931 100644 --- a/core/src/main/java/com/google/adk/tools/mcp/McpSessionManager.java +++ b/core/src/main/java/com/google/adk/tools/mcp/McpSessionManager.java @@ -36,7 +36,6 @@ *

This class provides methods for creating and initializing MCP client sessions, handling * different connection parameters and transport builders. */ -// TODO(b/413489523): Implement this class. public class McpSessionManager { private final Object connectionParams; // ServerParameters or SseServerParameters diff --git a/core/src/main/java/com/google/adk/tools/mcp/McpTool.java b/core/src/main/java/com/google/adk/tools/mcp/McpTool.java index 3c08643c8..8ec09a635 100644 --- a/core/src/main/java/com/google/adk/tools/mcp/McpTool.java +++ b/core/src/main/java/com/google/adk/tools/mcp/McpTool.java @@ -31,7 +31,6 @@ import org.slf4j.Logger; import org.slf4j.LoggerFactory; -// TODO(b/413489523): Add support for auth. This is a TODO for Python as well. /** * Initializes a MCP tool. * diff --git a/core/src/test/java/com/google/adk/agents/InvocationContextTest.java b/core/src/test/java/com/google/adk/agents/InvocationContextTest.java index 995387157..0bf6593cc 100644 --- a/core/src/test/java/com/google/adk/agents/InvocationContextTest.java +++ b/core/src/test/java/com/google/adk/agents/InvocationContextTest.java @@ -23,6 +23,7 @@ import com.google.adk.apps.ResumabilityConfig; import com.google.adk.artifacts.BaseArtifactService; import com.google.adk.events.Event; +import com.google.adk.events.EventActions; import com.google.adk.memory.BaseMemoryService; import com.google.adk.models.LlmCallsLimitExceededException; import com.google.adk.plugins.PluginManager; @@ -31,6 +32,7 @@ import com.google.adk.summarizer.EventsCompactionConfig; import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableMap; +import com.google.common.collect.ImmutableSet; import com.google.genai.types.Content; import com.google.genai.types.FunctionCall; import com.google.genai.types.FunctionResponse; @@ -74,30 +76,30 @@ public void setUp() { activeStreamingTools.put("test-tool", new ActiveStreamingTool(new LiveRequestQueue())); } - // The deprecated shim selects the same resumption behavior as resumable(true), so every - // resumability branch keyed on isResumable() must treat it identically. + // The two flags select different flows: the shim runs the legacy one and is not resumable. @Test - public void isResumable_shimOnly_reportsResumable() { - InvocationContext shimContext = contextWith(resumabilityConfigWithShim()); - InvocationContext resumableContext = contextWith(resumabilityConfigResumable()); + @SuppressWarnings("deprecation") // Exercises the deprecated shim. + public void isResumable_andIsLegacyResumability_separateTheTwoModes() { + InvocationContext shimContext = + contextWith(ResumabilityConfig.builder().plainTextContinuationAutoResume(true).build()); + InvocationContext resumableContext = + contextWith(ResumabilityConfig.builder().resumable(true).build()); + InvocationContext explicitlyOffContext = + contextWith(ResumabilityConfig.builder().resumable(false).build()); InvocationContext neitherContext = contextWith(null); - assertThat(shimContext.isResumable()).isTrue(); - assertThat(resumableContext.isResumable()).isTrue(); - assertThat(neitherContext.isResumable()).isFalse(); - } + assertThat(shimContext.isResumable()).isFalse(); + assertThat(shimContext.isLegacyResumability()).isTrue(); - @SuppressWarnings("deprecation") // Exercises the deprecated shim. - private static ResumabilityConfig resumabilityConfigWithShim() { - return ResumabilityConfig.builder().plainTextContinuationAutoResume(true).build(); - } + assertThat(resumableContext.isResumable()).isTrue(); + assertThat(resumableContext.isLegacyResumability()).isFalse(); - @SuppressWarnings("deprecation") // ResumabilityConfig is deprecated until durable resumability. - private static ResumabilityConfig resumabilityConfigResumable() { - return ResumabilityConfig.builder().resumable(true).build(); + // An explicit false and an absent config are the same answer, either side of the null guard. + assertThat(explicitlyOffContext.isResumable()).isFalse(); + assertThat(neitherContext.isResumable()).isFalse(); + assertThat(neitherContext.isLegacyResumability()).isFalse(); } - @SuppressWarnings("deprecation") // ResumabilityConfig is deprecated until durable resumability. private InvocationContext contextWith(ResumabilityConfig resumabilityConfig) { return InvocationContext.builder() .sessionService(mockSessionService) @@ -921,6 +923,19 @@ private static Event agentEvent(@Nullable String branch) { return Event.builder().author("some_agent").branch(branch).build(); } + private static Event agentEvent( + String invocationId, String author, EventActions actions, Content content) { + Event.Builder builder = + Event.builder().id(Event.generateEventId()).invocationId(invocationId).author(author); + if (actions != null) { + builder.actions(actions); + } + if (content != null) { + builder.content(content); + } + return builder.build(); + } + private static Event callEvent(@Nullable String branch, String callId) { return Event.builder() .author("some_agent") @@ -949,4 +964,465 @@ private static Event userResponseEvent(@Nullable String branch, String callId) { .build())) .build(); } + + // ---- Resumability: runtime checkpoint state. ---- + + private InvocationContext resumableContext(Session eventSession, String invocationId) { + return InvocationContext.builder() + .sessionService(mockSessionService) + .artifactService(mockArtifactService) + .agent(mockAgent) + .session(eventSession) + .invocationId(invocationId) + .resumabilityConfig(ResumabilityConfig.builder().resumable(true).build()) + .build(); + } + + @Test + public void setAgentState_storesStateAndClearsEnd() { + InvocationContext context = resumableContext(session, "inv"); + + context.setAgentState("a", ImmutableMap.of("k", "v"), /* endOfAgent= */ false); + + assertThat(context.agentStates()).containsEntry("a", ImmutableMap.of("k", "v")); + assertThat(context.endOfAgents()).containsEntry("a", false); + } + + @Test + public void setAgentState_endOfAgent_marksEndedAndDropsState() { + InvocationContext context = resumableContext(session, "inv"); + context.setAgentState("a", ImmutableMap.of("k", "v"), /* endOfAgent= */ false); + + context.setAgentState("a", /* agentState= */ null, /* endOfAgent= */ true); + + assertThat(context.endOfAgents()).containsEntry("a", true); + assertThat(context.agentStates()).doesNotContainKey("a"); + } + + @Test + public void setAgentState_nullStateNotEnded_clearsBoth() { + InvocationContext context = resumableContext(session, "inv"); + context.setAgentState("a", ImmutableMap.of("k", "v"), /* endOfAgent= */ false); + + context.setAgentState("a", /* agentState= */ null, /* endOfAgent= */ false); + + assertThat(context.agentStates()).doesNotContainKey("a"); + assertThat(context.endOfAgents()).doesNotContainKey("a"); + } + + @Test + public void resetSubAgentStates_recursivelyClearsDescendants() { + BaseAgent grandChild = SequentialAgent.builder().name("gc").build(); + BaseAgent child1 = + SequentialAgent.builder().name("c1").subAgents(ImmutableList.of(grandChild)).build(); + BaseAgent child2 = SequentialAgent.builder().name("c2").build(); + BaseAgent parent = + SequentialAgent.builder().name("p").subAgents(ImmutableList.of(child1, child2)).build(); + InvocationContext context = + InvocationContext.builder() + .sessionService(mockSessionService) + .artifactService(mockArtifactService) + .agent(parent) + .session(session) + .invocationId("inv") + .resumabilityConfig(ResumabilityConfig.builder().resumable(true).build()) + .build(); + context.setAgentState("c1", ImmutableMap.of("k", "v"), /* endOfAgent= */ false); + context.setAgentState("c2", ImmutableMap.of("k", "v"), /* endOfAgent= */ false); + context.setAgentState("gc", ImmutableMap.of("k", "v"), /* endOfAgent= */ false); + + context.resetSubAgentStates("p"); + + // Every descendant of p is cleared, including the grandchild reached recursively. + assertThat(context.agentStates()).doesNotContainKey("c1"); + assertThat(context.agentStates()).doesNotContainKey("c2"); + assertThat(context.agentStates()).doesNotContainKey("gc"); + } + + @Test + public void shouldPauseInvocation_resumableWithLongRunningCall_returnsTrue() { + InvocationContext context = resumableContext(session, "inv"); + Event event = + Event.builder() + .id("e1") + .invocationId("inv") + .author("a") + .content( + Content.fromParts( + Part.builder() + .functionCall(FunctionCall.builder().id("c1").name("tool").build()) + .build())) + .longRunningToolIds(ImmutableSet.of("c1")) + .build(); + + assertThat(context.shouldPauseInvocation(event)).isTrue(); + } + + @Test + public void shouldPauseInvocation_notResumable_returnsFalse() { + InvocationContext context = + InvocationContext.builder() + .sessionService(mockSessionService) + .artifactService(mockArtifactService) + .agent(mockAgent) + .session(session) + .invocationId("inv") + .build(); + Event event = + Event.builder() + .id("e1") + .invocationId("inv") + .author("a") + .content( + Content.fromParts( + Part.builder() + .functionCall(FunctionCall.builder().id("c1").name("tool").build()) + .build())) + .longRunningToolIds(ImmutableSet.of("c1")) + .build(); + + assertThat(context.shouldPauseInvocation(event)).isFalse(); + } + + @Test + public void shouldPauseInvocation_noLongRunningIds_returnsFalse() { + InvocationContext context = resumableContext(session, "inv"); + Event event = + Event.builder() + .id("e1") + .invocationId("inv") + .author("a") + .content( + Content.fromParts( + Part.builder() + .functionCall(FunctionCall.builder().id("c1").name("tool").build()) + .build())) + .build(); + + assertThat(context.shouldPauseInvocation(event)).isFalse(); + } + + @Test + public void shouldPauseInvocation_callIdNotInLongRunningSet_returnsFalse() { + InvocationContext context = resumableContext(session, "inv"); + Event event = + Event.builder() + .id("e1") + .invocationId("inv") + .author("a") + .content( + Content.fromParts( + Part.builder() + .functionCall(FunctionCall.builder().id("c1").name("tool").build()) + .build())) + .longRunningToolIds(ImmutableSet.of("other")) + .build(); + + assertThat(context.shouldPauseInvocation(event)).isFalse(); + } + + private static Event twoLongRunningCallsEvent() { + return Event.builder() + .id("m") + .invocationId("inv") + .author("root") + .content( + Content.fromParts( + Part.builder() + .functionCall(FunctionCall.builder().id("a").name("approve_a").build()) + .build(), + Part.builder() + .functionCall(FunctionCall.builder().id("b").name("approve_b").build()) + .build())) + .longRunningToolIds(ImmutableSet.of("a", "b")) + .build(); + } + + private static Event longRunningCallEvent(String callId, String name) { + return Event.builder() + .id("m-" + callId) + .invocationId("inv") + .author("root") + .content( + Content.fromParts( + Part.builder() + .functionCall(FunctionCall.builder().id(callId).name(name).build()) + .build())) + .longRunningToolIds(ImmutableSet.of(callId)) + .build(); + } + + private static Event functionResponseEvent(String id, String name) { + return Event.builder() + .id("r-" + id) + .invocationId("inv") + .author("user") + .content( + Content.fromParts( + Part.builder() + .functionResponse( + FunctionResponse.builder() + .id(id) + .name(name) + .response(ImmutableMap.of("status", "done")) + .build()) + .build())) + .build(); + } + + @Test + public void lastEventsPauseInvocation_callInsideWindow_returnsTrue() { + Session eventSession = Session.builder("s").build(); + eventSession.events().add(twoLongRunningCallsEvent()); + eventSession.events().add(functionResponseEvent("a", "approve_a")); + InvocationContext context = resumableContext(eventSession, "inv"); + + assertThat(context.lastEventsPauseInvocation()).isTrue(); + } + + @Test + public void lastEventsPauseInvocation_callPushedOutOfWindow_returnsFalse() { + Session eventSession = Session.builder("s").build(); + eventSession.events().add(twoLongRunningCallsEvent()); + eventSession.events().add(functionResponseEvent("a", "approve_a")); + eventSession.events().add(functionResponseEvent("b", "approve_b")); + InvocationContext context = resumableContext(eventSession, "inv"); + + assertThat(context.lastEventsPauseInvocation()).isFalse(); + } + + // Python looks only at whether the event carries a long-running call, never at whether it was + // answered, so a fully answered call still inside the window pauses. + @Test + public void lastEventsPauseInvocation_answeredCallInsideWindow_returnsTrue() { + Session eventSession = Session.builder("s").build(); + eventSession.events().add(longRunningCallEvent("c1", "approve_c1")); + eventSession.events().add(functionResponseEvent("c1", "approve_c1")); + InvocationContext context = resumableContext(eventSession, "inv"); + + assertThat(context.lastEventsPauseInvocation()).isTrue(); + } + + // The converse: an unanswered call older than the two-event window does not pause, as in Python. + @Test + public void lastEventsPauseInvocation_unansweredCallOlderThanWindow_returnsFalse() { + Session eventSession = Session.builder("s").build(); + eventSession.events().add(twoLongRunningCallsEvent()); + eventSession + .events() + .add(agentEvent("inv", "root", null, Content.fromParts(Part.fromText("thinking")))); + eventSession + .events() + .add(agentEvent("inv", "root", null, Content.fromParts(Part.fromText("still here")))); + InvocationContext context = resumableContext(eventSession, "inv"); + + assertThat(context.lastEventsPauseInvocation()).isFalse(); + } + + @Test + public void lastEventsPauseInvocation_noCall_returnsFalse() { + Session eventSession = Session.builder("s").build(); + eventSession + .events() + .add(agentEvent("inv", "user", null, Content.fromParts(Part.fromText("hi")))); + eventSession + .events() + .add(agentEvent("inv", "root", null, Content.fromParts(Part.fromText("answer")))); + InvocationContext context = resumableContext(eventSession, "inv"); + + assertThat(context.lastEventsPauseInvocation()).isFalse(); + } + + @Test + public void events_filtersByInvocationAndBranch() { + Session eventSession = Session.builder("s").build(); + Event thisInv = agentEvent("inv", "a", null, Content.fromParts(Part.fromText("x"))); + Event otherInv = agentEvent("other", "a", null, Content.fromParts(Part.fromText("y"))); + Event branchB = + Event.builder() + .id("e3") + .invocationId("inv") + .author("a") + .branch("branchB") + .content(Content.fromParts(Part.fromText("z"))) + .build(); + eventSession.events().add(thisInv); + eventSession.events().add(otherInv); + eventSession.events().add(branchB); + InvocationContext context = resumableContext(eventSession, "inv"); + + assertThat(context.events(/* currentInvocation= */ true, /* currentBranch= */ false)) + .containsExactly(thisInv, branchB) + .inOrder(); + // A null-branch event is visible on any branch; the "branchB" event is filtered out. + assertThat(context.events(/* currentInvocation= */ true, /* currentBranch= */ true)) + .containsExactly(thisInv); + } + + @Test + public void populateInvocationAgentStates_notResumable_doesNothing() { + Session eventSession = Session.builder("s").build(); + eventSession + .events() + .add( + agentEvent( + "inv", + "a", + EventActions.builder().agentState(ImmutableMap.of("k", "v")).build(), + Content.fromParts(Part.fromText("x")))); + InvocationContext context = + InvocationContext.builder() + .sessionService(mockSessionService) + .artifactService(mockArtifactService) + .agent(mockAgent) + .session(eventSession) + .invocationId("inv") + .build(); + + context.populateInvocationAgentStates(); + + assertThat(context.agentStates()).isEmpty(); + assertThat(context.endOfAgents()).isEmpty(); + } + + @Test + public void populateInvocationAgentStates_endOfAgentEvent_marksEndedAndRemovesState() { + Session eventSession = Session.builder("s").build(); + eventSession + .events() + .add( + agentEvent( + "inv", + "a", + EventActions.builder().agentState(ImmutableMap.of("k", "v")).build(), + Content.fromParts(Part.fromText("x")))); + eventSession + .events() + .add(agentEvent("inv", "a", EventActions.builder().endOfAgent(true).build(), null)); + InvocationContext context = resumableContext(eventSession, "inv"); + + context.populateInvocationAgentStates(); + + assertThat(context.endOfAgents()).containsEntry("a", true); + assertThat(context.agentStates()).doesNotContainKey("a"); + } + + @Test + public void populateInvocationAgentStates_agentStateEvent_setsStateAndClearsEnd() { + Session eventSession = Session.builder("s").build(); + eventSession + .events() + .add( + agentEvent( + "inv", + "a", + EventActions.builder().agentState(ImmutableMap.of("k", "v")).build(), + Content.fromParts(Part.fromText("x")))); + InvocationContext context = resumableContext(eventSession, "inv"); + + context.populateInvocationAgentStates(); + + assertThat(context.agentStates()).containsEntry("a", ImmutableMap.of("k", "v")); + assertThat(context.endOfAgents()).containsEntry("a", false); + } + + @Test + public void populateInvocationAgentStates_agentStateAndEndOfAgent_endOfAgentWins() { + Session eventSession = Session.builder("s").build(); + eventSession + .events() + .add( + agentEvent( + "inv", + "a", + EventActions.builder() + .endOfAgent(true) + .agentState(ImmutableMap.of("k", "v")) + .build(), + /* content= */ null)); + InvocationContext context = resumableContext(eventSession, "inv"); + + context.populateInvocationAgentStates(); + + assertThat(context.endOfAgents()).containsEntry("a", true); + assertThat(context.agentStates()).doesNotContainKey("a"); + } + + // EventActions.setEndInvocation aliases endOfAgent, so only a content-less event counts as done. + @Test + public void populateInvocationAgentStates_endOfAgentOnContentEvent_notTreatedAsCheckpoint() { + Session eventSession = Session.builder("s").build(); + eventSession + .events() + .add( + agentEvent( + "inv", + "a", + EventActions.builder().endOfAgent(true).build(), + Content.fromParts(Part.fromText("tool ended the invocation")))); + InvocationContext context = resumableContext(eventSession, "inv"); + + context.populateInvocationAgentStates(); + + assertThat(context.endOfAgents()).doesNotContainEntry("a", true); + } + + // An after-agent callback emits content after the end-of-agent marker; Python reopens the agent + // on it rather than leaving a resume with nothing to run. + @Test + public void populateInvocationAgentStates_contentAfterEndOfAgent_reopensAgent() { + Session eventSession = Session.builder("s").build(); + eventSession + .events() + .add(agentEvent("inv", "a", EventActions.builder().endOfAgent(true).build(), null)); + eventSession + .events() + .add(agentEvent("inv", "a", null, Content.fromParts(Part.fromText("callback output")))); + InvocationContext context = resumableContext(eventSession, "inv"); + + context.populateInvocationAgentStates(); + + assertThat(context.agentStates()).containsKey("a"); + assertThat(context.endOfAgents()).containsEntry("a", false); + } + + @Test + public void populateInvocationAgentStates_newContentFromNonUserAuthor_initializesEmptyState() { + Session eventSession = Session.builder("s").build(); + eventSession + .events() + .add(agentEvent("inv", "a", null, Content.fromParts(Part.fromText("hello")))); + InvocationContext context = resumableContext(eventSession, "inv"); + + context.populateInvocationAgentStates(); + + assertThat(context.agentStates()).containsKey("a"); + assertThat(context.agentStates().get("a")).isEmpty(); + assertThat(context.endOfAgents()).containsEntry("a", false); + } + + @Test + public void populateInvocationAgentStates_userMessage_ignoredForDefaultState() { + Session eventSession = Session.builder("s").build(); + eventSession + .events() + .add(agentEvent("inv", "user", null, Content.fromParts(Part.fromText("hi")))); + InvocationContext context = resumableContext(eventSession, "inv"); + + context.populateInvocationAgentStates(); + + assertThat(context.agentStates()).isEmpty(); + } + + @Test + public void populateInvocationAgentStates_noContentNoState_ignored() { + Session eventSession = Session.builder("s").build(); + eventSession.events().add(agentEvent("inv", "a", null, null)); + InvocationContext context = resumableContext(eventSession, "inv"); + + context.populateInvocationAgentStates(); + + assertThat(context.agentStates()).isEmpty(); + assertThat(context.endOfAgents()).isEmpty(); + } } diff --git a/core/src/test/java/com/google/adk/agents/LlmAgentTest.java b/core/src/test/java/com/google/adk/agents/LlmAgentTest.java index 26843bb56..d017860ce 100644 --- a/core/src/test/java/com/google/adk/agents/LlmAgentTest.java +++ b/core/src/test/java/com/google/adk/agents/LlmAgentTest.java @@ -19,13 +19,16 @@ import static com.google.adk.testing.TestUtils.assertEqualIgnoringFunctionIds; import static com.google.adk.testing.TestUtils.createInvocationContext; import static com.google.adk.testing.TestUtils.createLlmResponse; +import static com.google.adk.testing.TestUtils.createSubAgent; import static com.google.adk.testing.TestUtils.createTestAgent; import static com.google.adk.testing.TestUtils.createTestAgentBuilder; import static com.google.adk.testing.TestUtils.createTestLlm; import static com.google.adk.testing.TestUtils.createTextLlmResponse; +import static com.google.adk.testing.TestUtils.simplifyEvents; import static com.google.common.collect.Iterables.getOnlyElement; import static com.google.common.truth.Truth.assertThat; import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertThrows; import com.google.adk.agents.Callbacks.AfterModelCallback; import com.google.adk.agents.Callbacks.AfterToolCallback; @@ -33,7 +36,10 @@ import com.google.adk.agents.Callbacks.BeforeToolCallback; import com.google.adk.agents.Callbacks.OnModelErrorCallback; import com.google.adk.agents.Callbacks.OnToolErrorCallback; +import com.google.adk.apps.ResumabilityConfig; +import com.google.adk.artifacts.InMemoryArtifactService; import com.google.adk.events.Event; +import com.google.adk.events.EventActions; import com.google.adk.examples.Example; import com.google.adk.models.LlmRegistry; import com.google.adk.models.LlmRequest; @@ -49,8 +55,11 @@ import com.google.adk.tools.ExampleTool; import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableMap; +import com.google.common.collect.ImmutableSet; import com.google.errorprone.annotations.CanIgnoreReturnValue; import com.google.genai.types.Content; +import com.google.genai.types.FunctionCall; +import com.google.genai.types.FunctionResponse; import com.google.genai.types.Part; import com.google.genai.types.Schema; import com.google.genai.types.Type; @@ -632,4 +641,298 @@ public void run_withExampleTool_doesNotAddFunctionDeclarations() { var config = request.config().get(); assertThat(config.tools().isPresent()).isFalse(); } + + // ---- Resumability: resume into a transferred sub-agent (ResumableLlmAgentTest parity). ---- + + private static InvocationContext resumableContextWithSeededEvent( + LlmAgent rootAgent, Event seededEvent) { + InMemorySessionService sessionService = new InMemorySessionService(); + Session session = sessionService.createSession("app", "user").blockingGet(); + var unused = sessionService.appendEvent(session, seededEvent).blockingGet(); + return InvocationContext.builder() + .sessionService(sessionService) + .artifactService(new InMemoryArtifactService()) + .invocationId("inv") + .agent(rootAgent) + .session(session) + .userContent(Content.fromParts(Part.fromText("hi"))) + .runConfig(RunConfig.builder().build()) + .resumabilityConfig(ResumabilityConfig.builder().resumable(true).build()) + .build(); + } + + /** + * A transfer event shaped the way the flow really records one: the {@code transfer_to_agent} + * tool's function response, authored by the transferring agent and carrying the transfer action. + */ + private static Event transferEvent(String author, String targetAgent) { + return Event.builder() + .id("t1") + .invocationId("inv") + .author(author) + .actions(EventActions.builder().transferToAgent(targetAgent).build()) + .content( + Content.builder() + .role("user") + .parts( + Part.builder() + .functionResponse( + FunctionResponse.builder() + .id("transfer-call") + .name("transfer_to_agent") + .response(ImmutableMap.of()) + .build()) + .build()) + .build()) + .build(); + } + + @Test + public void runAsync_resumeFromTransferCall_runsTransferredSubAgent() { + LlmAgent sub = + createTestAgentBuilder(createTestLlm(createTextLlmResponse("sub response"))) + .name("sub") + .build(); + TestLlm rootLlm = createTestLlm(createTextLlmResponse("root should not run")); + LlmAgent root = createTestAgentBuilder(rootLlm).name("root").subAgents(sub).build(); + InvocationContext context = resumableContextWithSeededEvent(root, transferEvent("root", "sub")); + context.setAgentState("root", ImmutableMap.of(), /* endOfAgent= */ false); + + List events = root.runAsync(context).toList().blockingGet(); + + // The transferred sub-agent runs; the root model is not re-invoked; root marks end-of-agent. + assertThat(simplifyEvents(events)).contains("sub: sub response"); + assertThat(rootLlm.getRequests()).isEmpty(); + assertThat( + events.stream() + .anyMatch(event -> event.author().equals("root") && event.actions().endOfAgent())) + .isTrue(); + // Recorded in the invocation state too, so a later resume no-ops instead of re-running root. + assertThat(context.endOfAgents()).containsEntry("root", true); + } + + @Test + public void runAsync_resumeUserResponseWithNoMatchingCall_throws() { + // A user function response matching no prior call is rejected, not silently continued. + TestLlm rootLlm = createTestLlm(createTextLlmResponse("root should not run")); + LlmAgent root = createTestAgentBuilder(rootLlm).name("root").build(); + Event userResponse = + Event.builder() + .id("u1") + .invocationId("inv") + .author("user") + .content( + Content.fromParts( + Part.builder() + .functionResponse( + FunctionResponse.builder() + .id("nomatch") + .name("tool") + .response(ImmutableMap.of("k", "v")) + .build()) + .build())) + .build(); + InvocationContext context = resumableContextWithSeededEvent(root, userResponse); + context.setAgentState("root", ImmutableMap.of(), /* endOfAgent= */ false); + + var resumed = root.runAsync(context).toList(); + assertThrows(IllegalArgumentException.class, resumed::blockingGet); + assertThat(rootLlm.getRequests()).isEmpty(); + } + + @Test + public void runAsync_resumeFromTransfer_subAgentRepauses_rootDoesNotEndOfAgent() { + // The transferred sub-agent pauses again on its own long-running call when resumed. + Event subPause = + Event.builder() + .id("p1") + .invocationId("inv") + .author("sub") + .content( + Content.builder() + .parts( + Part.builder() + .functionCall( + FunctionCall.builder() + .id("lro") + .name("waitTool") + .args(ImmutableMap.of()))) + .role("model") + .build()) + .longRunningToolIds(ImmutableSet.of("lro")) + .build(); + BaseAgent sub = createSubAgent("sub", subPause); + LlmAgent root = + createTestAgentBuilder(createTestLlm(createTextLlmResponse("root should not run"))) + .name("root") + .subAgents(sub) + .build(); + InvocationContext context = resumableContextWithSeededEvent(root, transferEvent("root", "sub")); + context.setAgentState("root", ImmutableMap.of(), /* endOfAgent= */ false); + + List events = root.runAsync(context).toList().blockingGet(); + + // The sub-agent re-pauses, so root must not mark end-of-agent. + assertThat( + events.stream() + .anyMatch(event -> event.author().equals("root") && event.actions().endOfAgent())) + .isFalse(); + } + + @Test + public void runAsync_resumeTransferActionWithoutTransferResponse_continuesRootAgent() { + // Only a transfer_to_agent function response records a transfer, not the action alone. + LlmAgent sub = + createTestAgentBuilder(createTestLlm(createTextLlmResponse("sub should not run"))) + .name("sub") + .build(); + TestLlm rootLlm = createTestLlm(createTextLlmResponse("root continues")); + LlmAgent root = createTestAgentBuilder(rootLlm).name("root").subAgents(sub).build(); + Event actionOnPlainText = + Event.builder() + .id("t1") + .invocationId("inv") + .author("root") + .actions(EventActions.builder().transferToAgent("sub").build()) + .content(Content.fromParts(Part.fromText("just text"))) + .build(); + InvocationContext context = resumableContextWithSeededEvent(root, actionOnPlainText); + context.setAgentState("root", ImmutableMap.of(), /* endOfAgent= */ false); + + List events = root.runAsync(context).toList().blockingGet(); + + assertThat(simplifyEvents(events)).contains("root: root continues"); + assertThat(simplifyEvents(events)).doesNotContain("sub: sub should not run"); + } + + // A user event answering no call must not strand the transfer: the sub-agent still owns the turn. + @Test + public void runAsync_resumeAfterTransfer_userEventWithoutFunctionResponse_resumesSubAgent() { + LlmAgent sub = + createTestAgentBuilder(createTestLlm(createTextLlmResponse("sub resumes"))) + .name("sub") + .build(); + TestLlm rootLlm = createTestLlm(createTextLlmResponse("root should not run")); + LlmAgent root = createTestAgentBuilder(rootLlm).name("root").subAgents(sub).build(); + InvocationContext context = resumableContextWithSeededEvent(root, transferEvent("root", "sub")); + var unused = + context + .sessionService() + .appendEvent( + context.session(), + Event.builder() + .id("u1") + .invocationId("inv") + .author("user") + .content(Content.fromParts(Part.fromText("carry on"))) + .build()) + .blockingGet(); + context.setAgentState("root", ImmutableMap.of(), /* endOfAgent= */ false); + + List events = root.runAsync(context).toList().blockingGet(); + + assertThat(simplifyEvents(events)).contains("sub: sub resumes"); + assertThat(rootLlm.getRequests()).isEmpty(); + } + + // Mirrors AgentTransferTest.transferTargets_ordersSubAgentsThenParentThenPeers: the flow keeps a + // second copy of this rule that the agents package cannot reach, so both copies are pinned. + @Test + public void transferTargets_ordersSubAgentsThenParentThenPeers() { + LlmAgent child = + createTestAgentBuilder(createTestLlm(createTextLlmResponse("x"))).name("child").build(); + LlmAgent peerA = + createTestAgentBuilder(createTestLlm(createTextLlmResponse("x"))).name("peer_a").build(); + LlmAgent subject = + createTestAgentBuilder(createTestLlm(createTextLlmResponse("x"))) + .name("subject") + .subAgents(child) + .build(); + LlmAgent peerB = + createTestAgentBuilder(createTestLlm(createTextLlmResponse("x"))).name("peer_b").build(); + LlmAgent unused = + createTestAgentBuilder(createTestLlm(createTextLlmResponse("x"))) + .name("parent") + .subAgents(peerA, subject, peerB) + .build(); + + assertThat(subject.transferTargets().stream().map(BaseAgent::name)) + .containsExactly("child", "parent", "peer_a", "peer_b") + .inOrder(); + } + + @Test + public void transferTargets_noLlmAgentParent_returnsOnlySubAgents() { + LlmAgent child = + createTestAgentBuilder(createTestLlm(createTextLlmResponse("x"))).name("child").build(); + LlmAgent subject = + createTestAgentBuilder(createTestLlm(createTextLlmResponse("x"))) + .name("subject") + .subAgents(child) + .build(); + + assertThat(subject.transferTargets().stream().map(BaseAgent::name)).containsExactly("child"); + } + + @Test + public void runAsync_resumeFromTransferToMissingAgent_throws() { + // The transferred-to agent is gone: a resumed transfer must fail, as a live one does. + TestLlm rootLlm = createTestLlm(createTextLlmResponse("root should not run")); + LlmAgent root = createTestAgentBuilder(rootLlm).name("root").build(); + InvocationContext context = + resumableContextWithSeededEvent(root, transferEvent("root", "removed_agent")); + context.setAgentState("root", ImmutableMap.of(), /* endOfAgent= */ false); + + var resumed = root.runAsync(context).toList(); + assertThrows(IllegalStateException.class, resumed::blockingGet); + assertThat(rootLlm.getRequests()).isEmpty(); + } + + @Test + public void runAsync_resumeNoTransfer_continuesRootAgent() { + TestLlm rootLlm = createTestLlm(createTextLlmResponse("root continues")); + LlmAgent root = createTestAgentBuilder(rootLlm).name("root").build(); + Event priorModelResponse = + Event.builder() + .id("m1") + .invocationId("inv") + .author("root") + .content(Content.fromParts(Part.fromText("earlier response"))) + .build(); + InvocationContext context = resumableContextWithSeededEvent(root, priorModelResponse); + context.setAgentState("root", ImmutableMap.of(), /* endOfAgent= */ false); + + List events = root.runAsync(context).toList().blockingGet(); + + // No transfer recorded: the root agent continues by invoking its model. + assertThat(simplifyEvents(events)).contains("root: root continues"); + assertThat(rootLlm.getRequests()).hasSize(1); + } + + @Test + public void runAsync_resumeFromTransferToPeer_runsTransferredPeerAgent() { + LlmAgent root = + createTestAgentBuilder(createTestLlm(createTextLlmResponse("root"))) + .name("root") + .subAgents( + createTestAgentBuilder( + createTestLlm(createTextLlmResponse("agent A should not re-run"))) + .name("agent_a") + .build(), + createTestAgentBuilder(createTestLlm(createTextLlmResponse("agent B response"))) + .name("agent_b") + .build()) + .build(); + // agent_a transferred to its peer agent_b (not a descendant), then the invocation paused. + LlmAgent agentA = (LlmAgent) root.findAgent("agent_a").get(); + InvocationContext context = + resumableContextWithSeededEvent(agentA, transferEvent("agent_a", "agent_b")); + context.setAgentState("agent_a", ImmutableMap.of(), /* endOfAgent= */ false); + + List events = agentA.runAsync(context).toList().blockingGet(); + + // The transferred peer runs; agent_a does not re-run itself. + assertThat(simplifyEvents(events)).contains("agent_b: agent B response"); + assertThat(simplifyEvents(events)).doesNotContain("agent_a: agent A should not re-run"); + } } diff --git a/core/src/test/java/com/google/adk/agents/LoopAgentTest.java b/core/src/test/java/com/google/adk/agents/LoopAgentTest.java index b2d0778c6..6afd913e7 100644 --- a/core/src/test/java/com/google/adk/agents/LoopAgentTest.java +++ b/core/src/test/java/com/google/adk/agents/LoopAgentTest.java @@ -19,16 +19,21 @@ import static com.google.adk.testing.TestUtils.createEscalateEvent; import static com.google.adk.testing.TestUtils.createEvent; import static com.google.adk.testing.TestUtils.createInvocationContext; +import static com.google.adk.testing.TestUtils.createResumableInvocationContext; import static com.google.adk.testing.TestUtils.createSubAgent; import static com.google.common.collect.ImmutableList.toImmutableList; import static com.google.common.truth.Truth.assertThat; +import static org.junit.Assert.assertThrows; import com.google.adk.agents.Callbacks.BeforeAgentCallback; import com.google.adk.events.Event; import com.google.adk.testing.TestBaseAgent; import com.google.common.collect.ImmutableList; +import com.google.common.collect.ImmutableMap; +import com.google.common.collect.ImmutableSet; import com.google.common.collect.Iterables; import com.google.genai.types.Content; +import com.google.genai.types.FunctionCall; import com.google.genai.types.Part; import io.reactivex.rxjava3.core.Flowable; import io.reactivex.rxjava3.core.Maybe; @@ -240,4 +245,140 @@ public void runAsync_withEndInvocationInSubAgentCallback_stopsSubAgentButLoopCon assertThat(normalAgentRunCount.get()).isEqualTo(3); assertThat(subAgent2RunCount.get()).isEqualTo(1); } + + // ---- Resumability: durable checkpoint resume. ---- + + @Test + public void runAsync_notResumable_doesNotEmitEndOfAgent() { + TestBaseAgent subAgent = + createSubAgent("sub", createEvent("e").toBuilder().author("sub").build()); + LoopAgent loopAgent = + LoopAgent.builder().name("loop").subAgents(subAgent).maxIterations(1).build(); + InvocationContext context = createInvocationContext(loopAgent); + + List events = loopAgent.runAsync(context).toList().blockingGet(); + + assertThat(events.stream().anyMatch(event -> event.actions().endOfAgent())).isFalse(); + } + + @Test + public void runAsync_resumingFromMiddle_restoresIterationAndAgent() { + TestBaseAgent agent1 = + createSubAgent( + "agent1", () -> Flowable.just(createEvent("a1").toBuilder().author("agent1").build())); + TestBaseAgent agent2 = + createSubAgent( + "agent2", () -> Flowable.just(createEvent("a2").toBuilder().author("agent2").build())); + LoopAgent loopAgent = + LoopAgent.builder().name("loop").subAgents(agent1, agent2).maxIterations(3).build(); + InvocationContext context = createResumableInvocationContext(loopAgent); + // Resume mid-loop: iteration 1, at agent2. + context.setAgentState( + "loop", + ImmutableMap.of("current_sub_agent", "agent2", "times_looped", 1), + /* endOfAgent= */ false); + + List events = loopAgent.runAsync(context).toList().blockingGet(); + + // The first sub-agent to run is agent2 (agent1 is skipped in the resumed iteration). + String firstSubAgentAuthor = + events.stream() + .filter(event -> event.content().isPresent()) + .map(Event::author) + .filter(author -> author.equals("agent1") || author.equals("agent2")) + .findFirst() + .orElse(null); + assertThat(firstSubAgentAuthor).isEqualTo("agent2"); + // The loop still finishes (reaches maxIterations) and marks end-of-agent. + assertThat( + events.stream() + .anyMatch(event -> event.author().equals("loop") && event.actions().endOfAgent())) + .isTrue(); + } + + // Restarting the count would loop past the cap, so an unreadable value fails the resume instead. + @Test + public void runAsync_resumingWithNonNumericTimesLooped_throws() { + TestBaseAgent subAgent = + createSubAgent("sub", createEvent("e").toBuilder().author("sub").build()); + LoopAgent loopAgent = + LoopAgent.builder().name("loop").subAgents(subAgent).maxIterations(2).build(); + InvocationContext context = createResumableInvocationContext(loopAgent); + context.setAgentState( + "loop", + ImmutableMap.of("current_sub_agent", "sub", "times_looped", "two"), + /* endOfAgent= */ false); + + var resumed = loopAgent.runAsync(context).toList(); + + assertThrows(IllegalStateException.class, resumed::blockingGet); + } + + @Test + public void runAsync_resumingAtIterationCap_endsWithoutRunningSubAgent() { + TestBaseAgent subAgent = + createSubAgent("sub", createEvent("e").toBuilder().author("sub").build()); + LoopAgent loopAgent = + LoopAgent.builder().name("loop").subAgents(subAgent).maxIterations(2).build(); + InvocationContext context = createResumableInvocationContext(loopAgent); + // Resume with the loop already at the iteration cap (2 of 2 done). + context.setAgentState( + "loop", + ImmutableMap.of("current_sub_agent", "sub", "times_looped", 2), + /* endOfAgent= */ false); + + List events = loopAgent.runAsync(context).toList().blockingGet(); + + // Ends immediately: end-of-agent is emitted and no sub-agent runs. + assertThat( + events.stream() + .anyMatch(event -> event.author().equals("loop") && event.actions().endOfAgent())) + .isTrue(); + assertThat(events.stream().anyMatch(event -> event.author().equals("sub"))).isFalse(); + } + + @Test + public void runAsync_resumable_belowIterationCap_loopsExactlyMaxIterations() { + // Guards the cap comparison timesLooped >= maxIterations from below the cap. + TestBaseAgent subAgent = + createSubAgent( + "sub", () -> Flowable.just(createEvent("e").toBuilder().author("sub").build())); + LoopAgent loopAgent = + LoopAgent.builder().name("loop").subAgents(subAgent).maxIterations(2).build(); + InvocationContext context = createResumableInvocationContext(loopAgent); + + List events = loopAgent.runAsync(context).toList().blockingGet(); + + assertThat(events.stream().filter(event -> event.author().equals("sub")).count()).isEqualTo(2L); + assertThat( + events.stream() + .anyMatch(event -> event.author().equals("loop") && event.actions().endOfAgent())) + .isTrue(); + } + + @Test + public void runAsync_resumable_pausesOnLongRunningCall_doesNotEmitEndOfAgent() { + Event longRunningCall = + Event.builder() + .id("lro") + .invocationId("invocationId") + .author("sub") + .content( + Content.fromParts( + Part.builder() + .functionCall(FunctionCall.builder().id("c1").name("tool").build()) + .build())) + .longRunningToolIds(ImmutableSet.of("c1")) + .build(); + TestBaseAgent subAgent = createSubAgent("sub", longRunningCall); + LoopAgent loopAgent = + LoopAgent.builder().name("loop").subAgents(subAgent).maxIterations(5).build(); + InvocationContext context = createResumableInvocationContext(loopAgent); + + List events = loopAgent.runAsync(context).toList().blockingGet(); + + // Paused on the long-running call: no end-of-agent, and the loop did not run 5 iterations. + assertThat(events.stream().anyMatch(event -> event.actions().endOfAgent())).isFalse(); + assertThat(events.stream().filter(event -> event.author().equals("sub")).count()).isEqualTo(1L); + } } diff --git a/core/src/test/java/com/google/adk/agents/ParallelAgentTest.java b/core/src/test/java/com/google/adk/agents/ParallelAgentTest.java index e51240c45..97d107aa9 100644 --- a/core/src/test/java/com/google/adk/agents/ParallelAgentTest.java +++ b/core/src/test/java/com/google/adk/agents/ParallelAgentTest.java @@ -16,12 +16,22 @@ package com.google.adk.agents; +import static com.google.adk.testing.TestUtils.createEscalateEvent; +import static com.google.adk.testing.TestUtils.createFunctionCallLlmResponse; import static com.google.adk.testing.TestUtils.createInvocationContext; +import static com.google.adk.testing.TestUtils.createResumableInvocationContext; +import static com.google.adk.testing.TestUtils.createSubAgent; +import static com.google.adk.testing.TestUtils.createTestAgentBuilder; +import static com.google.adk.testing.TestUtils.createTestLlm; +import static com.google.adk.testing.TestUtils.createTextLlmResponse; +import static com.google.adk.testing.TestUtils.simplifyEvents; import static com.google.common.truth.Truth.assertThat; import static java.util.concurrent.TimeUnit.MILLISECONDS; import com.google.adk.events.Event; +import com.google.adk.tools.FunctionTool; import com.google.common.collect.ImmutableList; +import com.google.common.collect.ImmutableMap; import com.google.genai.types.Content; import com.google.genai.types.Part; import io.reactivex.rxjava3.core.Flowable; @@ -30,6 +40,7 @@ import io.reactivex.rxjava3.schedulers.TestScheduler; import io.reactivex.rxjava3.subscribers.TestSubscriber; import java.util.List; +import org.jspecify.annotations.Nullable; import org.junit.Test; import org.junit.runner.RunWith; import org.junit.runners.JUnit4; @@ -194,4 +205,195 @@ public void runAsync_withTestScheduler_usesVirtualTime() { testSubscriber.assertValueCount(1); testSubscriber.assertComplete(); } + + // ---- Resumability: checkpoint / skip-completed / pause. ---- + + /** Tools for the resumability tests. */ + public static final class Tools { + private Tools() {} + + // A long-running tool awaiting an external result returns nothing yet, so a branch pauses. + @SuppressWarnings("unused") // Invoked reflectively by FunctionTool. + public static @Nullable ImmutableMap waitForApproval(String reason) { + return null; + } + } + + @Test + public void runAsync_resumable_allSubAgentsEnd_emitsEndOfAgent() { + LlmAgent sub1 = + createTestAgentBuilder(createTestLlm(createTextLlmResponse("sub1 done"))) + .name("sub1") + .build(); + LlmAgent sub2 = + createTestAgentBuilder(createTestLlm(createTextLlmResponse("sub2 done"))) + .name("sub2") + .build(); + ParallelAgent parallelAgent = + ParallelAgent.builder().name("parallel").subAgents(sub1, sub2).build(); + InvocationContext context = createResumableInvocationContext(parallelAgent); + + List events = parallelAgent.runAsync(context).toList().blockingGet(); + + assertThat( + events.stream() + .anyMatch( + event -> event.author().equals("parallel") && event.actions().endOfAgent())) + .isTrue(); + } + + @Test + public void runAsync_resumable_sequentialSubAgentEnds_emitsEndOfAgent() { + LlmAgent leaf = + createTestAgentBuilder(createTestLlm(createTextLlmResponse("leaf done"))) + .name("leaf") + .build(); + SequentialAgent sequential = + SequentialAgent.builder().name("sequential").subAgents(leaf).build(); + ParallelAgent parallelAgent = + ParallelAgent.builder().name("parallel").subAgents(sequential).build(); + InvocationContext context = createResumableInvocationContext(parallelAgent); + + List events = parallelAgent.runAsync(context).toList().blockingGet(); + + // A SequentialAgent child records end-of-agent on completion, so the parallel agent ends too. + assertThat( + events.stream() + .anyMatch( + event -> event.author().equals("parallel") && event.actions().endOfAgent())) + .isTrue(); + } + + @Test + public void runAsync_resumable_subAgentEscalates_emitsEndOfAgent() { + // Authored by the sub-agent itself: only a direct sub-agent's escalation addresses this agent. + BaseAgent escalating = + createSubAgent( + "escalating", createEscalateEvent("esc").toBuilder().author("escalating").build()); + ParallelAgent parallelAgent = + ParallelAgent.builder().name("parallel").subAgents(escalating).build(); + InvocationContext context = createResumableInvocationContext(parallelAgent); + + List events = parallelAgent.runAsync(context).toList().blockingGet(); + + // A sub-agent escalation ends the parallel agent even though no branch recorded end-of-agent. + assertThat( + events.stream() + .anyMatch( + event -> event.author().equals("parallel") && event.actions().endOfAgent())) + .isTrue(); + } + + @Test + public void runAsync_notResumable_doesNotEmitEndOfAgent() { + LlmAgent sub1 = + createTestAgentBuilder(createTestLlm(createTextLlmResponse("sub1 done"))) + .name("sub1") + .build(); + ParallelAgent parallelAgent = ParallelAgent.builder().name("parallel").subAgents(sub1).build(); + InvocationContext context = createInvocationContext(parallelAgent); + + List events = parallelAgent.runAsync(context).toList().blockingGet(); + + assertThat(events.stream().anyMatch(event -> event.actions().endOfAgent())).isFalse(); + } + + @Test + public void runAsync_resumable_oneBranchPauses_doesNotEmitEndOfAgent() { + LlmAgent pausing = + createTestAgentBuilder( + createTestLlm( + createFunctionCallLlmResponse( + "c1", "waitForApproval", ImmutableMap.of("reason", "x")))) + .name("pausing") + .tools( + FunctionTool.create( + Tools.class, + "waitForApproval", + /* requireConfirmation= */ false, + /* isLongRunning= */ true)) + .build(); + LlmAgent completing = + createTestAgentBuilder(createTestLlm(createTextLlmResponse("done"))) + .name("completing") + .build(); + ParallelAgent parallelAgent = + ParallelAgent.builder().name("parallel").subAgents(pausing, completing).build(); + InvocationContext context = createResumableInvocationContext(parallelAgent); + + List events = parallelAgent.runAsync(context).toList().blockingGet(); + + // One branch paused: the parallel agent does not end. + assertThat( + events.stream() + .anyMatch( + event -> event.author().equals("parallel") && event.actions().endOfAgent())) + .isFalse(); + } + + @Test + public void runAsync_resumable_firstRun_emitsInitialStateCheckpoint() { + LlmAgent sub1 = + createTestAgentBuilder(createTestLlm(createTextLlmResponse("sub1 done"))) + .name("sub1") + .build(); + ParallelAgent parallelAgent = ParallelAgent.builder().name("parallel").subAgents(sub1).build(); + InvocationContext context = createResumableInvocationContext(parallelAgent); + // No prior state for "parallel": this is a first run, so an initial checkpoint is recorded. + + List events = parallelAgent.runAsync(context).toList().blockingGet(); + + assertThat( + events.stream() + .anyMatch( + event -> + event.author().equals("parallel") + && event.actions().agentState().isPresent())) + .isTrue(); + } + + @Test + public void runAsync_resumable_resume_doesNotReEmitInitialStateCheckpoint() { + LlmAgent sub1 = + createTestAgentBuilder(createTestLlm(createTextLlmResponse("sub1 done"))) + .name("sub1") + .build(); + ParallelAgent parallelAgent = ParallelAgent.builder().name("parallel").subAgents(sub1).build(); + InvocationContext context = createResumableInvocationContext(parallelAgent); + // The parallel agent already checkpointed that it started in a prior run. + context.setAgentState("parallel", ImmutableMap.of(), /* endOfAgent= */ false); + + List events = parallelAgent.runAsync(context).toList().blockingGet(); + + // No new initial checkpoint is emitted on resume (only the final end-of-agent checkpoint). + assertThat( + events.stream() + .anyMatch( + event -> + event.author().equals("parallel") + && event.actions().agentState().isPresent())) + .isFalse(); + } + + @Test + public void runAsync_resumable_skipsCompletedBranchesOnResume() { + LlmAgent sub1 = + createTestAgentBuilder(createTestLlm(createTextLlmResponse("sub1 done"))) + .name("sub1") + .build(); + LlmAgent sub2 = + createTestAgentBuilder(createTestLlm(createTextLlmResponse("sub2 done"))) + .name("sub2") + .build(); + ParallelAgent parallelAgent = + ParallelAgent.builder().name("parallel").subAgents(sub1, sub2).build(); + InvocationContext context = createResumableInvocationContext(parallelAgent); + // sub1 already finished in a prior run. + context.setAgentState("sub1", /* agentState= */ null, /* endOfAgent= */ true); + + List events = parallelAgent.runAsync(context).toList().blockingGet(); + + assertThat(simplifyEvents(events)).doesNotContain("sub1: sub1 done"); + assertThat(simplifyEvents(events)).contains("sub2: sub2 done"); + } } diff --git a/core/src/test/java/com/google/adk/agents/SequentialAgentTest.java b/core/src/test/java/com/google/adk/agents/SequentialAgentTest.java index ea3e524fa..2989acbcf 100644 --- a/core/src/test/java/com/google/adk/agents/SequentialAgentTest.java +++ b/core/src/test/java/com/google/adk/agents/SequentialAgentTest.java @@ -19,12 +19,19 @@ import static com.google.adk.testing.TestUtils.createEvent; import static com.google.adk.testing.TestUtils.createInvocationContext; import static com.google.adk.testing.TestUtils.createLlmResponse; +import static com.google.adk.testing.TestUtils.createResumableInvocationContext; import static com.google.adk.testing.TestUtils.createSubAgent; import static com.google.adk.testing.TestUtils.createTestAgent; +import static com.google.adk.testing.TestUtils.createTestAgentBuilder; import static com.google.adk.testing.TestUtils.createTestLlm; +import static com.google.adk.testing.TestUtils.createTextLlmResponse; +import static com.google.adk.testing.TestUtils.simplifyEvents; import static com.google.common.collect.Iterables.getOnlyElement; import static com.google.common.truth.Truth.assertThat; +import static org.junit.Assert.assertThrows; +import com.google.adk.apps.ResumabilityConfig; +import com.google.adk.artifacts.InMemoryArtifactService; import com.google.adk.events.Event; import com.google.adk.sessions.InMemorySessionService; import com.google.adk.sessions.Session; @@ -32,11 +39,14 @@ import com.google.adk.testing.TestLlm; import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableMap; +import com.google.common.collect.ImmutableSet; import com.google.genai.types.Content; import com.google.genai.types.FunctionCall; import com.google.genai.types.FunctionResponse; import com.google.genai.types.Part; +import java.util.HashMap; import java.util.List; +import java.util.Map; import org.junit.Test; import org.junit.runner.RunWith; import org.junit.runners.JUnit4; @@ -341,4 +351,160 @@ private static InvocationContext contextResumingBranchedCall( context.branch(contextBranch); return context; } + + // ---- Resumability: durable checkpoint resume. ---- + + @Test + public void runAsync_resumingFromMiddle_startsFromCorrectAgent() { + LlmAgent agent1 = + createTestAgentBuilder(createTestLlm(createTextLlmResponse("a1 done"))) + .name("agent1") + .build(); + LlmAgent agent2 = + createTestAgentBuilder(createTestLlm(createTextLlmResponse("a2 done"))) + .name("agent2") + .build(); + LlmAgent agent3 = + createTestAgentBuilder(createTestLlm(createTextLlmResponse("a3 done"))) + .name("agent3") + .build(); + SequentialAgent sequentialAgent = + SequentialAgent.builder().name("seq").subAgents(agent1, agent2, agent3).build(); + InvocationContext context = createResumableInvocationContext(sequentialAgent); + // Seed the checkpoint: resume at agent2. + context.setAgentState( + "seq", ImmutableMap.of("current_sub_agent", "agent2"), /* endOfAgent= */ false); + + List events = sequentialAgent.runAsync(context).toList().blockingGet(); + + assertThat(simplifyEvents(events)).doesNotContain("agent1: a1 done"); + assertThat(simplifyEvents(events)).contains("agent2: a2 done"); + assertThat(simplifyEvents(events)).contains("agent3: a3 done"); + } + + // A null name is not "finished": Python's model rejects it, so a resume must not silently skip + // every sub-agent. + @Test + public void runAsync_resumingFromNullSubAgentName_throws() { + LlmAgent agent1 = + createTestAgentBuilder(createTestLlm(createTextLlmResponse("a1 done"))) + .name("agent1") + .build(); + SequentialAgent sequentialAgent = + SequentialAgent.builder().name("seq").subAgents(agent1).build(); + InvocationContext context = createResumableInvocationContext(sequentialAgent); + Map checkpoint = new HashMap<>(); + checkpoint.put("current_sub_agent", null); + context.setAgentState("seq", checkpoint, /* endOfAgent= */ false); + + var resumed = sequentialAgent.runAsync(context).toList(); + + assertThrows(IllegalStateException.class, resumed::blockingGet); + } + + @Test + public void runAsync_resumable_emitsEndOfAgent() { + LlmAgent agent1 = + createTestAgentBuilder(createTestLlm(createTextLlmResponse("done"))).name("agent1").build(); + SequentialAgent sequentialAgent = + SequentialAgent.builder().name("seq").subAgents(agent1).build(); + InvocationContext context = createResumableInvocationContext(sequentialAgent); + + List events = sequentialAgent.runAsync(context).toList().blockingGet(); + + assertThat( + events.stream() + .anyMatch(event -> event.author().equals("seq") && event.actions().endOfAgent())) + .isTrue(); + } + + @Test + public void runAsync_notResumable_doesNotEmitEndOfAgent() { + LlmAgent agent1 = + createTestAgentBuilder(createTestLlm(createTextLlmResponse("done"))).name("agent1").build(); + SequentialAgent sequentialAgent = + SequentialAgent.builder().name("seq").subAgents(agent1).build(); + InvocationContext context = createInvocationContext(sequentialAgent); + + List events = sequentialAgent.runAsync(context).toList().blockingGet(); + + assertThat(events.stream().anyMatch(event -> event.actions().endOfAgent())).isFalse(); + } + + // Legacy resumption fast-forwards from history; the resumable flow uses its own checkpoints. + @Test + @SuppressWarnings("deprecation") // The legacy shim is deprecated by design. + public void runAsync_legacyResumption_doesNotRerunCompletedSubAgents() { + TestBaseAgent agent1 = + createSubAgent("agent1", createEvent("a1").toBuilder().author("agent1").build()); + TestBaseAgent agent2 = + createSubAgent("agent2", createEvent("a2").toBuilder().author("agent2").build()); + TestBaseAgent agent3 = + createSubAgent("agent3", createEvent("a3").toBuilder().author("agent3").build()); + SequentialAgent sequentialAgent = + SequentialAgent.builder().name("seq").subAgents(agent1, agent2, agent3).build(); + + // A legacy paused session: agent2's long-running call, now answered, with no checkpoints. + InMemorySessionService sessionService = new InMemorySessionService(); + Session session = sessionService.createSession("app", "user").blockingGet(); + var unusedCall = + sessionService + .appendEvent( + session, + Event.builder() + .id("fc") + .invocationId("inv") + .author("agent2") + .content( + Content.fromParts( + Part.builder() + .functionCall(FunctionCall.builder().id("c1").name("tool").build()) + .build())) + .longRunningToolIds(ImmutableSet.of("c1")) + .build()) + .blockingGet(); + var unusedResponse = + sessionService + .appendEvent( + session, + Event.builder() + .id("fr") + .invocationId("inv") + .author("user") + .content( + Content.fromParts( + Part.builder() + .functionResponse( + FunctionResponse.builder() + .id("c1") + .name("tool") + .response(ImmutableMap.of()) + .build()) + .build())) + .build()) + .blockingGet(); + InvocationContext context = + InvocationContext.builder() + .sessionService(sessionService) + .artifactService(new InMemoryArtifactService()) + .invocationId("inv") + .agent(sequentialAgent) + .session(session) + .userContent(Content.fromParts(Part.fromText("resume"))) + .runConfig(RunConfig.builder().build()) + .resumabilityConfig( + ResumabilityConfig.builder() + .resumable(false) + .plainTextContinuationAutoResume(true) + .build()) + .build(); + // No populate/checkpoint state seeded -- this is the legacy case. + + var unused = sequentialAgent.runAsync(context).toList().blockingGet(); + + // agent1 (already completed before the pause) must not re-run; agent2 and agent3 do. + assertThat(agent1.getInvocationCount()).isEqualTo(0); + assertThat(agent2.getInvocationCount()).isEqualTo(1); + assertThat(agent3.getInvocationCount()).isEqualTo(1); + } } diff --git a/core/src/test/java/com/google/adk/events/EventActionsTest.java b/core/src/test/java/com/google/adk/events/EventActionsTest.java index e547cac9e..6de5568ce 100644 --- a/core/src/test/java/com/google/adk/events/EventActionsTest.java +++ b/core/src/test/java/com/google/adk/events/EventActionsTest.java @@ -90,6 +90,7 @@ public void merge_mergesAllFields() { new ConcurrentHashMap<>(ImmutableMap.of("tool2", TOOL_CONFIRMATION))) .endOfAgent(true) .setModelResponse(ImmutableMap.of("field1", "value1")) + .rewindBeforeInvocationId("inv1") .build(); EventActions merged = eventActions1.toBuilder().merge(eventActions2).build(); @@ -111,6 +112,58 @@ public void merge_mergesAllFields() { assertThat(merged.endOfAgent()).isTrue(); assertThat(merged.compaction()).hasValue(COMPACTION); assertThat(merged.setModelResponse()).hasValue(ImmutableMap.of("field1", "value1")); + assertThat(merged.rewindBeforeInvocationId()).hasValue("inv1"); + } + + @Test + public void rewindBeforeInvocationId_roundTripsThroughToBuilder() { + EventActions actions = EventActions.builder().rewindBeforeInvocationId("inv1").build(); + + EventActions rebuilt = actions.toBuilder().build(); + + assertThat(rebuilt).isEqualTo(actions); + assertThat(rebuilt).isNotEqualTo(EventActions.builder().build()); + assertThat(rebuilt.rewindBeforeInvocationId()).hasValue("inv1"); + } + + @Test + public void rewindBeforeInvocationId_roundTripsThroughJson() { + EventActions actions = EventActions.builder().rewindBeforeInvocationId("inv1").build(); + + String json = actions.toJson(); + EventActions deserialized = EventActions.fromJsonString(json, EventActions.class); + + assertThat(json).contains("\"rewindBeforeInvocationId\":\"inv1\""); + assertThat(deserialized.rewindBeforeInvocationId()).hasValue("inv1"); + } + + @Test + public void rewindBeforeInvocationId_absentByDefault_andOmittedFromJson() { + EventActions actions = EventActions.builder().build(); + + assertThat(actions.rewindBeforeInvocationId()).isEmpty(); + assertThat(actions.toJson()).doesNotContain("rewindBeforeInvocationId"); + } + + @Test + public void merge_rewindBeforeInvocationId_absentInOther_keepsThis() { + EventActions first = EventActions.builder().rewindBeforeInvocationId("inv1").build(); + EventActions second = EventActions.builder().build(); + + EventActions merged = first.toBuilder().merge(second).build(); + + assertThat(merged.rewindBeforeInvocationId()).hasValue("inv1"); + } + + @Test + public void setRewindBeforeInvocationId_setsAndClears() { + EventActions actions = new EventActions(); + + actions.setRewindBeforeInvocationId("inv1"); + assertThat(actions.rewindBeforeInvocationId()).hasValue("inv1"); + + actions.setRewindBeforeInvocationId(null); + assertThat(actions.rewindBeforeInvocationId()).isEmpty(); } @Test diff --git a/core/src/test/java/com/google/adk/examples/ExampleUtilsTest.java b/core/src/test/java/com/google/adk/examples/ExampleUtilsTest.java index 2d22ed3f1..7ca9a2782 100644 --- a/core/src/test/java/com/google/adk/examples/ExampleUtilsTest.java +++ b/core/src/test/java/com/google/adk/examples/ExampleUtilsTest.java @@ -45,7 +45,6 @@ public List getExamples(String query) { } } - // TODO: sduskis - Should this 0 examples use case actually return ""? @Test public void buildFewShotFewShot_noExamples() { TestExampleProvider exampleProvider = new TestExampleProvider(ImmutableList.of()); diff --git a/core/src/test/java/com/google/adk/flows/llmflows/AgentTransferTest.java b/core/src/test/java/com/google/adk/flows/llmflows/AgentTransferTest.java index 79552520b..fa7abf0a9 100644 --- a/core/src/test/java/com/google/adk/flows/llmflows/AgentTransferTest.java +++ b/core/src/test/java/com/google/adk/flows/llmflows/AgentTransferTest.java @@ -20,9 +20,11 @@ import static com.google.adk.testing.TestUtils.createLlmResponse; import static com.google.adk.testing.TestUtils.createTestAgentBuilder; import static com.google.adk.testing.TestUtils.createTestLlm; +import static com.google.adk.testing.TestUtils.createTextLlmResponse; import static com.google.adk.testing.TestUtils.simplifyEvents; import static com.google.common.truth.Truth.assertThat; +import com.google.adk.agents.BaseAgent; import com.google.adk.agents.InvocationContext; import com.google.adk.agents.LiveRequest; import com.google.adk.agents.LiveRequestQueue; @@ -86,17 +88,76 @@ public Single> runAsync(Map args, ToolContex } } + // Order is load-bearing: this list is rendered into the model's transfer instruction. @Test - public void exitLoopTool_exitsLoop() { - Content generatedContent = - Content.fromParts( - Part.fromText("Mock LLM Response:I will call the exit_loop tool."), - Part.fromFunctionCall("exit_loop", ImmutableMap.of())); - - TestLlm unusedTestLlm = createTestLlm(createLlmResponse(generatedContent)); - // InvocationContext unusedInvocationContext = - // createInvocationContext(createTestAgent(testLlm)); - // TODO: b/413488103 - complete when LoopAgent is implemented. + public void transferTargets_ordersSubAgentsThenParentThenPeers() { + LlmAgent child = + createTestAgentBuilder(createTestLlm(createTextLlmResponse("unused"))) + .name("child") + .build(); + LlmAgent peerA = + createTestAgentBuilder(createTestLlm(createTextLlmResponse("unused"))) + .name("peer_a") + .build(); + LlmAgent subject = + createTestAgentBuilder(createTestLlm(createTextLlmResponse("unused"))) + .name("subject") + .subAgents(child) + .build(); + LlmAgent peerB = + createTestAgentBuilder(createTestLlm(createTextLlmResponse("unused"))) + .name("peer_b") + .build(); + LlmAgent unused = + createTestAgentBuilder(createTestLlm(createTextLlmResponse("unused"))) + .name("parent") + .subAgents(peerA, subject, peerB) + .build(); + + assertThat(AgentTransfer.transferTargets(subject).stream().map(BaseAgent::name)) + .containsExactly("child", "parent", "peer_a", "peer_b") + .inOrder(); + } + + @Test + public void transferTargets_disallowedParentAndPeers_returnsOnlySubAgents() { + LlmAgent child = + createTestAgentBuilder(createTestLlm(createTextLlmResponse("unused"))) + .name("child") + .build(); + LlmAgent peer = + createTestAgentBuilder(createTestLlm(createTextLlmResponse("unused"))).name("peer").build(); + LlmAgent subject = + createTestAgentBuilder(createTestLlm(createTextLlmResponse("unused"))) + .name("subject") + .subAgents(child) + .disallowTransferToParent(true) + .disallowTransferToPeers(true) + .build(); + LlmAgent unused = + createTestAgentBuilder(createTestLlm(createTextLlmResponse("unused"))) + .name("parent") + .subAgents(subject, peer) + .build(); + + assertThat(AgentTransfer.transferTargets(subject).stream().map(BaseAgent::name)) + .containsExactly("child"); + } + + @Test + public void transferTargets_noLlmAgentParent_returnsOnlySubAgents() { + LlmAgent child = + createTestAgentBuilder(createTestLlm(createTextLlmResponse("unused"))) + .name("child") + .build(); + LlmAgent subject = + createTestAgentBuilder(createTestLlm(createTextLlmResponse("unused"))) + .name("subject") + .subAgents(child) + .build(); + + assertThat(AgentTransfer.transferTargets(subject).stream().map(BaseAgent::name)) + .containsExactly("child"); } @Test diff --git a/core/src/test/java/com/google/adk/flows/llmflows/FunctionsTest.java b/core/src/test/java/com/google/adk/flows/llmflows/FunctionsTest.java index 8e8555114..c4aec95c0 100644 --- a/core/src/test/java/com/google/adk/flows/llmflows/FunctionsTest.java +++ b/core/src/test/java/com/google/adk/flows/llmflows/FunctionsTest.java @@ -399,6 +399,13 @@ public void getAskUserConfirmationFunctionCalls_eventWithConfirmationFunctionCal assertThat(result).containsExactly(confirmationCall1, confirmationCall2); } + @Test + public void hasPendingLongRunningCall_singleEventWithFunctionResponse_returnsFalse() { + // One event cannot hold both a paused call and its response, so nothing is pending. + assertThat(Functions.hasPendingLongRunningCall(ImmutableList.of(functionResponseEvent("c1")))) + .isFalse(); + } + // Default ToolExecutionMode.NONE behaves like PARALLEL: blocking tools still execute serially // on the caller thread (no worker scheduler is used), preserving the historical default. @Test @@ -568,10 +575,36 @@ public void hasPendingLongRunningCall_emptyList_returnsFalse() { assertThat(Functions.hasPendingLongRunningCall(ImmutableList.of())).isFalse(); } + @Test + public void hasPendingLongRunningCall_list_matchingResponse_stillReturnsTrue() { + // This legacy predicate does not match responses to calls; StepResume does that when resumable. + ImmutableList events = + ImmutableList.of(longRunningCallEvent("call1"), functionResponseEvent("call1")); + assertThat(Functions.hasPendingLongRunningCall(events)).isTrue(); + } + private static Event longRunningCallEvent(String callId) { return functionCallEvent(callId, callId); } + private static Event functionResponseEvent(String callId) { + return Event.builder() + .id("response_" + callId) + .invocationId("invocation1") + .author("user") + .content( + Content.fromParts( + Part.builder() + .functionResponse( + FunctionResponse.builder() + .id(callId) + .name("tool") + .response(ImmutableMap.of()) + .build()) + .build())) + .build(); + } + // Event with a function call; longRunningId, when non-null, is marked long-running. private static Event functionCallEvent(String callId, String longRunningId) { Event.Builder builder = diff --git a/core/src/test/java/com/google/adk/models/chat/ChatCompletionsResponseTest.java b/core/src/test/java/com/google/adk/models/chat/ChatCompletionsResponseTest.java index 9d7ba4154..e7ccedcf9 100644 --- a/core/src/test/java/com/google/adk/models/chat/ChatCompletionsResponseTest.java +++ b/core/src/test/java/com/google/adk/models/chat/ChatCompletionsResponseTest.java @@ -943,7 +943,7 @@ public void testToLlmResponse_nonStreamingText_messageLevelSignatureStaysOnDtoBu // future fix (propagating the signature to the text Part, mirroring streaming) flips // this test from passing to failing -- forcing an intentional, documented update. // - // TODO(b/...): consider attaching message.extraContent.google.thought_signature to the + // TODO: consider attaching message.extraContent.google.thought_signature to the // output text Part to match the streaming-path contract. If/when that fix lands, this // test should be updated to assert that textPart.thoughtSignature() has the decoded // bytes (compare with diff --git a/core/src/test/java/com/google/adk/plugins/agentanalytics/BigQueryAgentAnalyticsPluginTest.java b/core/src/test/java/com/google/adk/plugins/agentanalytics/BigQueryAgentAnalyticsPluginTest.java index 90ea29866..fd7f3355f 100644 --- a/core/src/test/java/com/google/adk/plugins/agentanalytics/BigQueryAgentAnalyticsPluginTest.java +++ b/core/src/test/java/com/google/adk/plugins/agentanalytics/BigQueryAgentAnalyticsPluginTest.java @@ -702,8 +702,7 @@ public void onEventCallback_emptyAuthorNoCurrentAgent_fallsBackToUnknownSentinel // An empty author is not a usable fallback: withFallbackAgent guards on // `author != null && !author.isEmpty()`, so fallbackAgentName stays unset and resolveAgentName // yields the "unknown" sentinel rather than an empty agent name. Pins the `&&` against a - // `||`-mutation (go/mutation-testing), which would stamp "" as the agent for empty-author - // events. + // `||`-mutation, which would stamp "" as the agent for empty-author events. when(mockInvocationContext.agent()).thenReturn(null); Event event = Event.builder() @@ -726,7 +725,7 @@ public void onEventCallback_nullAuthorNoCurrentAgent_fallsBackToUnknownSentinelW // A null author must be short-circuited by the `author != null` half of withFallbackAgent's // guard so `author.isEmpty()` is never dereferenced. Exercised via the AGENT_RESPONSE path // (whose extraAttributes tolerate a null author, unlike the STATE_DELTA map). Pins the `&&` - // against a `||`-mutation (go/mutation-testing), which would NPE on null authors. + // against a `||`-mutation, which would NPE on null authors. when(mockInvocationContext.agent()).thenReturn(null); Event event = Event.builder() diff --git a/core/src/test/java/com/google/adk/runner/RunnerLegacyResumabilityTest.java b/core/src/test/java/com/google/adk/runner/RunnerLegacyResumabilityTest.java index 7f0302758..b00583053 100644 --- a/core/src/test/java/com/google/adk/runner/RunnerLegacyResumabilityTest.java +++ b/core/src/test/java/com/google/adk/runner/RunnerLegacyResumabilityTest.java @@ -16,28 +16,36 @@ package com.google.adk.runner; +import static com.google.adk.testing.ResumabilityTestUtils.answerCall; +import static com.google.adk.testing.ResumabilityTestUtils.newSession; +import static com.google.adk.testing.ResumabilityTestUtils.pendingFunctionTool; +import static com.google.adk.testing.ResumabilityTestUtils.runTurn; +import static com.google.adk.testing.ResumabilityTestUtils.shimRunner; import static com.google.adk.testing.TestUtils.createFunctionCallLlmResponse; +import static com.google.adk.testing.TestUtils.createLlmResponse; import static com.google.adk.testing.TestUtils.createTestAgentBuilder; import static com.google.adk.testing.TestUtils.createTestLlm; import static com.google.adk.testing.TestUtils.createTextLlmResponse; import static com.google.adk.testing.TestUtils.simplifyEvents; +import static com.google.common.collect.ImmutableList.toImmutableList; import static com.google.common.truth.Truth.assertThat; import com.google.adk.agents.LlmAgent; import com.google.adk.agents.LoopAgent; import com.google.adk.agents.ParallelAgent; +import com.google.adk.agents.RunConfig; import com.google.adk.agents.SequentialAgent; -import com.google.adk.apps.App; -import com.google.adk.apps.ResumabilityConfig; import com.google.adk.events.Event; import com.google.adk.flows.llmflows.Functions; import com.google.adk.sessions.Session; import com.google.adk.telemetry.Tracing; +import com.google.adk.testing.ResumabilityTestUtils.Tools; import com.google.adk.testing.TestLlm; import com.google.adk.tools.FunctionTool; import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableMap; import com.google.common.collect.Iterables; +import com.google.common.collect.Streams; import com.google.genai.types.Content; import com.google.genai.types.FunctionCall; import com.google.genai.types.FunctionResponse; @@ -45,9 +53,9 @@ import io.opentelemetry.api.trace.Tracer; import io.opentelemetry.sdk.testing.junit4.OpenTelemetryRule; import io.reactivex.rxjava3.core.Flowable; -import java.util.List; +import java.util.Objects; +import java.util.Optional; import java.util.concurrent.atomic.AtomicInteger; -import org.jspecify.annotations.Nullable; import org.junit.After; import org.junit.Before; import org.junit.Rule; @@ -62,6 +70,7 @@ * shim, so the shim keeps behaving exactly as resumability did before it was split in two. */ @RunWith(JUnit4.class) +@SuppressWarnings("deprecation") // The class exists to exercise the deprecated shim flag. public final class RunnerLegacyResumabilityTest { @Rule public final OpenTelemetryRule openTelemetryRule = OpenTelemetryRule.create(); @@ -79,33 +88,396 @@ public void tearDown() { Tracing.setTracerForTesting(originalTracer); } - public static class Tools { - private Tools() {} + // Regression: after a transfer, a finished-check keyed on the root wedged later plain-text turns. + @Test + public void runAsync_resumableTransferWithPlainTextAutoResume_laterTurnsRunSubAgent() { + Content transferCall = + Content.fromParts( + Part.fromFunctionCall( + "transfer_to_agent", ImmutableMap.of("agent_name", "sub_agent_1"))); + TestLlm testLlm = + createTestLlm( + createLlmResponse(transferCall), + createTextLlmResponse("r1"), + createTextLlmResponse("r2"), + createTextLlmResponse("r3"), + createTextLlmResponse("r4"), + createTextLlmResponse("r5")); + LlmAgent subAgent1 = createTestAgentBuilder(testLlm).name("sub_agent_1").build(); + LlmAgent rootAgent = + createTestAgentBuilder(testLlm) + .name("root_agent") + .subAgents(ImmutableList.of(subAgent1)) + .build(); + Runner runner = shimRunner(rootAgent); + Session session = newSession(runner); - public static ImmutableMap echoTool(String message) { - return ImmutableMap.of("message", message); + // Turn 1 transfers; turns 2-5 must each be answered by the sub-agent, not an empty stream. + var unused = runTurn(runner, session, "m1"); + for (String expected : new String[] {"r2", "r3", "r4", "r5"}) { + ImmutableList turn = runTurn(runner, session, "m"); + assertThat(simplifyEvents(turn)).contains("sub_agent_1: " + expected); } + } - // A long-running tool awaiting an external result has nothing to return yet; FunctionTool - // coerces the absent return into an empty response. - @SuppressWarnings("unused") // Invoked reflectively by FunctionTool. - public static @Nullable ImmutableMap pendingTool(String message) { - return null; - } + // The shim must not re-attach to a checkpoint-less session: it starts a new invocation instead. + @Test + public void runAsync_plainTextAutoResume_checkpointlessSession_startsNewInvocation() { + LlmAgent agent = + createTestAgentBuilder(createTestLlm(createTextLlmResponse("second answer"))) + .name("agent") + .build(); + Runner runner = shimRunner(agent); + Session session = newSession(runner); + + // Seed a completed prior turn with no checkpoints, as a pre-checkpoint session looks. + String priorInvocationId = "pre_checkpoint_invocation"; + Event unusedUserEvent = + runner + .sessionService() + .appendEvent( + session, + Event.builder() + .id(Event.generateEventId()) + .invocationId(priorInvocationId) + .author("user") + .content(Content.fromParts(Part.fromText("first turn"))) + .build()) + .blockingGet(); + Event unusedModelEvent = + runner + .sessionService() + .appendEvent( + session, + Event.builder() + .id(Event.generateEventId()) + .invocationId(priorInvocationId) + .author("agent") + .content(Content.fromParts(Part.fromText("first answer"))) + .build()) + .blockingGet(); - static final AtomicInteger pendingProgressToolCalls = new AtomicInteger(0); + ImmutableList secondTurn = runTurn(runner, session, "second turn"); - // A long-running tool that reports progress: it returns a non-empty "pending" status on the - // initial call. Counts executions so a test can assert it runs exactly once across turns. - @SuppressWarnings("unused") // Invoked reflectively by FunctionTool. - public static ImmutableMap pendingProgressTool(String message) { - pendingProgressToolCalls.incrementAndGet(); - return ImmutableMap.of("status", "pending"); - } + // The agent runs under a fresh invocation id, not the checkpoint-less prior one. + assertThat(simplifyEvents(secondTurn)).contains("agent: second answer"); + assertThat( + secondTurn.stream() + .map(Event::invocationId) + .filter(priorInvocationId::equals) + .collect(toImmutableList())) + .isEmpty(); + } + + // The shim selects the legacy flow but does not resume on plain text; that starts a new one. + @Test + @SuppressWarnings( + "deprecation") // Resumability + the auto-resume shim are intentionally deprecated. + public void runAsync_plainTextContinuation_shimOn_startsNewInvocation() { + TestLlm testLlm = + createTestLlm( + createFunctionCallLlmResponse( + "lro_call_id", "pendingTool", ImmutableMap.of("message", "draft")), + createTextLlmResponse("continued")); + LlmAgent agent = + createTestAgentBuilder(testLlm).name("agent").tools(pendingFunctionTool()).build(); + Runner runner = shimRunner(agent); + Session session = newSession(runner); + + ImmutableList turn1 = runTurn(runner, session, "draft the note"); + String invocationId = turn1.get(0).invocationId(); + assertThat(testLlm.getRequests()).hasSize(1); // paused after a single model call + + ImmutableList turn2 = runTurn(runner, session, "Proceed"); + + // The continuation opens a new invocation and the agent plans again. + assertThat(testLlm.getRequests()).hasSize(2); + assertThat(turn2).isNotEmpty(); + Session reloaded = + runner + .sessionService() + .getSession("test", "user", session.id(), Optional.empty()) + .blockingGet(); + Event lastUserEvent = + Streams.findLast( + reloaded.events().stream().filter(event -> Objects.equals(event.author(), "user"))) + .orElse(null); + assertThat(lastUserEvent).isNotNull(); + assertThat(lastUserEvent.invocationId()).isNotEqualTo(invocationId); + assertThat(turn2.stream().noneMatch(event -> invocationId.equals(event.invocationId()))) + .isTrue(); + } + + // The shim resumes only unfinished invocations, so text after a finished turn starts a new one. + @Test + @SuppressWarnings( + "deprecation") // Resumability + the auto-resume shim are intentionally deprecated. + public void runAsync_plainText_autoResumeFlagOn_afterCompletedTurn_startsNewInvocation() { + TestLlm testLlm = + createTestLlm( + createTextLlmResponse("first answer"), createTextLlmResponse("second answer")); + LlmAgent agent = createTestAgentBuilder(testLlm).name("agent").build(); + Runner runner = shimRunner(agent); + Session session = newSession(runner); + + ImmutableList turn1 = runTurn(runner, session, "hi"); + String invocationId = turn1.get(0).invocationId(); + assertThat(simplifyEvents(turn1)).contains("agent: first answer"); + + ImmutableList turn2 = runTurn(runner, session, "again"); + + assertThat(simplifyEvents(turn2)).contains("agent: second answer"); + assertThat(turn2.get(0).invocationId()).isNotEqualTo(invocationId); + } + + // A pause nested in a SequentialAgent is the flat case: the shim does not resume on plain text. + @Test + @SuppressWarnings( + "deprecation") // Resumability + the auto-resume shim are intentionally deprecated. + public void runAsync_plainTextContinuation_inSequentialAgent_shimOn_startsNewInvocation() { + TestLlm testLlm = + createTestLlm( + createFunctionCallLlmResponse( + "lro_call_id", "pendingTool", ImmutableMap.of("message", "draft")), + createTextLlmResponse("should not re-plan")); + LlmAgent childAgent = + createTestAgentBuilder(testLlm).name("child_agent").tools(pendingFunctionTool()).build(); + SequentialAgent workflowAgent = + SequentialAgent.builder() + .name("workflow_agent") + .subAgents(ImmutableList.of(childAgent)) + .build(); + Runner runner = shimRunner(workflowAgent); + Session session = newSession(runner); + + ImmutableList turn1 = runTurn(runner, session, "draft the note"); + String invocationId = turn1.get(0).invocationId(); + assertThat(testLlm.getRequests()).hasSize(1); // paused inside the workflow after one model call + + Object unused = runTurn(runner, session, "Proceed"); + + // The plain-text turn opens a new invocation rather than resuming the paused one. + assertThat(testLlm.getRequests()).hasSize(2); + Session reloaded = + runner + .sessionService() + .getSession("test", "user", session.id(), Optional.empty()) + .blockingGet(); + Event lastUserEvent = + Streams.findLast( + reloaded.events().stream().filter(event -> Objects.equals(event.author(), "user"))) + .orElse(null); + assertThat(lastUserEvent).isNotNull(); + assertThat(lastUserEvent.invocationId()).isNotEqualTo(invocationId); + } + + // Shim-only: a value-returning long-running tool still pauses BaseLlmFlow after one model call. + @Test + @SuppressWarnings( + "deprecation") // Resumability + the auto-resume shim are intentionally deprecated. + public void runAsync_shimOnly_valueReturningLongRunningCall_pausesAfterSingleModelCall() { + TestLlm testLlm = + createTestLlm( + createFunctionCallLlmResponse( + "lro_call_id", "echoTool", ImmutableMap.of("message", "hi")), + createTextLlmResponse("should not be reached")); + LlmAgent agent = + createTestAgentBuilder(testLlm) + .name("agent") + .tools( + FunctionTool.create( + Tools.class, + "echoTool", + /* requireConfirmation= */ false, + /* isLongRunning= */ true)) + .build(); + Runner runner = shimRunner(agent); + Session session = newSession(runner); + + ImmutableList events = runTurn(runner, session, "go"); + + // The tool answered in-turn, but the flow still pauses rather than summarizing. + assertThat(testLlm.getRequests()).hasSize(1); + assertThat(simplifyEvents(events)).doesNotContain("agent: should not be reached"); + } + + // Shim-only: the legacy LoopAgent body stops on a pending call and writes no checkpoint. + @Test + @SuppressWarnings( + "deprecation") // Resumability + the auto-resume shim are intentionally deprecated. + public void runAsync_shimOnlyInLoopAgent_stopsOnPendingCallAndWritesNoCheckpoints() { + TestLlm llm = + createTestLlm( + createFunctionCallLlmResponse( + "lro_call_id", "pendingTool", ImmutableMap.of("message", "draft")), + createTextLlmResponse("second iteration")); + LlmAgent child = + createTestAgentBuilder(llm).name("looped_agent").tools(pendingFunctionTool()).build(); + LoopAgent loopAgent = + LoopAgent.builder().name("loop").maxIterations(3).subAgents(child).build(); + Runner runner = shimRunner(loopAgent); + Session session = newSession(runner); + + var unused = runTurn(runner, session, "start"); + + // The pending call stops the loop after the first iteration, and nothing is checkpointed. + assertThat(llm.getRequests()).hasSize(1); + Session reloaded = + runner + .sessionService() + .getSession("test", "user", session.id(), Optional.empty()) + .blockingGet(); + assertThat( + reloaded.events().stream() + .filter( + event -> + event.actions().endOfAgent() || event.actions().agentState().isPresent()) + .collect(toImmutableList())) + .isEmpty(); + } + + // Shim-only: a mid-sequence pause restarts the sequence and persists no checkpoint. + @Test + @SuppressWarnings( + "deprecation") // Resumability + the auto-resume shim are intentionally deprecated. + public void runAsync_shimOnlyInSequentialAgent_restartsSequenceAndWritesNoCheckpoints() { + TestLlm llmA = createTestLlm(createTextLlmResponse("A ran"), createTextLlmResponse("A replay")); + TestLlm llmB = + createTestLlm( + createFunctionCallLlmResponse( + "lro_call_id", "pendingTool", ImmutableMap.of("message", "draft")), + createTextLlmResponse("B done")); + TestLlm llmC = createTestLlm(createTextLlmResponse("C ran")); + LlmAgent agentA = createTestAgentBuilder(llmA).name("agent_a").build(); + LlmAgent agentB = + createTestAgentBuilder(llmB).name("agent_b").tools(pendingFunctionTool()).build(); + LlmAgent agentC = createTestAgentBuilder(llmC).name("agent_c").build(); + SequentialAgent workflowAgent = + SequentialAgent.builder() + .name("workflow_agent") + .subAgents(ImmutableList.of(agentA, agentB, agentC)) + .build(); + Runner runner = shimRunner(workflowAgent); + Session session = newSession(runner); + + ImmutableList turn1 = runTurn(runner, session, "draft the note"); + + // Turn 1 pauses at agent_b: agent_c never runs. + assertThat(llmA.getRequests()).hasSize(1); + assertThat(llmB.getRequests()).hasSize(1); + assertThat(llmC.getRequests()).isEmpty(); + assertThat(simplifyEvents(turn1)).doesNotContain("agent_c: C ran"); + + ImmutableList turn2 = runTurn(runner, session, "Proceed"); + + // No durable state to resume from, so the sequence restarts at agent_a and runs through. + assertThat(llmA.getRequests()).hasSize(2); + assertThat(llmB.getRequests()).hasSize(2); + assertThat(llmC.getRequests()).hasSize(1); + assertThat(simplifyEvents(turn2)).contains("agent_c: C ran"); + + // The shim never writes durable state: nothing was checkpointed into the session. + Session reloaded = + runner + .sessionService() + .getSession("test", "user", session.id(), Optional.empty()) + .blockingGet(); + assertThat( + reloaded.events().stream() + .filter( + event -> + event.actions().endOfAgent() || event.actions().agentState().isPresent()) + .collect(toImmutableList())) + .isEmpty(); + } + + // Shim-only: a grandchild-authored pause also restarts the sequence rather than resuming. + @Test + @SuppressWarnings( + "deprecation") // Resumability + the auto-resume shim are intentionally deprecated. + public void runAsync_shimOnlyInNestedSequentialAgent_restartsSequence() { + TestLlm llmA = createTestLlm(createTextLlmResponse("A ran"), createTextLlmResponse("A replay")); + TestLlm llmB = + createTestLlm( + createFunctionCallLlmResponse( + "lro_call_id", "pendingTool", ImmutableMap.of("message", "draft")), + createTextLlmResponse("B done")); + LlmAgent agentA = createTestAgentBuilder(llmA).name("agent_a").build(); + LlmAgent agentB = + createTestAgentBuilder(llmB).name("agent_b").tools(pendingFunctionTool()).build(); + SequentialAgent innerWorkflow = + SequentialAgent.builder() + .name("inner_workflow") + .subAgents(ImmutableList.of(agentB)) + .build(); + SequentialAgent workflowAgent = + SequentialAgent.builder() + .name("workflow_agent") + .subAgents(ImmutableList.of(agentA, innerWorkflow)) + .build(); + Runner runner = shimRunner(workflowAgent); + Session session = newSession(runner); + + var unused = runTurn(runner, session, "draft the note"); + ImmutableList turn2 = runTurn(runner, session, "Proceed"); + + assertThat(llmA.getRequests()).hasSize(2); // the sequence restarts at agent_a + assertThat(llmB.getRequests()).hasSize(2); + assertThat(simplifyEvents(turn2)).contains("agent_b: B done"); } + // With the shim, a plain-text continuation starts a new invocation and still merges stateDelta. + @Test + public void runAsync_plainTextWithShim_withStateDelta_mergesStateIntoSession() { + TestLlm testLlm = + createTestLlm( + createFunctionCallLlmResponse( + "lro_call_id", "pendingTool", ImmutableMap.of("message", "hello")), + createTextLlmResponse("should not be reached")); + LlmAgent agent = + createTestAgentBuilder(testLlm).name("agent").tools(pendingFunctionTool()).build(); + Runner runner = shimRunner(agent); + Session session = newSession(runner); + + ImmutableList pausedTurn = runTurn(runner, session, "start"); + String pausedInvocationId = pausedTurn.get(0).invocationId(); + assertThat(testLlm.getRequests()).hasSize(1); // paused after a single model call + + // Plain-text "Proceed" via the non-resume overload; the flag resumes the paused invocation. + ImmutableMap stateDelta = ImmutableMap.of("key1", "value1", "key2", 42); + ImmutableList resumed = + ImmutableList.copyOf( + runner + .runAsync( + "user", + session.id(), + Content.fromParts(Part.fromText("Proceed")), + RunConfig.builder().build(), + stateDelta) + .toList() + .blockingGet()); + + // The shim starts a new invocation for the plain-text turn; the stateDelta still lands. + assertThat(testLlm.getRequests()).hasSize(2); + assertThat(resumed).isNotEmpty(); + Session finalSession = + runner + .sessionService() + .getSession("test", "user", session.id(), Optional.empty()) + .blockingGet(); + Event continuation = + Streams.findLast( + finalSession.events().stream() + .filter(event -> Objects.equals(event.author(), "user"))) + .orElseThrow(); + assertThat(continuation.invocationId()).isNotEqualTo(pausedInvocationId); + assertThat(continuation.actions().stateDelta()).containsAtLeastEntriesIn(stateDelta); + assertThat(finalSession.state()).containsAtLeastEntriesIn(stateDelta); + } + + // ===== CL1-parity: every CL1 resumable(true) test, re-run under the text-only shim ===== + @Test - @SuppressWarnings("deprecation") // Resumability flag is intentionally deprecated (partial). public void runAsync_withToolConfirmation_inSequentialAgent_runsLaterSubAgentsAfterResume_legacyShim() { LlmAgent agentA = @@ -134,23 +506,10 @@ public static ImmutableMap pendingProgressTool(String message) { .name("workflow_agent") .subAgents(ImmutableList.of(agentA, agentB, agentC)) .build(); - Runner runner = - Runner.builder() - .app( - App.builder() - .name("test") - .rootAgent(workflowAgent) - .resumabilityConfig( - ResumabilityConfig.builder().plainTextContinuationAutoResume(true).build()) - .build()) - .build(); - Session session = runner.sessionService().createSession("test", "user").blockingGet(); + Runner runner = shimRunner(workflowAgent); + Session session = newSession(runner); - List eventsBeforeConfirmation = - runner - .runAsync("user", session.id(), Content.fromParts(Part.fromText("from user"))) - .toList() - .blockingGet(); + ImmutableList eventsBeforeConfirmation = runTurn(runner, session, "from user"); // Turn 1: A runs, B pauses for confirmation, and C must not run yet. assertThat(simplifyEvents(eventsBeforeConfirmation)).contains("a_agent: agent A done"); @@ -163,21 +522,22 @@ public static ImmutableMap pendingProgressTool(String message) { .filter(functionCalls -> !functionCalls.isEmpty()) .findFirst() .get()); - List eventsAfterConfirmation = - runner - .runAsync( - "user", - session.id(), - Content.fromParts( - Part.builder() - .functionResponse( - FunctionResponse.builder() - .id(askUserConfirmationFunctionCall.id().get()) - .name(askUserConfirmationFunctionCall.name().get()) - .response(ImmutableMap.of("confirmed", true))) - .build())) - .toList() - .blockingGet(); + ImmutableList eventsAfterConfirmation = + ImmutableList.copyOf( + runner + .runAsync( + "user", + session.id(), + Content.fromParts( + Part.builder() + .functionResponse( + FunctionResponse.builder() + .id(askUserConfirmationFunctionCall.id().get()) + .name(askUserConfirmationFunctionCall.name().get()) + .response(ImmutableMap.of("confirmed", true))) + .build())) + .toList() + .blockingGet()); // Turn 2: B resumes and executes the tool, then C runs. A is not re-run. assertThat(simplifyEvents(eventsAfterConfirmation)) @@ -189,7 +549,6 @@ public static ImmutableMap pendingProgressTool(String message) { } @Test - @SuppressWarnings("deprecation") // Resumability flag is intentionally deprecated (partial). public void runAsync_withLongRunningCall_inSequentialAgent_runsLaterSubAgentsAfterResume_legacyShim() { LlmAgent agentA = @@ -222,23 +581,10 @@ public static ImmutableMap pendingProgressTool(String message) { .name("workflow_agent") .subAgents(ImmutableList.of(agentA, agentB, agentC)) .build(); - Runner runner = - Runner.builder() - .app( - App.builder() - .name("test") - .rootAgent(workflowAgent) - .resumabilityConfig( - ResumabilityConfig.builder().plainTextContinuationAutoResume(true).build()) - .build()) - .build(); - Session session = runner.sessionService().createSession("test", "user").blockingGet(); + Runner runner = shimRunner(workflowAgent); + Session session = newSession(runner); - List eventsBeforeResume = - runner - .runAsync("user", session.id(), Content.fromParts(Part.fromText("from user"))) - .toList() - .blockingGet(); + ImmutableList eventsBeforeResume = runTurn(runner, session, "from user"); // Turn 1: A runs, B issues the long-running call and pauses; C must not run yet. B must not // make @@ -247,21 +593,8 @@ public static ImmutableMap pendingProgressTool(String message) { assertThat(simplifyEvents(eventsBeforeResume)).doesNotContain("b_agent: agent B resumed"); assertThat(simplifyEvents(eventsBeforeResume)).doesNotContain("c_agent: agent C done"); - List eventsAfterResume = - runner - .runAsync( - "user", - session.id(), - Content.fromParts( - Part.builder() - .functionResponse( - FunctionResponse.builder() - .id("lro_call_id") - .name("echoTool") - .response(ImmutableMap.of("message", "hello"))) - .build())) - .toList() - .blockingGet(); + ImmutableList eventsAfterResume = + answerCall(runner, session, "lro_call_id", "echoTool", ImmutableMap.of("message", "hello")); // Turn 2: B resumes from the long-running response, then C runs. A is not re-run. assertThat(simplifyEvents(eventsAfterResume)) @@ -270,7 +603,6 @@ public static ImmutableMap pendingProgressTool(String message) { } @Test - @SuppressWarnings("deprecation") // Resumability flag is intentionally deprecated (partial). public void runAsync_withLongRunningCall_resumable_pausesAfterSingleModelCall_legacyShim() { TestLlm testLlm = createTestLlm( @@ -290,23 +622,10 @@ public void runAsync_withLongRunningCall_resumable_pausesAfterSingleModelCall_le /* requireConfirmation= */ false, /* isLongRunning= */ true)) .build(); - Runner runner = - Runner.builder() - .app( - App.builder() - .name("test") - .rootAgent(agent) - .resumabilityConfig( - ResumabilityConfig.builder().plainTextContinuationAutoResume(true).build()) - .build()) - .build(); - Session session = runner.sessionService().createSession("test", "user").blockingGet(); + Runner runner = shimRunner(agent); + Session session = newSession(runner); - List events = - runner - .runAsync("user", session.id(), Content.fromParts(Part.fromText("from user"))) - .toList() - .blockingGet(); + ImmutableList events = runTurn(runner, session, "from user"); // The flow paused after the single long-running call instead of re-calling the model. assertThat(testLlm.getRequests()).hasSize(1); @@ -314,7 +633,6 @@ public void runAsync_withLongRunningCall_resumable_pausesAfterSingleModelCall_le } @Test - @SuppressWarnings("deprecation") // Resumability flag is intentionally deprecated (partial). public void runAsync_loopAgentWithLongRunningSubAgent_resumable_stopsAfterFirstIteration_legacyShim() { AtomicInteger calls = new AtomicInteger(); @@ -342,30 +660,16 @@ public void runAsync_withLongRunningCall_resumable_pausesAfterSingleModelCall_le .subAgents(ImmutableList.of(inner)) .maxIterations(3) .build(); - Runner runner = - Runner.builder() - .app( - App.builder() - .name("test") - .rootAgent(loop) - .resumabilityConfig( - ResumabilityConfig.builder().plainTextContinuationAutoResume(true).build()) - .build()) - .build(); - Session session = runner.sessionService().createSession("test", "user").blockingGet(); + Runner runner = shimRunner(loop); + Session session = newSession(runner); - List unused = - runner - .runAsync("user", session.id(), Content.fromParts(Part.fromText("from user"))) - .toList() - .blockingGet(); + ImmutableList unused = runTurn(runner, session, "from user"); // Paused after the first iteration: one model call, not maxIterations. assertThat(loopLlm.getRequests()).hasSize(1); } @Test - @SuppressWarnings("deprecation") // Resumability flag is intentionally deprecated (partial). public void runAsync_parallelAgentWithLongRunningBranch_resumable_otherBranchCompletes_legacyShim() { TestLlm longRunningLlm = @@ -392,23 +696,10 @@ public void runAsync_withLongRunningCall_resumable_pausesAfterSingleModelCall_le .name("parallel") .subAgents(ImmutableList.of(longRunningBranch, plainBranch)) .build(); - Runner runner = - Runner.builder() - .app( - App.builder() - .name("test") - .rootAgent(parallel) - .resumabilityConfig( - ResumabilityConfig.builder().plainTextContinuationAutoResume(true).build()) - .build()) - .build(); - Session session = runner.sessionService().createSession("test", "user").blockingGet(); + Runner runner = shimRunner(parallel); + Session session = newSession(runner); - List events = - runner - .runAsync("user", session.id(), Content.fromParts(Part.fromText("from user"))) - .toList() - .blockingGet(); + ImmutableList events = runTurn(runner, session, "from user"); // The long-running branch paused after one model call; the other branch still completed. assertThat(longRunningLlm.getRequests()).hasSize(1); diff --git a/core/src/test/java/com/google/adk/runner/RunnerResumabilityTest.java b/core/src/test/java/com/google/adk/runner/RunnerResumabilityTest.java new file mode 100644 index 000000000..c612c52e5 --- /dev/null +++ b/core/src/test/java/com/google/adk/runner/RunnerResumabilityTest.java @@ -0,0 +1,1864 @@ +/* + * Copyright 2025 Google LLC + * + * Licensed 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 com.google.adk.runner; + +import static com.google.adk.testing.ResumabilityTestUtils.answerCall; +import static com.google.adk.testing.ResumabilityTestUtils.assertEndOfAgent; +import static com.google.adk.testing.ResumabilityTestUtils.assertNoEndOfAgent; +import static com.google.adk.testing.ResumabilityTestUtils.confirmingEchoFunctionTool; +import static com.google.adk.testing.ResumabilityTestUtils.longRunningEchoFunctionTool; +import static com.google.adk.testing.ResumabilityTestUtils.newSession; +import static com.google.adk.testing.ResumabilityTestUtils.pauseThenSay; +import static com.google.adk.testing.ResumabilityTestUtils.pausingAgent; +import static com.google.adk.testing.ResumabilityTestUtils.pendingFunctionTool; +import static com.google.adk.testing.ResumabilityTestUtils.reloadSession; +import static com.google.adk.testing.ResumabilityTestUtils.resumableRunner; +import static com.google.adk.testing.ResumabilityTestUtils.resume; +import static com.google.adk.testing.ResumabilityTestUtils.resumeFlowable; +import static com.google.adk.testing.ResumabilityTestUtils.runTurn; +import static com.google.adk.testing.ResumabilityTestUtils.textAgent; +import static com.google.adk.testing.TestUtils.createFunctionCallLlmResponse; +import static com.google.adk.testing.TestUtils.createLlmResponse; +import static com.google.adk.testing.TestUtils.createTestAgentBuilder; +import static com.google.adk.testing.TestUtils.createTestLlm; +import static com.google.adk.testing.TestUtils.createTextLlmResponse; +import static com.google.adk.testing.TestUtils.simplifyEvents; +import static com.google.adk.testing.TestUtils.simplifyResumableEvents; +import static com.google.common.collect.ImmutableList.toImmutableList; +import static com.google.common.truth.Truth.assertThat; +import static org.junit.Assert.assertThrows; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.CALLS_REAL_METHODS; +import static org.mockito.Mockito.atLeastOnce; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.times; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +import com.google.adk.agents.InvocationContext; +import com.google.adk.agents.LlmAgent; +import com.google.adk.agents.LoopAgent; +import com.google.adk.agents.ParallelAgent; +import com.google.adk.agents.RunConfig; +import com.google.adk.agents.SequentialAgent; +import com.google.adk.apps.App; +import com.google.adk.events.Event; +import com.google.adk.events.EventActions; +import com.google.adk.flows.llmflows.Functions; +import com.google.adk.plugins.BasePlugin; +import com.google.adk.sessions.Session; +import com.google.adk.telemetry.Tracing; +import com.google.adk.testing.ResumabilityTestUtils.Tools; +import com.google.adk.testing.TestBaseAgent; +import com.google.adk.testing.TestLlm; +import com.google.adk.tools.FunctionTool; +import com.google.common.collect.ImmutableList; +import com.google.common.collect.ImmutableMap; +import com.google.common.collect.Iterables; +import com.google.common.collect.Streams; +import com.google.genai.types.Content; +import com.google.genai.types.FunctionCall; +import com.google.genai.types.FunctionResponse; +import com.google.genai.types.Part; +import io.opentelemetry.api.trace.Tracer; +import io.opentelemetry.sdk.testing.junit4.OpenTelemetryRule; +import io.reactivex.rxjava3.core.Flowable; +import io.reactivex.rxjava3.core.Maybe; +import java.util.List; +import java.util.Objects; +import java.util.Optional; +import java.util.concurrent.atomic.AtomicInteger; +import org.jspecify.annotations.Nullable; +import org.junit.After; +import org.junit.Before; +import org.junit.Rule; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +/** Runner tests for durable resumability: {@code ResumabilityConfig.resumable(true)}. */ +@RunWith(JUnit4.class) +@SuppressWarnings("deprecation") // The class exists to exercise the deprecated resumability flags. +public final class RunnerResumabilityTest { + @Rule public final OpenTelemetryRule openTelemetryRule = OpenTelemetryRule.create(); + + private Tracer originalTracer; + + private BasePlugin mockPlugin(String name) { + // CALLS_REAL_METHODS keeps the default Maybe.empty() returns; a bare mock would return null. + BasePlugin plugin = mock(BasePlugin.class, CALLS_REAL_METHODS); + when(plugin.getName()).thenReturn(name); + return plugin; + } + + @Before + public void setUp() { + this.originalTracer = Tracing.getTracer(); + Tracing.setTracerForTesting( + openTelemetryRule.getOpenTelemetry().getTracer("RunnerResumabilityTest")); + } + + @After + public void tearDown() { + Tracing.setTracerForTesting(originalTracer); + } + + // A confirmation answer resumes B in SequentialAgent(A, B, C): C runs, A does not re-run. + @Test + public void runAsync_withToolConfirmation_inSequentialAgent_runsLaterSubAgentsAfterResume() { + LlmAgent agentA = + createTestAgentBuilder(createTestLlm(createTextLlmResponse("agent A done"))) + .name("a_agent") + .build(); + // B pauses right after requesting confirmation, so one follow-up response covers the resume. + TestLlm bTestLlm = + createTestLlm( + createFunctionCallLlmResponse( + "tool_call_id", "echoTool", ImmutableMap.of("message", "hello")), + createTextLlmResponse("Response after user confirmed.")); + LlmAgent agentB = + createTestAgentBuilder(bTestLlm) + .name("b_agent") + .tools(confirmingEchoFunctionTool()) + .build(); + LlmAgent agentC = + createTestAgentBuilder(createTestLlm(createTextLlmResponse("agent C done"))) + .name("c_agent") + .build(); + SequentialAgent workflowAgent = + SequentialAgent.builder() + .name("workflow_agent") + .subAgents(ImmutableList.of(agentA, agentB, agentC)) + .build(); + Runner runner = resumableRunner(workflowAgent); + Session session = newSession(runner); + + ImmutableList eventsBeforeConfirmation = runTurn(runner, session, "from user"); + + // Turn 1: A runs, B pauses for confirmation, and C must not run yet. + assertThat(simplifyEvents(eventsBeforeConfirmation)).contains("a_agent: agent A done"); + assertThat(simplifyEvents(eventsBeforeConfirmation)).doesNotContain("c_agent: agent C done"); + + FunctionCall askUserConfirmationFunctionCall = + Iterables.getOnlyElement( + eventsBeforeConfirmation.stream() + .map(Functions::getAskUserConfirmationFunctionCalls) + .filter(functionCalls -> !functionCalls.isEmpty()) + .findFirst() + .get()); + ImmutableList eventsAfterConfirmation = + ImmutableList.copyOf( + runner + .runAsync( + "user", + session.id(), + Content.fromParts( + Part.builder() + .functionResponse( + FunctionResponse.builder() + .id(askUserConfirmationFunctionCall.id().get()) + .name(askUserConfirmationFunctionCall.name().get()) + .response(ImmutableMap.of("confirmed", true))) + .build())) + .toList() + .blockingGet()); + + // Turn 2: B resumes and runs the tool, then C runs; A is not re-run. + assertThat(simplifyResumableEvents(eventsAfterConfirmation)) + .containsExactly( + "b_agent: FunctionResponse(name=echoTool, response={message=hello})", + "b_agent: Response after user confirmed.", + "b_agent: end_of_agent", + "workflow_agent: agent_state={current_sub_agent=c_agent}", + "c_agent: agent C done", + "c_agent: end_of_agent", + "workflow_agent: end_of_agent") + .inOrder(); + } + + // A pending long-running call pauses SequentialAgent(A, B, C) at B; on resume C runs, not A. + @Test + public void runAsync_withLongRunningCall_inSequentialAgent_runsLaterSubAgentsAfterResume() { + LlmAgent agentA = + createTestAgentBuilder(createTestLlm(createTextLlmResponse("agent A done"))) + .name("a_agent") + .build(); + // B pauses right after the no-result long-running call, so one response covers the resume. + TestLlm bTestLlm = + createTestLlm( + createFunctionCallLlmResponse( + "lro_call_id", "pendingTool", ImmutableMap.of("message", "hello")), + createTextLlmResponse("agent B resumed")); + LlmAgent agentB = + createTestAgentBuilder(bTestLlm).name("b_agent").tools(pendingFunctionTool()).build(); + LlmAgent agentC = + createTestAgentBuilder(createTestLlm(createTextLlmResponse("agent C done"))) + .name("c_agent") + .build(); + SequentialAgent workflowAgent = + SequentialAgent.builder() + .name("workflow_agent") + .subAgents(ImmutableList.of(agentA, agentB, agentC)) + .build(); + Runner runner = resumableRunner(workflowAgent); + Session session = newSession(runner); + + ImmutableList eventsBeforeResume = runTurn(runner, session, "from user"); + + // Turn 1: A runs, B issues the long-running call and pauses; C must not run yet. + assertThat(simplifyEvents(eventsBeforeResume)).contains("a_agent: agent A done"); + assertThat(simplifyEvents(eventsBeforeResume)).doesNotContain("b_agent: agent B resumed"); + assertThat(simplifyEvents(eventsBeforeResume)).doesNotContain("c_agent: agent C done"); + + ImmutableList eventsAfterResume = + answerCall( + runner, session, "lro_call_id", "pendingTool", ImmutableMap.of("message", "hello")); + + // Turn 2: B resumes from the long-running response and C runs; A is not re-run. + assertThat(simplifyResumableEvents(eventsAfterResume)) + .containsExactly( + "b_agent: agent B resumed", + "b_agent: end_of_agent", + "workflow_agent: agent_state={current_sub_agent=c_agent}", + "c_agent: agent C done", + "c_agent: end_of_agent", + "workflow_agent: end_of_agent") + .inOrder(); + } + + // A LoopAgent(w1, w2) paused on w1 resumes w1, then advances the loop to w2 and closes it. + @Test + public void runAsync_withLongRunningCall_inLoopAgent_runsRemainingSubAgentsAfterResume() { + TestLlm w1TestLlm = + createTestLlm( + createFunctionCallLlmResponse( + "lro_call_id", "pendingTool", ImmutableMap.of("message", "hello")), + createTextLlmResponse("w1 resumed")); + LlmAgent w1 = + createTestAgentBuilder(w1TestLlm).name("w1_agent").tools(pendingFunctionTool()).build(); + LlmAgent w2 = + createTestAgentBuilder(createTestLlm(createTextLlmResponse("w2 done"))) + .name("w2_agent") + .build(); + LoopAgent workflowAgent = + LoopAgent.builder() + .name("loop_agent") + .subAgents(ImmutableList.of(w1, w2)) + .maxIterations(1) + .build(); + Runner runner = resumableRunner(workflowAgent); + Session session = newSession(runner); + + ImmutableList eventsBeforeResume = runTurn(runner, session, "from user"); + + // Turn 1: w1 issues the long-running call and pauses; w2 must not run yet. + assertThat(simplifyEvents(eventsBeforeResume)).doesNotContain("w1_agent: w1 resumed"); + assertThat(simplifyEvents(eventsBeforeResume)).doesNotContain("w2_agent: w2 done"); + + ImmutableList eventsAfterResume = + answerCall( + runner, session, "lro_call_id", "pendingTool", ImmutableMap.of("message", "hello")); + + // Turn 2: w1 resumes, then the loop advances to w2 and closes; w1 is not re-run. + assertThat(simplifyResumableEvents(eventsAfterResume)) + .containsExactly( + "w1_agent: w1 resumed", + "w1_agent: end_of_agent", + "loop_agent: agent_state={current_sub_agent=w2_agent, times_looped=0}", + "w2_agent: w2 done", + "w2_agent: end_of_agent", + "loop_agent: end_of_agent") + .inOrder(); + } + + // A transfer runs the sub-agent nested, as in Python: the root closes only after it completes. + @Test + public void runAsync_resumable_transferToSubAgent_closesRootAfterSubAgentCompletes() { + Content transferCall = + Content.fromParts( + Part.fromFunctionCall( + "transfer_to_agent", ImmutableMap.of("agent_name", "sub_agent_1"))); + TestLlm testLlm = + createTestLlm( + createLlmResponse(transferCall), + createTextLlmResponse("response1"), + createTextLlmResponse("response2")); + LlmAgent subAgent1 = createTestAgentBuilder(testLlm).name("sub_agent_1").build(); + LlmAgent rootAgent = + createTestAgentBuilder(testLlm) + .name("root_agent") + .subAgents(ImmutableList.of(subAgent1)) + .build(); + Runner runner = resumableRunner(rootAgent); + Session session = newSession(runner); + + ImmutableList turn1 = runTurn(runner, session, "hi"); + + // The sub-agent runs inside the root's turn; the root closes last. + assertThat(simplifyResumableEvents(turn1)) + .containsExactly( + "root_agent: FunctionCall(name=transfer_to_agent, args={agent_name=sub_agent_1})", + "root_agent: FunctionResponse(name=transfer_to_agent, response={})", + "sub_agent_1: response1", + "sub_agent_1: end_of_agent", + "root_agent: end_of_agent") + .inOrder(); + + ImmutableList turn2 = runTurn(runner, session, "again"); + + // A new turn resumes at the transferred sub-agent, not the finished root. + assertThat(simplifyEvents(turn2)).contains("sub_agent_1: response2"); + } + + // A transferred-to sub-agent that pauses on a long-running call resumes in the same invocation. + @Test + public void runAsync_resumable_transferredSubAgentPausesOnLongRunningCall_resumesSubAgent() { + Content transferCall = + Content.fromParts( + Part.fromFunctionCall( + "transfer_to_agent", ImmutableMap.of("agent_name", "sub_agent_1"))); + TestLlm testLlm = + createTestLlm( + createLlmResponse(transferCall), + createFunctionCallLlmResponse( + "lro_call_id", "pendingTool", ImmutableMap.of("message", "hello")), + createTextLlmResponse("resumed answer")); + LlmAgent subAgent1 = + createTestAgentBuilder(testLlm).name("sub_agent_1").tools(pendingFunctionTool()).build(); + LlmAgent rootAgent = + createTestAgentBuilder(testLlm) + .name("root_agent") + .subAgents(ImmutableList.of(subAgent1)) + .build(); + Runner runner = resumableRunner(rootAgent); + Session session = newSession(runner); + + ImmutableList pausedTurn = runTurn(runner, session, "hi"); + String pausedInvocationId = pausedTurn.get(0).invocationId(); + + // The sub-agent runs nested and pauses, so neither it nor the root closes. + ImmutableList pausedEvents = simplifyResumableEvents(pausedTurn); + assertThat(pausedEvents) + .contains("sub_agent_1: FunctionCall(name=pendingTool, args={message=hello})"); + assertThat(pausedEvents).doesNotContain("root_agent: end_of_agent"); + assertThat(pausedEvents).doesNotContain("sub_agent_1: end_of_agent"); + assertThat(simplifyEvents(pausedTurn)).doesNotContain("sub_agent_1: resumed answer"); + + ImmutableList resumed = + ImmutableList.copyOf( + runner + .runAsync( + "user", + session.id(), + /* invocationId= */ null, + Content.fromParts( + Part.builder() + .functionResponse( + FunctionResponse.builder() + .id("lro_call_id") + .name("pendingTool") + .response(ImmutableMap.of("message", "hello"))) + .build()), + RunConfig.builder().build(), + /* stateDelta= */ null) + .toList() + .blockingGet()); + + // The resume continues the transferred sub-agent (same invocation) to completion. + assertThat(resumed).isNotEmpty(); + assertThat(resumed.stream().allMatch(event -> event.invocationId().equals(pausedInvocationId))) + .isTrue(); + assertThat(simplifyEvents(resumed)).contains("sub_agent_1: resumed answer"); + // The root is an LlmAgent ancestor of the paused agent, so the resume re-enters it rather than + // routing past it, and it closes itself once the sub-agent finishes. + assertThat(simplifyResumableEvents(resumed)).contains("root_agent: end_of_agent"); + } + + // An invocation paused on two long-running calls stays paused until both are answered. + @Test + public void runAsync_withTwoLongRunningCalls_pausesUntilBothAnswered() { + TestLlm testLlm = + createTestLlm( + createLlmResponse( + Content.builder() + .role("model") + .parts( + Part.builder() + .functionCall( + FunctionCall.builder() + .id("call_a") + .name("pendingTool") + .args(ImmutableMap.of("message", "a"))) + .build(), + Part.builder() + .functionCall( + FunctionCall.builder() + .id("call_b") + .name("pendingTool") + .args(ImmutableMap.of("message", "b"))) + .build()) + .build()), + createTextLlmResponse("both approved")); + LlmAgent agent = + createTestAgentBuilder(testLlm).name("root_agent").tools(pendingFunctionTool()).build(); + Runner runner = resumableRunner(agent); + Session session = newSession(runner); + + runTurn(runner, session, "start"); + // Turn 1: both calls are issued and the invocation pauses after one model call. + assertThat(testLlm.getRequests()).hasSize(1); + + ImmutableList afterFirstAnswer = + answerCall(runner, session, "call_a", "pendingTool", ImmutableMap.of("message", "a")); + // One call answered is not enough: nothing runs and the model is not re-invoked. + assertThat(afterFirstAnswer).isEmpty(); + assertThat(testLlm.getRequests()).hasSize(1); + + ImmutableList afterSecondAnswer = + answerCall(runner, session, "call_b", "pendingTool", ImmutableMap.of("message", "b")); + // Both answered: the model is re-invoked and summarizes. + assertThat(testLlm.getRequests()).hasSize(2); + assertThat(simplifyEvents(afterSecondAnswer)).contains("root_agent: both approved"); + } + + // A long-running tool that returns a value resolves in-turn, so the flow summarizes, not pauses. + @Test + public void runAsync_withLongRunningCall_resumable_valueReturn_continuesAndSummarizes() { + TestLlm testLlm = + createTestLlm( + createFunctionCallLlmResponse( + "lro_call_id", "echoTool", ImmutableMap.of("message", "hello")), + createTextLlmResponse("summarized echo")); + LlmAgent agent = + createTestAgentBuilder(testLlm).name("agent").tools(longRunningEchoFunctionTool()).build(); + Runner runner = resumableRunner(agent); + Session session = newSession(runner); + + ImmutableList events = runTurn(runner, session, "from user"); + + // Summarized in-turn with no pause, and no end-of-agent: the invocation stays resumable. + assertThat(testLlm.getRequests()).hasSize(2); + assertThat(simplifyEvents(events)).contains("agent: summarized echo"); + assertThat(events.stream().anyMatch(event -> event.actions().endOfAgent())).isFalse(); + } + + // A resume runs the whole plugin bracket, not only on-event, as the new-invocation path does. + @Test + public void runAsync_resume_runsFullPluginBracket() { + BasePlugin resumePlugin = mockPlugin("resume"); + TestLlm testLlm = + createTestLlm( + createFunctionCallLlmResponse( + "lro_call_id", "pendingTool", ImmutableMap.of("message", "hello")), + createTextLlmResponse("resumed and summarized")); + LlmAgent agent = + createTestAgentBuilder(testLlm).name("agent").tools(pendingFunctionTool()).build(); + Runner runner = resumableRunner(agent, resumePlugin); + Session session = newSession(runner); + + runTurn(runner, session, "start"); + + ImmutableList resumed = + answerCall( + runner, session, "lro_call_id", "pendingTool", ImmutableMap.of("message", "hello")); + + assertThat(simplifyEvents(resumed)).contains("agent: resumed and summarized"); + // Twice: once on the initial turn, once on the resume turn, which used to run only on-event. + verify(resumePlugin, times(2)).onUserMessageCallback(any(), any()); + verify(resumePlugin, times(2)).beforeRunCallback(any()); + verify(resumePlugin, times(2)).afterRunCallback(any()); + verify(resumePlugin, atLeastOnce()).onEventCallback(any(), any()); + } + + // An unanswered call keeps the invocation paused: plain text and no message are both inert. + @Test + public void resume_pausedCallUnanswered_staysPausedForPlainTextAndNoMessage() { + TestLlm testLlm = + createTestLlm( + createFunctionCallLlmResponse( + "lro_call_id", "pendingTool", ImmutableMap.of("message", "hello")), + createTextLlmResponse("should not be reached")); + LlmAgent agent = + createTestAgentBuilder(testLlm).name("agent").tools(pendingFunctionTool()).build(); + Runner runner = resumableRunner(agent); + Session session = newSession(runner); + + ImmutableList firstTurn = runTurn(runner, session, "start"); + String invocationId = firstTurn.get(0).invocationId(); + + // A plain-text message with an explicit invocation id returns instead of throwing. + ImmutableList resumedWithText = + resume( + runner, + session, + invocationId, + Content.fromParts(Part.fromText("please continue")), + /* stateDelta= */ null); + assertThat(resumedWithText).isEmpty(); + + ImmutableList resumedWithNoMessage = + resume(runner, session, invocationId, /* newMessage= */ null, /* stateDelta= */ null); + assertThat(resumedWithNoMessage).isEmpty(); + // The call is still unanswered, so the model was never re-invoked past the pausing turn. + assertThat(testLlm.getRequests()).hasSize(1); + } + + // A silent pause inside a nested Sequential leaves the outer checkpoint on the inner workflow. + @Test + public void runAsync_resume_nestedSequentialSilentPause_checkpointStaysOnUnfinishedInner() { + LlmAgent a1 = pausingAgent("a1_agent", pauseThenSay("lro_call_id", "a1 resumed")); + LlmAgent a2 = textAgent("a2_agent", "a2 done"); + SequentialAgent inner = + SequentialAgent.builder().name("inner_agent").subAgents(ImmutableList.of(a1, a2)).build(); + LlmAgent b = textAgent("b_agent", "b done"); + SequentialAgent outer = + SequentialAgent.builder().name("outer_agent").subAgents(ImmutableList.of(inner, b)).build(); + Runner runner = resumableRunner(outer); + Session session = newSession(runner); + + ImmutableList turn1 = runTurn(runner, session, "start"); + // a1 pauses, so neither a2 nor b runs, and no workflow closes. + assertThat(simplifyEvents(turn1)).doesNotContain("a2_agent: a2 done"); + assertThat(simplifyEvents(turn1)).doesNotContain("b_agent: b done"); + assertNoEndOfAgent(turn1, "outer_agent"); + // The outer checkpoint still names the unfinished inner workflow. + assertThat(simplifyResumableEvents(turn1)) + .contains("outer_agent: agent_state={current_sub_agent=inner_agent}"); + + ImmutableList resumed = + answerCall(runner, session, "lro_call_id", "pendingTool", ImmutableMap.of("message", "hi")); + + // The resume finishes the inner workflow before advancing the outer one; nothing re-runs. + assertThat(simplifyEvents(resumed)) + .containsAtLeast("a1_agent: a1 resumed", "a2_agent: a2 done", "b_agent: b done") + .inOrder(); + assertEndOfAgent(resumed, "outer_agent"); + } + + // A leaf paused under a ParallelAgent resumes, and the enclosing sequence advances past it. + @Test + public void runAsync_resume_pausedUnderParallelAgent_advancesEnclosingSequential() { + LlmAgent leaf = + createTestAgentBuilder( + createTestLlm( + createFunctionCallLlmResponse( + "lro_call_id", "pendingTool", ImmutableMap.of("message", "hi")), + createTextLlmResponse("leaf resumed"))) + .name("leaf_agent") + .tools(pendingFunctionTool()) + .build(); + ParallelAgent parallel = + ParallelAgent.builder().name("parallel_agent").subAgents(ImmutableList.of(leaf)).build(); + LlmAgent next = + createTestAgentBuilder(createTestLlm(createTextLlmResponse("next ran"))) + .name("next_agent") + .build(); + SequentialAgent root = + SequentialAgent.builder() + .name("seq_agent") + .subAgents(ImmutableList.of(parallel, next)) + .build(); + Runner runner = resumableRunner(root); + Session session = newSession(runner); + + runTurn(runner, session, "start"); + + ImmutableList resumed = + answerCall(runner, session, "lro_call_id", "pendingTool", ImmutableMap.of("message", "hi")); + + assertThat(simplifyEvents(resumed)).contains("leaf_agent: leaf resumed"); + // The enclosing sequence then advances past the parallel block. + assertThat(simplifyEvents(resumed)).contains("next_agent: next ran"); + } + + // A leaf under a ParallelAgent paused on two calls, answered once, stays paused. + @Test + public void runAsync_resume_pausedUnderParallelAgent_partiallyAnswered_staysPaused() { + Content twoLongRunningCalls = + Content.builder() + .role("model") + .parts( + Part.builder() + .functionCall( + FunctionCall.builder() + .id("c1") + .name("pendingTool") + .args(ImmutableMap.of("message", "a")) + .build()) + .build(), + Part.builder() + .functionCall( + FunctionCall.builder() + .id("c2") + .name("pendingTool") + .args(ImmutableMap.of("message", "b")) + .build()) + .build()) + .build(); + LlmAgent leaf = + createTestAgentBuilder( + createTestLlm( + createLlmResponse(twoLongRunningCalls), createTextLlmResponse("leaf summary"))) + .name("leaf_agent") + .tools(pendingFunctionTool()) + .build(); + ParallelAgent parallel = + ParallelAgent.builder().name("parallel_agent").subAgents(ImmutableList.of(leaf)).build(); + SequentialAgent root = + SequentialAgent.builder().name("seq_agent").subAgents(ImmutableList.of(parallel)).build(); + Runner runner = resumableRunner(root); + Session session = newSession(runner); + + runTurn(runner, session, "start"); + + // Answer only c1; c2 remains unanswered. + ImmutableList resumed = + answerCall(runner, session, "c1", "pendingTool", ImmutableMap.of("message", "a")); + + // c2 is still unanswered, so the model is not re-invoked: no "leaf summary", no new events. + assertThat(simplifyEvents(resumed)).doesNotContain("leaf_agent: leaf summary"); + assertThat(resumed).isEmpty(); + } + + // Two paused sibling leaves: answering one resumes it while the other holds the sequence. + @Test + public void runAsync_resume_parallelSiblingsPartiallyAnswered_answeredSiblingResumes() { + LlmAgent leafA = + createTestAgentBuilder( + createTestLlm( + createFunctionCallLlmResponse( + "c1", "pendingTool", ImmutableMap.of("message", "a")), + createTextLlmResponse("a resumed"))) + .name("leaf_a") + .tools(pendingFunctionTool()) + .build(); + LlmAgent leafB = + createTestAgentBuilder( + createTestLlm( + createFunctionCallLlmResponse( + "c2", "pendingTool", ImmutableMap.of("message", "b")), + createTextLlmResponse("b resumed"))) + .name("leaf_b") + .tools(pendingFunctionTool()) + .build(); + ParallelAgent parallel = + ParallelAgent.builder() + .name("parallel_agent") + .subAgents(ImmutableList.of(leafA, leafB)) + .build(); + LlmAgent next = + createTestAgentBuilder(createTestLlm(createTextLlmResponse("next ran"))) + .name("next_agent") + .build(); + SequentialAgent root = + SequentialAgent.builder() + .name("seq_agent") + .subAgents(ImmutableList.of(parallel, next)) + .build(); + Runner runner = resumableRunner(root); + Session session = newSession(runner); + + runTurn(runner, session, "start"); + + // Answer only leaf_a's call; leaf_b's stays unanswered. + ImmutableList resumed = + answerCall(runner, session, "c1", "pendingTool", ImmutableMap.of("message", "a")); + + assertThat(simplifyEvents(resumed)).contains("leaf_a: a resumed"); + assertThat(simplifyEvents(resumed)).doesNotContain("leaf_b: b resumed"); + // leaf_b is still pending, so the sequence must not advance past the parallel block. + assertThat(simplifyEvents(resumed)).doesNotContain("next_agent: next ran"); + // leaf_b must stay resumable, so nothing above it may close. + assertNoEndOfAgent(resumed, "leaf_b"); + assertNoEndOfAgent(resumed, "parallel_agent"); + assertNoEndOfAgent(resumed, "seq_agent"); + assertEndOfAgent(resumed, "leaf_a"); + } + + // With the shim off, a plain-text continuation after a pause starts a new invocation. + @Test + public void runAsync_plainTextContinuation_autoResumeFlagOff_startsNewInvocation() { + TestLlm testLlm = + createTestLlm( + createFunctionCallLlmResponse( + "lro_call_id", "pendingTool", ImmutableMap.of("message", "draft")), + createTextLlmResponse("re-planned")); + LlmAgent agent = + createTestAgentBuilder(testLlm).name("agent").tools(pendingFunctionTool()).build(); + Runner runner = resumableRunner(agent); + Session session = newSession(runner); + + ImmutableList turn1 = runTurn(runner, session, "draft the note"); + String invocationId = turn1.get(0).invocationId(); + + ImmutableList turn2 = runTurn(runner, session, "Proceed"); + + assertThat(testLlm.getRequests()).hasSize(2); // new invocation re-invoked the model + assertThat(turn2.get(0).invocationId()).isNotEqualTo(invocationId); + } + + // With resumability off, a completed LlmAgent emits no end-of-agent checkpoint. + @Test + public void runAsync_resumabilityDisabled_completedLlmAgent_emitsNoEndOfAgent() { + LlmAgent agent = + createTestAgentBuilder(createTestLlm(createTextLlmResponse("all done"))) + .name("agent") + .build(); + Runner runner = + Runner.builder().app(App.builder().name("test").rootAgent(agent).build()).build(); + Session session = newSession(runner); + + ImmutableList events = runTurn(runner, session, "from user"); + + assertThat(events.stream().anyMatch(event -> event.actions().endOfAgent())).isFalse(); + } + + // A completed resumable LlmAgent emits a trailing end-of-agent; resuming past it is a no-op. + @Test + public void resume_completedInvocation_emitsEndOfAgentThenIsNoOp() { + LlmAgent agent = + createTestAgentBuilder(createTestLlm(createTextLlmResponse("all done"))) + .name("agent") + .build(); + Runner runner = resumableRunner(agent); + Session session = newSession(runner); + + ImmutableList firstTurn = runTurn(runner, session, "from user"); + + // The end-of-agent checkpoint is how a later run can tell the invocation finished. + Event last = Iterables.getLast(firstTurn); + assertThat(last.author()).isEqualTo("agent"); + assertThat(last.actions().endOfAgent()).isTrue(); + + ImmutableList resumed = + resume( + runner, + session, + firstTurn.get(0).invocationId(), + /* newMessage= */ null, + /* stateDelta= */ null); + + assertThat(resumed).isEmpty(); + } + + // A function-response resume continues the matching invocation and merges the stateDelta. + @Test + public void resume_withFunctionResponseAndStateDelta_resumesSameInvocationAndMergesState() { + TestLlm testLlm = + createTestLlm( + createFunctionCallLlmResponse( + "lro_call_id", "pendingTool", ImmutableMap.of("message", "hello")), + createTextLlmResponse("resumed answer")); + LlmAgent agent = + createTestAgentBuilder(testLlm).name("agent").tools(pendingFunctionTool()).build(); + Runner runner = resumableRunner(agent); + Session session = newSession(runner); + + ImmutableList pausedTurn = runTurn(runner, session, "from user"); + String pausedInvocationId = pausedTurn.get(0).invocationId(); + assertThat(simplifyEvents(pausedTurn)).doesNotContain("agent: resumed answer"); + + ImmutableMap stateDelta = ImmutableMap.of("key1", "value1", "key2", 42); + ImmutableList resumed = + ImmutableList.copyOf( + runner + .runAsync( + "user", + session.id(), + /* invocationId= */ null, + Content.fromParts( + Part.builder() + .functionResponse( + FunctionResponse.builder() + .id("lro_call_id") + .name("pendingTool") + .response(ImmutableMap.of("message", "hello"))) + .build()), + RunConfig.builder().build(), + stateDelta) + .toList() + .blockingGet()); + + // The resumed events belong to the original (paused) invocation, not a fresh one. + assertThat(resumed).isNotEmpty(); + assertThat(resumed.stream().allMatch(event -> event.invocationId().equals(pausedInvocationId))) + .isTrue(); + assertThat(simplifyEvents(resumed)).contains("agent: resumed answer"); + assertThat(resumed.stream().anyMatch(event -> event.actions().endOfAgent())).isTrue(); + Session finalSession = + runner + .sessionService() + .getSession("test", "user", session.id(), Optional.empty()) + .blockingGet(); + assertThat(finalSession.state()).containsAtLeastEntriesIn(stateDelta); + // The delta is also stamped on the appended (function-response) event, for history rehydration. + Event lastUserEvent = + Streams.findLast( + finalSession.events().stream() + .filter(event -> Objects.equals(event.author(), "user"))) + .orElseThrow(); + assertThat(lastUserEvent.actions().stateDelta()).containsAtLeastEntriesIn(stateDelta); + } + + // Python parity (REPLAY_CALLS): a resume runs a persisted, unexecuted call, not the model. + @Test + public void resume_unexecutedFunctionCall_replaysCallInsteadOfCallingModel() { + TestLlm testLlm = createTestLlm(createTextLlmResponse("should not be reached")); + LlmAgent agent = + createTestAgentBuilder(testLlm) + .name("agent") + .tools(FunctionTool.create(Tools.class, "echoTool")) + .build(); + Runner runner = resumableRunner(agent); + Session session = newSession(runner); + String invocationId = "inv-replay"; + var unusedUser = + runner + .sessionService() + .appendEvent( + session, + Event.builder() + .id("u1") + .invocationId(invocationId) + .author("user") + .content(Content.fromParts(Part.fromText("go"))) + .build()) + .blockingGet(); + // The agent's call was persisted but never executed: no function response follows it. + var unusedCall = + runner + .sessionService() + .appendEvent( + session, + Event.builder() + .id("c1") + .invocationId(invocationId) + .author("agent") + .content( + Content.fromParts( + Part.builder() + .functionCall( + FunctionCall.builder() + .id("call_1") + .name("echoTool") + .args(ImmutableMap.of("message", "hi")) + .build()) + .build())) + .build()) + .blockingGet(); + + ImmutableList resumed = + resume(runner, session, invocationId, /* newMessage= */ null, /* stateDelta= */ null); + + // The scripted model only returns text, so an echoTool response can only come from replay. + assertThat( + resumed.stream() + .flatMap(event -> event.functionResponses().stream()) + .map(response -> response.name().orElse("")) + .collect(toImmutableList())) + .contains("echoTool"); + // Replay precedes the model, so the first request already carries the tool's response. + assertThat(testLlm.getRequests().get(0).contents().toString()).contains("echoTool"); + } + + // Python restricts a transfer to the agent's declared targets; a replay must not get around it. + @Test + public void resume_replayedTransferToUndeclaredAgent_fails() { + TestLlm peerLlm = createTestLlm(createTextLlmResponse("peer child should not run")); + LlmAgent peerChild = createTestAgentBuilder(peerLlm).name("peer_child").build(); + LlmAgent peer = + createTestAgentBuilder(createTestLlm(createTextLlmResponse("peer"))) + .name("peer") + .subAgents(peerChild) + .build(); + // No sub-agents and no peer transfers, so peer_child is not a declared target of agent. + LlmAgent agent = + createTestAgentBuilder(createTestLlm(createTextLlmResponse("should not be reached"))) + .name("agent") + .disallowTransferToPeers(true) + .build(); + LlmAgent root = + createTestAgentBuilder(createTestLlm(createTextLlmResponse("root"))) + .name("root") + .subAgents(agent, peer) + .build(); + Runner runner = resumableRunner(root); + Session session = newSession(runner); + String invocationId = "inv-bad-transfer"; + var unusedUser = + runner + .sessionService() + .appendEvent( + session, + Event.builder() + .id("u1") + .invocationId(invocationId) + .author("user") + .content(Content.fromParts(Part.fromText("go"))) + .build()) + .blockingGet(); + // A transfer call persisted but never executed, naming an agent this one may not transfer to. + var unusedCall = + runner + .sessionService() + .appendEvent( + session, + Event.builder() + .id("c1") + .invocationId(invocationId) + .author("agent") + .content( + Content.fromParts( + Part.builder() + .functionCall( + FunctionCall.builder() + .id("call_1") + .name("transfer_to_agent") + .args(ImmutableMap.of("agent_name", "peer_child")) + .build()) + .build())) + .build()) + .blockingGet(); + + resumeFlowable(runner, session, invocationId, /* newMessage= */ null, /* stateDelta= */ null) + .test() + .assertError( + error -> + error instanceof IllegalStateException + && error.getMessage().contains("is not allowed to transfer")); + assertThat(peerLlm.getRequests()).isEmpty(); + } + + // Python parity: a stateDelta resumed with no message is persisted as a content-less event. + @Test + public void resume_withStateDeltaAndNoMessage_persistsStateDeltaEvent() { + TestLlm testLlm = + createTestLlm( + createFunctionCallLlmResponse( + "lro_call_id", "pendingTool", ImmutableMap.of("message", "draft")), + createTextLlmResponse("done")); + LlmAgent agent = + createTestAgentBuilder(testLlm).name("agent").tools(pendingFunctionTool()).build(); + Runner runner = resumableRunner(agent); + Session session = newSession(runner); + ImmutableList turn1 = runTurn(runner, session, "start"); + String invocationId = turn1.get(0).invocationId(); + + ImmutableMap stateDelta = ImmutableMap.of("key1", "value1"); + var unused = resume(runner, session, invocationId, /* newMessage= */ null, stateDelta); + + Session reloaded = + runner + .sessionService() + .getSession("test", "user", session.id(), Optional.empty()) + .blockingGet(); + assertThat(reloaded.state()).containsAtLeastEntriesIn(stateDelta); + assertThat( + reloaded.events().stream() + .anyMatch( + event -> + event.content().isEmpty() + && event.actions().stateDelta().containsKey("key1"))) + .isTrue(); + } + + // ResumeInvocationTest parity: a function response resumes an older invocation, not the latest. + @Test + public void resume_resumesAnyInvocation_notJustTheLatest() { + TestLlm testLlm = + createTestLlm( + createFunctionCallLlmResponse( + "call-1", "pendingTool", ImmutableMap.of("message", "hi")), + createTextLlmResponse("llm response in invocation 2"), + createFunctionCallLlmResponse( + "call-3", "pendingTool", ImmutableMap.of("message", "hi")), + createTextLlmResponse("llm response after resuming invocation 1")); + LlmAgent agent = + createTestAgentBuilder(testLlm).name("agent").tools(pendingFunctionTool()).build(); + Runner runner = resumableRunner(agent); + Session session = newSession(runner); + + // Invocation 1 pauses on the long-running call. + ImmutableList inv1 = runTurn(runner, session, "q1"); + String inv1Id = inv1.get(0).invocationId(); + // Invocation 2 finishes; invocation 3 pauses again. + runTurn(runner, session, "q2"); + runTurn(runner, session, "q3"); + + // Resume invocation 1 (the oldest, not the latest) via its function response. + ImmutableList resumed = + ImmutableList.copyOf( + runner + .runAsync( + "user", + session.id(), + inv1Id, + Content.fromParts( + Part.builder() + .functionResponse( + FunctionResponse.builder() + .id("call-1") + .name("pendingTool") + .response(ImmutableMap.of("message", "hi"))) + .build()), + RunConfig.builder().build(), + /* stateDelta= */ null) + .toList() + .blockingGet()); + + assertThat(simplifyEvents(resumed)).contains("agent: llm response after resuming invocation 1"); + assertThat(resumed.stream().allMatch(event -> event.invocationId().equals(inv1Id))).isTrue(); + } + + // InMemoryRunnerTest parity: resume by invocationId rehydrates the agent's state from history. + @Test + public void resume_restoresAgentStateFromHistory() { + TestBaseAgent agent = + new TestBaseAgent( + "test_agent", + "desc", + () -> Flowable.empty(), + /* subAgents= */ null, + /* beforeAgentCallbacks= */ null, + /* afterAgentCallbacks= */ null); + Runner runner = resumableRunner(agent); + Session session = newSession(runner); + Object unusedUser = + runner + .sessionService() + .appendEvent( + session, + Event.builder() + .id("u1") + .invocationId("test-inv") + .author("user") + .content(Content.fromParts(Part.fromText("hi"))) + .build()) + .blockingGet(); + Object unusedState = + runner + .sessionService() + .appendEvent( + session, + Event.builder() + .id("s1") + .invocationId("test-inv") + .author("test_agent") + .actions( + EventActions.builder() + .agentState(ImmutableMap.of("saved", "state")) + .build()) + .content(Content.fromParts(Part.fromText("previous response"))) + .build()) + .blockingGet(); + + Object unused = + runner + .runAsync( + "user", + session.id(), + "test-inv", + /* newMessage= */ null, + RunConfig.builder().build(), + /* stateDelta= */ null) + .toList() + .blockingGet(); + + assertThat(agent.getLastInvocationContext().agentStates()) + .containsEntry("test_agent", ImmutableMap.of("saved", "state")); + } + + // Python parity: resume text is dropped when the invocation already has a user message. + @Test + public void resume_withTextMessage_invocationAlreadyHasUserMessage_doesNotAppendIt() { + TestBaseAgent agent = + new TestBaseAgent("test_agent", "desc", () -> Flowable.empty(), null, null, null); + Runner runner = resumableRunner(agent); + Session session = newSession(runner); + Object unusedUser = + runner + .sessionService() + .appendEvent( + session, + Event.builder() + .id("u1") + .invocationId("test-inv") + .author("user") + .content(Content.fromParts(Part.fromText("hi"))) + .build()) + .blockingGet(); + + Object unused = + runner + .runAsync( + "user", + session.id(), + "test-inv", + Content.fromParts(Part.fromText("New message")), + RunConfig.builder().build(), + /* stateDelta= */ null) + .toList() + .blockingGet(); + + Session reloaded = + runner + .sessionService() + .getSession("test", "user", session.id(), Optional.empty()) + .blockingGet(); + // Only the original user message remains: the resume text was not appended. + assertThat(reloaded.events()).hasSize(1); + assertThat( + Iterables.getLast(reloaded.events()) + .content() + .flatMap(Content::parts) + .get() + .get(0) + .text()) + .hasValue("hi"); + } + + // The same resume does append the text when the invocation has no user message to duplicate. + @Test + public void resume_withTextMessage_invocationHasNoUserMessage_appendsIt() { + TestBaseAgent agent = + new TestBaseAgent("test_agent", "desc", () -> Flowable.empty(), null, null, null); + Runner runner = resumableRunner(agent); + Session session = newSession(runner); + Object unusedSeed = + runner + .sessionService() + .appendEvent( + session, + Event.builder() + .id("a1") + .invocationId("test-inv") + .author("test_agent") + .content(Content.fromParts(Part.fromText("agent spoke first"))) + .build()) + .blockingGet(); + + Object unused = + resume( + runner, + session, + "test-inv", + Content.fromParts(Part.fromText("New message")), + /* stateDelta= */ null); + + Session reloaded = reloadSession(runner, session); + assertThat(reloaded.events()).hasSize(2); + assertThat( + Iterables.getLast(reloaded.events()) + .content() + .flatMap(Content::parts) + .get() + .get(0) + .text()) + .hasValue("New message"); + } + + // RunnerTest parity (disabled counterpart): resuming a non-resumable app throws. + @Test + public void resume_notResumable_throwsException() { + LlmAgent agent = + createTestAgentBuilder(createTestLlm(createTextLlmResponse("x"))).name("agent").build(); + Runner runner = + Runner.builder().app(App.builder().name("test").rootAgent(agent).build()).build(); + Session session = newSession(runner); + String sessionId = session.id(); + + RunConfig runConfig = RunConfig.builder().build(); + assertThrows( + IllegalStateException.class, + () -> + runner.runAsync( + "user", + sessionId, + "some-inv", + /* newMessage= */ null, + runConfig, + /* stateDelta= */ null)); + } + + // With no invocation id, an orphan function response is rejected rather than fed to the model. + @Test + public void resume_functionResponseWithNoMatchingCall_throwsIllegalArgument() { + LlmAgent agent = + createTestAgentBuilder(createTestLlm(createTextLlmResponse("x"))).name("agent").build(); + Runner runner = resumableRunner(agent); + Session session = newSession(runner); + + Content orphanResponse = + Content.fromParts( + Part.builder() + .functionResponse( + FunctionResponse.builder() + .id("no_such_call") + .name("pendingTool") + .response(ImmutableMap.of("status", "done"))) + .build()); + + runner + .runAsync( + "user", + session.id(), + /* invocationId= */ null, + orphanResponse, + RunConfig.builder().build(), + /* stateDelta= */ null) + .test() + .assertError(IllegalArgumentException.class); + } + + // Python defaults a role-less message to "user"; without it the LLM request drops the response. + @Test + public void resume_messageWithoutRole_persistsItAsAUserTurnTheModelCanSee() { + TestLlm testLlm = + createTestLlm( + createFunctionCallLlmResponse( + "lro_call_id", "pendingTool", ImmutableMap.of("message", "hi")), + createTextLlmResponse("resumed answer")); + LlmAgent agent = + createTestAgentBuilder(testLlm).name("agent").tools(pendingFunctionTool()).build(); + Runner runner = resumableRunner(agent); + Session session = newSession(runner); + var unused = runTurn(runner, session, "hi"); + + // Built with no role, the shape a caller assembling a function response naturally produces. + Content roleless = + Content.builder() + .parts( + ImmutableList.of( + Part.builder() + .functionResponse( + FunctionResponse.builder() + .id("lro_call_id") + .name("pendingTool") + .response(ImmutableMap.of("status", "done"))) + .build())) + .build(); + var resumed = + resume(runner, session, /* invocationId= */ null, roleless, /* stateDelta= */ null); + + assertThat(simplifyEvents(resumed)).contains("agent: resumed answer"); + Content lastRequestContent = + Iterables.getLast(Iterables.getLast(testLlm.getRequests()).contents()); + assertThat(lastRequestContent.role()).hasValue("user"); + assertThat(lastRequestContent.parts().get().get(0).functionResponse()).isPresent(); + } + + // The ordinary runAsync overload routes into the resumable flow too, so it needs the same + // role default as the resume overload. + @Test + public void runAsync_resumableApp_messageWithoutRole_reachesTheModel() { + TestLlm testLlm = createTestLlm(createTextLlmResponse("answered")); + LlmAgent agent = createTestAgentBuilder(testLlm).name("agent").build(); + Runner runner = resumableRunner(agent); + Session session = newSession(runner); + Content roleless = + Content.builder().parts(ImmutableList.of(Part.fromText("no role here"))).build(); + + var unused = + runner + .runAsync("user", session.id(), roleless, RunConfig.builder().build()) + .toList() + .blockingGet(); + + Content sent = Iterables.getLast(Iterables.getLast(testLlm.getRequests()).contents()); + assertThat(sent.role()).hasValue("user"); + assertThat(sent.parts().get().get(0).text()).hasValue("no role here"); + } + + // The declared target wins over a same-named agent elsewhere in the tree, as in Python. + @Test + public void resume_replayedTransfer_prefersDeclaredTargetOverSameNamedAgent() { + TestLlm declaredLlm = createTestLlm(createTextLlmResponse("declared ran")); + TestLlm strangerLlm = createTestLlm(createTextLlmResponse("stranger ran")); + // Two agents share the name "helper": one is agent's own sub-agent, one is under a peer. + LlmAgent declaredHelper = createTestAgentBuilder(declaredLlm).name("helper").build(); + LlmAgent agent = + createTestAgentBuilder(createTestLlm(createTextLlmResponse("unused"))) + .name("agent") + .subAgents(declaredHelper) + .build(); + LlmAgent stranger = createTestAgentBuilder(strangerLlm).name("helper").build(); + LlmAgent peer = + createTestAgentBuilder(createTestLlm(createTextLlmResponse("unused"))) + .name("peer") + .subAgents(stranger) + .build(); + LlmAgent root = + createTestAgentBuilder(createTestLlm(createTextLlmResponse("unused"))) + .name("root") + // peer first, so a tree-wide search reaches its "helper" before the declared one. + .subAgents(peer, agent) + .build(); + Runner runner = resumableRunner(root); + Session session = newSession(runner); + String invocationId = "inv-dup-name"; + var unusedUser = + runner + .sessionService() + .appendEvent( + session, + Event.builder() + .id("u1") + .invocationId(invocationId) + .author("user") + .content(Content.fromParts(Part.fromText("go"))) + .build()) + .blockingGet(); + var unusedCall = + runner + .sessionService() + .appendEvent( + session, + Event.builder() + .id("c1") + .invocationId(invocationId) + .author("agent") + .content( + Content.fromParts( + Part.builder() + .functionCall( + FunctionCall.builder() + .id("call_1") + .name("transfer_to_agent") + .args(ImmutableMap.of("agent_name", "helper")) + .build()) + .build())) + .build()) + .blockingGet(); + + var unused = + resume(runner, session, invocationId, /* newMessage= */ null, /* stateDelta= */ null); + + assertThat(declaredLlm.getRequests()).isNotEmpty(); + assertThat(strangerLlm.getRequests()).isEmpty(); + } + + // A message mixing text with a function response resumes on the response; the text is ignored. + @Test + public void resume_functionResponseMixedWithText_resumesIgnoringTheText() { + LlmAgent agent = + createTestAgentBuilder( + createTestLlm( + createFunctionCallLlmResponse( + "lro_call_id", "pendingTool", ImmutableMap.of("message", "hi")), + createTextLlmResponse("done"))) + .name("agent") + .tools(pendingFunctionTool()) + .build(); + Runner runner = resumableRunner(agent); + Session session = newSession(runner); + runTurn(runner, session, "start"); + + Content mixed = + Content.fromParts( + Part.builder() + .functionResponse( + FunctionResponse.builder() + .id("lro_call_id") + .name("pendingTool") + .response(ImmutableMap.of("message", "hi"))) + .build(), + Part.fromText("and also this")); + + ImmutableList resumed = + ImmutableList.copyOf( + runner + .runAsync( + "user", + session.id(), + /* invocationId= */ null, + mixed, + RunConfig.builder().build(), + /* stateDelta= */ null) + .blockingIterable()); + + assertThat(simplifyEvents(resumed)).contains("agent: done"); + } + + // Responses answering calls from two invocations are rejected, not resolved to the newest. + // An id-less function response cannot be routed to a call, so the resume is rejected rather than + // opening a new invocation carrying an orphan response, as in Python. + @Test + public void resume_functionResponseWithoutId_throwsIllegalArgument() { + LlmAgent agent = + createTestAgentBuilder( + createTestLlm( + createFunctionCallLlmResponse( + "lro_call_id", "pendingTool", ImmutableMap.of("message", "hi")), + createTextLlmResponse("done"))) + .name("agent") + .tools(pendingFunctionTool()) + .build(); + Runner runner = resumableRunner(agent); + Session session = newSession(runner); + runTurn(runner, session, "start"); + + Content idLess = + Content.fromParts( + Part.builder() + .functionResponse( + FunctionResponse.builder() + .name("pendingTool") + .response(ImmutableMap.of("message", "hi"))) + .build()); + + runner + .runAsync( + "user", + session.id(), + /* invocationId= */ null, + idLess, + RunConfig.builder().build(), + /* stateDelta= */ null) + .test() + .assertError(IllegalArgumentException.class); + } + + @Test + public void resume_functionResponsesSpanningTwoInvocations_throwsIllegalArgument() { + LlmAgent agent = + createTestAgentBuilder( + createTestLlm( + createFunctionCallLlmResponse( + "call_a", "pendingTool", ImmutableMap.of("message", "a")), + createFunctionCallLlmResponse( + "call_b", "pendingTool", ImmutableMap.of("message", "b")))) + .name("agent") + .tools(pendingFunctionTool()) + .build(); + Runner runner = resumableRunner(agent); + Session session = newSession(runner); + + ImmutableList turn1 = runTurn(runner, session, "start"); + // Plain text starts a second invocation, which pauses on a long-running call of its own. + ImmutableList turn2 = runTurn(runner, session, "another"); + assertThat(turn2.get(0).invocationId()).isNotEqualTo(turn1.get(0).invocationId()); + + Content answersBoth = + Content.fromParts( + Part.builder() + .functionResponse( + FunctionResponse.builder() + .id("call_a") + .name("pendingTool") + .response(ImmutableMap.of("message", "a"))) + .build(), + Part.builder() + .functionResponse( + FunctionResponse.builder() + .id("call_b") + .name("pendingTool") + .response(ImmutableMap.of("message", "b"))) + .build()); + + runner + .runAsync( + "user", + session.id(), + /* invocationId= */ null, + answersBoth, + RunConfig.builder().build(), + /* stateDelta= */ null) + .test() + .assertError(IllegalArgumentException.class); + } + + // Resuming by invocation id alone recovers the invocation's original user message. + @Test + public void resume_withNoNewMessage_restoresOriginalUserContent() { + LlmAgent agent = + createTestAgentBuilder( + createTestLlm( + createFunctionCallLlmResponse( + "lro_call_id", "pendingTool", ImmutableMap.of("message", "hi")), + createTextLlmResponse("done"))) + .name("agent") + .tools(pendingFunctionTool()) + .build(); + UserContentCapturingPlugin plugin = new UserContentCapturingPlugin(); + Runner runner = resumableRunner(agent, plugin); + Session session = newSession(runner); + + ImmutableList pausedTurn = runTurn(runner, session, "the original ask"); + String pausedInvocationId = pausedTurn.get(0).invocationId(); + + resume(runner, session, pausedInvocationId, /* newMessage= */ null, /* stateDelta= */ null); + + assertThat(plugin.lastUserText).isEqualTo("the original ask"); + } + + // A leaf under a SequentialAgent paused on two calls, answered once, must not look completed. + @Test + public void runAsync_resume_pausedUnderSequential_partiallyAnswered_doesNotAdvance() { + Content twoLongRunningCalls = + Content.builder() + .role("model") + .parts( + Part.builder() + .functionCall( + FunctionCall.builder() + .id("c1") + .name("pendingTool") + .args(ImmutableMap.of("message", "a")) + .build()) + .build(), + Part.builder() + .functionCall( + FunctionCall.builder() + .id("c2") + .name("pendingTool") + .args(ImmutableMap.of("message", "b")) + .build()) + .build()) + .build(); + LlmAgent leaf = + createTestAgentBuilder( + createTestLlm( + createLlmResponse(twoLongRunningCalls), createTextLlmResponse("leaf summary"))) + .name("leaf_agent") + .tools(pendingFunctionTool()) + .build(); + LlmAgent next = + createTestAgentBuilder(createTestLlm(createTextLlmResponse("next ran"))) + .name("next_agent") + .build(); + SequentialAgent root = + SequentialAgent.builder().name("seq_agent").subAgents(ImmutableList.of(leaf, next)).build(); + Runner runner = resumableRunner(root); + Session session = newSession(runner); + + runTurn(runner, session, "start"); + + // Answer only c1; c2 remains unanswered. + ImmutableList resumed = + answerCall(runner, session, "c1", "pendingTool", ImmutableMap.of("message", "a")); + + // The sequence holds: neither the leaf's summary nor the following sub-agent runs. + assertThat(simplifyEvents(resumed)).doesNotContain("leaf_agent: leaf summary"); + assertThat(simplifyEvents(resumed)).doesNotContain("next_agent: next ran"); + } + + // Python parity: each ParallelAgent sub-agent runs on its own parallel. branch. + @Test + public void runAsync_resumable_parallelSubAgents_runOnTheirOwnBranches() { + LlmAgent branchA = + createTestAgentBuilder(pauseThenSay("call_a", "branch A resumed")) + .name("branch_a") + .tools(pendingFunctionTool()) + .build(); + LlmAgent branchB = textAgent("branch_b", "branch B done"); + ParallelAgent parallel = + ParallelAgent.builder() + .name("parallel_agent") + .subAgents(ImmutableList.of(branchA, branchB)) + .build(); + Runner runner = resumableRunner(parallel); + Session session = newSession(runner); + + runTurn(runner, session, "start"); + Session reloaded = reloadSession(runner, session); + + assertThat(branchesOf(reloaded, "branch_a")).containsExactly("parallel_agent.branch_a"); + assertThat(branchesOf(reloaded, "branch_b")).containsExactly("parallel_agent.branch_b"); + // The answer inherits the branch of the call it answers, so it lands on branch A alone. + answerCall(runner, session, "call_a", "pendingTool", ImmutableMap.of("message", "hello")); + assertThat(branchesOf(reloadSession(runner, session), "user")) + .contains("parallel_agent.branch_a"); + } + + /** The distinct branch strings carried by {@code author}'s events in {@code session}. */ + private static ImmutableList branchesOf(Session session, String author) { + return session.events().stream() + .filter(event -> author.equals(event.author())) + .map(event -> event.branch().orElse("-")) + .distinct() + .collect(toImmutableList()); + } + + // An answer on a parallel sub-branch the LlmAgent cannot see must still resume the leaf. + @Test + public void runAsync_resume_hitlAnswerOnSubBranch_resumesInsteadOfStalling() { + LlmAgent leaf = + createTestAgentBuilder( + createTestLlm( + createFunctionCallLlmResponse( + "lro_call_id", "pendingTool", ImmutableMap.of("message", "hi")), + createTextLlmResponse("leaf resumed"))) + .name("leaf_agent") + .tools(pendingFunctionTool()) + .build(); + ParallelAgent parallel = + ParallelAgent.builder().name("parallel_agent").subAgents(ImmutableList.of(leaf)).build(); + LlmAgent root = + createTestAgentBuilder( + createTestLlm( + createFunctionCallLlmResponse( + "transfer_call", + "transfer_to_agent", + ImmutableMap.of("agent_name", "parallel_agent")))) + .name("root_agent") + .subAgents(parallel) + .build(); + Runner runner = resumableRunner(root); + Session session = newSession(runner); + + ImmutableList firstTurn = runTurn(runner, session, "start"); + assertThat(simplifyEvents(firstTurn)).isNotEmpty(); + + ImmutableList resumed = + answerCall(runner, session, "lro_call_id", "pendingTool", ImmutableMap.of("message", "hi")); + + // The paused leaf resumes: the answer's branch does not strand the invocation. + assertThat(simplifyEvents(resumed)).contains("leaf_agent: leaf resumed"); + } + + // A transfer inside a workflow: the sequence must advance to the sub-agent after the transferring + // agent once the transferred-to leaf resumes, instead of stranding on the leaf. + @Test + public void runAsync_resume_transferInsideSequential_sequenceAdvancesAfterLeafResumes() { + LlmAgent leaf = + createTestAgentBuilder( + createTestLlm( + createFunctionCallLlmResponse( + "lro_call_id", "pendingTool", ImmutableMap.of("message", "hi")), + createTextLlmResponse("leaf resumed"))) + .name("leaf_agent") + .tools(pendingFunctionTool()) + .build(); + LlmAgent transferring = + createTestAgentBuilder( + createTestLlm( + createFunctionCallLlmResponse( + "transfer_call", + "transfer_to_agent", + ImmutableMap.of("agent_name", "leaf_agent")))) + .name("transferring_agent") + .subAgents(leaf) + .build(); + LlmAgent next = + createTestAgentBuilder(createTestLlm(createTextLlmResponse("next ran"))) + .name("next_agent") + .build(); + SequentialAgent root = + SequentialAgent.builder() + .name("seq_agent") + .subAgents(ImmutableList.of(transferring, next)) + .build(); + Runner runner = resumableRunner(root); + Session session = newSession(runner); + + runTurn(runner, session, "start"); + + ImmutableList resumed = + answerCall(runner, session, "lro_call_id", "pendingTool", ImmutableMap.of("message", "hi")); + + assertThat(simplifyEvents(resumed)).contains("leaf_agent: leaf resumed"); + // The enclosing sequence must not stall on the leaf. + assertThat(simplifyEvents(resumed)).contains("next_agent: next ran"); + } + + /** Records the user content each run is started with, so a resumed run can be checked. */ + private static final class UserContentCapturingPlugin extends BasePlugin { + private @Nullable String lastUserText; + + UserContentCapturingPlugin() { + super("user-content-capturing"); + } + + @Override + public Maybe beforeRunCallback(InvocationContext invocationContext) { + lastUserText = + invocationContext.userContent().flatMap(Content::parts).stream() + .flatMap(List::stream) + .map(part -> part.text().orElse("")) + .findFirst() + .orElse(null); + return Maybe.empty(); + } + } + + // Resuming a non-existent invocation with no message throws instead of calling the model. + @Test + public void resume_nonExistentInvocationId_throwsIllegalArgument() { + LlmAgent agent = + createTestAgentBuilder(createTestLlm(createTextLlmResponse("x"))).name("agent").build(); + Runner runner = resumableRunner(agent); + Session session = newSession(runner); + + runner + .runAsync( + "user", + session.id(), + "does-not-exist", + /* newMessage= */ null, + RunConfig.builder().build(), + /* stateDelta= */ null) + .test() + .assertError(IllegalArgumentException.class); + } + + // An explicit invocation id skips message resolution, so the orphan check must fire here too. + @Test + public void resume_orphanFunctionResponseWithProvidedInvocationId_throwsIllegalArgument() { + LlmAgent agent = + createTestAgentBuilder(createTestLlm(createTextLlmResponse("x"))).name("agent").build(); + Runner runner = resumableRunner(agent); + Session session = newSession(runner); + + Content orphanResponse = + Content.fromParts( + Part.builder() + .functionResponse( + FunctionResponse.builder() + .id("no_such_call") + .name("pendingTool") + .response(ImmutableMap.of("status", "done"))) + .build()); + + runner + .runAsync( + "user", + session.id(), + "some-inv", + orphanResponse, + RunConfig.builder().build(), + /* stateDelta= */ null) + .test() + .assertError(IllegalArgumentException.class); + } + + // InMemoryRunnerTest parity: the appended function response inherits the call's branch. + @Test + public void resume_withFunctionResponse_copiesBranchFromMatchingCall() { + TestBaseAgent agent = + new TestBaseAgent("test_agent", "desc", () -> Flowable.empty(), null, null, null); + Runner runner = resumableRunner(agent); + Session session = newSession(runner); + Object unusedFc = + runner + .sessionService() + .appendEvent( + session, + Event.builder() + .id("fc1") + .invocationId("test-inv") + .author("test_agent") + .branch("my_special_branch") + .content( + Content.fromParts( + Part.builder() + .functionCall( + FunctionCall.builder().id("call_abc").name("test_func").build()) + .build())) + .build()) + .blockingGet(); + + Object unused = + runner + .runAsync( + "user", + session.id(), + /* invocationId= */ null, + Content.fromParts( + Part.builder() + .functionResponse( + FunctionResponse.builder() + .id("call_abc") + .name("test_func") + .response(ImmutableMap.of("result", "ok"))) + .build()), + RunConfig.builder().build(), + /* stateDelta= */ null) + .toList() + .blockingGet(); + + Session reloaded = + runner + .sessionService() + .getSession("test", "user", session.id(), Optional.empty()) + .blockingGet(); + Event lastUser = + Streams.findLast(reloaded.events().stream().filter(event -> event.author().equals("user"))) + .get(); + assertThat(lastUser.branch()).hasValue("my_special_branch"); + } + + // A pending long-running call stops a resumable LoopAgent after the iteration, as Python ADK v1. + @Test + public void runAsync_loopAgentWithLongRunningSubAgent_resumable_stopsAfterFirstIteration() { + AtomicInteger calls = new AtomicInteger(); + TestLlm loopLlm = + createTestLlm( + () -> + calls.incrementAndGet() <= 5 + ? Flowable.just( + createFunctionCallLlmResponse( + "lro_call_id", "pendingTool", ImmutableMap.of("message", "hello"))) + : Flowable.just(createTextLlmResponse("stop"))); + LlmAgent inner = + createTestAgentBuilder(loopLlm).name("inner").tools(pendingFunctionTool()).build(); + LoopAgent loop = + LoopAgent.builder() + .name("loop") + .subAgents(ImmutableList.of(inner)) + .maxIterations(3) + .build(); + Runner runner = resumableRunner(loop); + Session session = newSession(runner); + + ImmutableList unused = runTurn(runner, session, "from user"); + + // Paused after the first iteration: one model call, not maxIterations. + assertThat(loopLlm.getRequests()).hasSize(1); + } + + // A pending long-running call pauses only its own ParallelAgent branch, as Python ADK v1 does. + @Test + public void runAsync_parallelAgentWithLongRunningBranch_resumable_otherBranchCompletes() { + TestLlm longRunningLlm = + createTestLlm( + createFunctionCallLlmResponse( + "lro_call_id", "pendingTool", ImmutableMap.of("message", "hello")), + createTextLlmResponse("unexpected")); + LlmAgent longRunningBranch = + createTestAgentBuilder(longRunningLlm) + .name("long_running_branch") + .tools(pendingFunctionTool()) + .build(); + LlmAgent plainBranch = + createTestAgentBuilder(createTestLlm(createTextLlmResponse("plain branch done"))) + .name("plain_branch") + .build(); + ParallelAgent parallel = + ParallelAgent.builder() + .name("parallel") + .subAgents(ImmutableList.of(longRunningBranch, plainBranch)) + .build(); + Runner runner = resumableRunner(parallel); + Session session = newSession(runner); + + ImmutableList events = runTurn(runner, session, "from user"); + + // The long-running branch paused after one model call; the other branch still completed. + assertThat(longRunningLlm.getRequests()).hasSize(1); + assertThat(simplifyEvents(events)).contains("plain_branch: plain branch done"); + } +} diff --git a/core/src/test/java/com/google/adk/sessions/InMemorySessionServiceTest.java b/core/src/test/java/com/google/adk/sessions/InMemorySessionServiceTest.java index 58c445641..2f3a5bcae 100644 --- a/core/src/test/java/com/google/adk/sessions/InMemorySessionServiceTest.java +++ b/core/src/test/java/com/google/adk/sessions/InMemorySessionServiceTest.java @@ -16,13 +16,16 @@ package com.google.adk.sessions; import static com.google.common.truth.Truth.assertThat; +import static org.junit.Assert.assertThrows; import com.google.adk.events.Event; import com.google.adk.events.EventActions; +import com.google.adk.sessions.InMemorySessionService.DuplicateSessionIdBehavior; import io.reactivex.rxjava3.core.Single; import java.lang.reflect.Field; import java.time.Instant; import java.util.HashMap; +import java.util.Map; import java.util.Optional; import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.ConcurrentMap; @@ -340,4 +343,88 @@ public void getSession_numRecentEventsAndAfterTimestamp_appliesBothFilters() { .containsExactly(300L, 400L, 500L) .inOrder(); } + + @Test + public void createSession_duplicateSessionId_byDefaultReplacesExistingSession() { + InMemorySessionService sessionService = new InMemorySessionService(); + Map state = new HashMap<>(); + state.put("sessionKey", "sessionValue"); + Session session = sessionService.createSession("app", "user", state, "session1").blockingGet(); + var unused = + sessionService.appendEvent(session, Event.builder().timestamp(100).build()).blockingGet(); + + Session replacement = + sessionService.createSession("app", "user", new HashMap<>(), "session1").blockingGet(); + + assertThat(replacement.events()).isEmpty(); + assertThat(replacement.state()).doesNotContainKey("sessionKey"); + Session retrieved = + sessionService.getSession("app", "user", "session1", Optional.empty()).blockingGet(); + assertThat(retrieved.events()).isEmpty(); + assertThat(retrieved.state()).doesNotContainKey("sessionKey"); + } + + @Test + public void createSession_duplicateSessionId_whenRejecting_errorsAndKeepsExistingSession() { + InMemorySessionService sessionService = + new InMemorySessionService(DuplicateSessionIdBehavior.REJECT); + Map state = new HashMap<>(); + state.put("sessionKey", "sessionValue"); + Session session = sessionService.createSession("app", "user", state, "session1").blockingGet(); + var unused = + sessionService.appendEvent(session, Event.builder().timestamp(100).build()).blockingGet(); + + Single duplicate = + sessionService.createSession("app", "user", new HashMap<>(), "session1"); + + SessionException exception = assertThrows(SessionException.class, duplicate::blockingGet); + assertThat(exception).hasMessageThat().isEqualTo(SessionException.SESSION_ALREADY_EXISTS); + Session retrieved = + sessionService.getSession("app", "user", "session1", Optional.empty()).blockingGet(); + assertThat(retrieved.events().stream().map(Event::timestamp)).containsExactly(100L); + assertThat(retrieved.state()).containsEntry("sessionKey", "sessionValue"); + } + + @Test + public void createSession_distinctSessionIds_whenRejecting_createsBoth() { + InMemorySessionService sessionService = + new InMemorySessionService(DuplicateSessionIdBehavior.REJECT); + + Session first = + sessionService.createSession("app", "user", new HashMap<>(), "session1").blockingGet(); + Session second = + sessionService.createSession("app", "user", new HashMap<>(), "session2").blockingGet(); + + assertThat(first.id()).isEqualTo("session1"); + assertThat(second.id()).isEqualTo("session2"); + assertThat(sessionService.listSessions("app", "user").blockingGet().sessions()).hasSize(2); + } + + @Test + public void createSession_afterDelete_whenRejecting_reusesTheSameId() { + InMemorySessionService sessionService = + new InMemorySessionService(DuplicateSessionIdBehavior.REJECT); + var unused = + sessionService.createSession("app", "user", new HashMap<>(), "session1").blockingGet(); + sessionService.deleteSession("app", "user", "session1").blockingAwait(); + + Session recreated = + sessionService.createSession("app", "user", new HashMap<>(), "session1").blockingGet(); + + assertThat(recreated.id()).isEqualTo("session1"); + } + + @Test + public void createSession_sameSessionIdDifferentUser_whenRejecting_isAllowed() { + InMemorySessionService sessionService = + new InMemorySessionService(DuplicateSessionIdBehavior.REJECT); + var unused = + sessionService.createSession("app", "user-a", new HashMap<>(), "session1").blockingGet(); + + Session other = + sessionService.createSession("app", "user-b", new HashMap<>(), "session1").blockingGet(); + + assertThat(other.id()).isEqualTo("session1"); + assertThat(other.userId()).isEqualTo("user-b"); + } } diff --git a/core/src/test/java/com/google/adk/sessions/SessionJsonConverterTest.java b/core/src/test/java/com/google/adk/sessions/SessionJsonConverterTest.java index a63e3b38d..d52ef93da 100644 --- a/core/src/test/java/com/google/adk/sessions/SessionJsonConverterTest.java +++ b/core/src/test/java/com/google/adk/sessions/SessionJsonConverterTest.java @@ -276,6 +276,107 @@ public void fromApiEvent_complexActions_success() { assertThat(eventActions.endOfAgent()).isTrue(); } + @Test + public void convertEventToJson_agentState_success() throws JsonProcessingException { + EventActions actions = + EventActions.builder() + .agentState(ImmutableMap.of("current_sub_agent", "b_agent", "times_looped", 2)) + .build(); + Event event = + Event.builder() + .author("agent") + .invocationId("inv-1") + .timestamp(Instant.parse("2023-01-01T00:00:00.123Z").toEpochMilli()) + .actions(actions) + .build(); + + String json = SessionJsonConverter.convertEventToJson(event, true); + JsonNode actionsNode = objectMapper.readTree(json).get("actions"); + + assertThat(actionsNode.get("agentState").get("current_sub_agent").asText()) + .isEqualTo("b_agent"); + assertThat(actionsNode.get("agentState").get("times_looped").asInt()).isEqualTo(2); + } + + // rawEvent supplies only the keys "actions" lacks: a rawEvent carrying endOfAgent alone must not + // shadow an agentState that actions does carry. + @Test + public void fromApiEvent_rawEventMissingAgentState_doesNotShadowActions() { + Map apiEvent = new HashMap<>(); + apiEvent.put("name", "sessions/123/events/456"); + apiEvent.put("invocationId", "inv-1"); + apiEvent.put("author", "agent"); + apiEvent.put("timestamp", "2023-01-01T00:00:00.123Z"); + Map actions = new HashMap<>(); + actions.put("agentState", ImmutableMap.of("current_sub_agent", "b_agent")); + apiEvent.put("actions", actions); + apiEvent.put("rawEvent", ImmutableMap.of("actions", ImmutableMap.of("endOfAgent", true))); + + Event event = SessionJsonConverter.fromApiEvent(apiEvent); + + assertThat(event.actions().agentState()).isPresent(); + assertThat(event.actions().agentState().get()).containsEntry("current_sub_agent", "b_agent"); + assertThat(event.actions().endOfAgent()).isTrue(); + } + + // Python's Vertex session service keeps agentState and endOfAgent only under rawEvent, so a + // session it wrote must still load with its checkpoints. + @Test + public void fromApiEvent_agentStateOnlyUnderRawEvent_readsCheckpoint() { + Map apiEvent = new HashMap<>(); + apiEvent.put("name", "sessions/123/events/456"); + apiEvent.put("invocationId", "inv-1"); + apiEvent.put("author", "agent"); + apiEvent.put("timestamp", "2023-01-01T00:00:00.123Z"); + // The six-key actions block Python writes, with no agentState and no endOfAgent. + apiEvent.put("actions", new HashMap<>(ImmutableMap.of("stateDelta", new HashMap<>()))); + Map rawActions = new HashMap<>(); + rawActions.put("agentState", ImmutableMap.of("current_sub_agent", "b_agent")); + rawActions.put("endOfAgent", true); + apiEvent.put("rawEvent", ImmutableMap.of("actions", rawActions)); + + Event event = SessionJsonConverter.fromApiEvent(apiEvent); + + assertThat(event.actions().agentState()).isPresent(); + assertThat(event.actions().agentState().get()).containsEntry("current_sub_agent", "b_agent"); + assertThat(event.actions().endOfAgent()).isTrue(); + } + + @Test + public void fromApiEvent_agentState_success() { + Map apiEvent = new HashMap<>(); + apiEvent.put("name", "sessions/123/events/456"); + apiEvent.put("invocationId", "inv-1"); + apiEvent.put("author", "agent"); + apiEvent.put("timestamp", "2023-01-01T00:00:00.123Z"); + Map actions = new HashMap<>(); + actions.put("agentState", ImmutableMap.of("current_sub_agent", "b_agent", "times_looped", 2)); + apiEvent.put("actions", actions); + + Event event = SessionJsonConverter.fromApiEvent(apiEvent); + + assertThat(event.actions().agentState()).isPresent(); + assertThat(event.actions().agentState().get()).containsEntry("current_sub_agent", "b_agent"); + assertThat(event.actions().agentState().get()).containsEntry("times_looped", 2); + } + + @Test + public void fromApiEvent_agentStateNotAMap_isIgnored() { + // Another runtime may write a non-map agentState; dropping it keeps the session loadable. + Map apiEvent = new HashMap<>(); + apiEvent.put("name", "sessions/123/events/456"); + apiEvent.put("invocationId", "inv-1"); + apiEvent.put("author", "agent"); + apiEvent.put("timestamp", "2023-01-01T00:00:00.123Z"); + Map actions = new HashMap<>(); + actions.put("agentState", "not-a-map"); + apiEvent.put("actions", actions); + + Event event = SessionJsonConverter.fromApiEvent(apiEvent); + + assertThat(event.actions().agentState()).isEmpty(); + } + @Test public void fromApiEvent_minimalEvent_success() { Map apiEvent = new HashMap<>(); diff --git a/core/src/test/java/com/google/adk/testing/ResumabilityTestUtils.java b/core/src/test/java/com/google/adk/testing/ResumabilityTestUtils.java new file mode 100644 index 000000000..840c6ea9e --- /dev/null +++ b/core/src/test/java/com/google/adk/testing/ResumabilityTestUtils.java @@ -0,0 +1,294 @@ +/* + * Copyright 2026 Google LLC + * + * Licensed 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 com.google.adk.testing; + +import static com.google.adk.testing.TestUtils.createTestAgentBuilder; +import static com.google.adk.testing.TestUtils.createTestLlm; +import static com.google.adk.testing.TestUtils.createTextLlmResponse; +import static com.google.adk.testing.TestUtils.simplifyResumableEvents; +import static com.google.common.collect.ImmutableList.toImmutableList; +import static com.google.common.truth.Truth.assertThat; +import static com.google.common.truth.Truth.assertWithMessage; +import static java.util.Arrays.stream; + +import com.google.adk.agents.BaseAgent; +import com.google.adk.agents.LlmAgent; +import com.google.adk.agents.RunConfig; +import com.google.adk.apps.App; +import com.google.adk.apps.ResumabilityConfig; +import com.google.adk.events.Event; +import com.google.adk.models.LlmResponse; +import com.google.adk.plugins.BasePlugin; +import com.google.adk.runner.Runner; +import com.google.adk.sessions.Session; +import com.google.adk.tools.FunctionTool; +import com.google.common.collect.ImmutableList; +import com.google.common.collect.ImmutableMap; +import com.google.errorprone.annotations.CanIgnoreReturnValue; +import com.google.genai.types.Content; +import com.google.genai.types.FunctionCall; +import com.google.genai.types.FunctionResponse; +import com.google.genai.types.Part; +import io.reactivex.rxjava3.core.Flowable; +import java.util.List; +import java.util.Map; +import java.util.Objects; +import java.util.Optional; +import org.jspecify.annotations.Nullable; + +/** + * Helpers shared by the resumability tests: building a runner in each resumption mode, driving + * turns and resumes through it, and asserting on the checkpoints it emits. + * + *

Separate from {@link TestUtils} because these reach the {@link Runner} and {@link App} layer, + * which the agent-level helpers there do not. + */ +public final class ResumabilityTestUtils { + + /** Tools the resumability tests pause on. */ + public static final class Tools { + private Tools() {} + + /** Returns its argument, so a call to it completes within the turn. */ + public static ImmutableMap echoTool(String message) { + return ImmutableMap.of("message", message); + } + + /** Awaits a result supplied later, so a call to it leaves the invocation paused. */ + public static @Nullable ImmutableMap pendingTool(String message) { + return null; + } + } + + /** {@code pendingTool} wired long-running: the shape that pauses an invocation. */ + public static FunctionTool pendingFunctionTool() { + return FunctionTool.create( + Tools.class, "pendingTool", /* requireConfirmation= */ false, /* isLongRunning= */ true); + } + + /** {@code echoTool} wired long-running but returning in-turn, so it does not pause. */ + public static FunctionTool longRunningEchoFunctionTool() { + return FunctionTool.create( + Tools.class, "echoTool", /* requireConfirmation= */ false, /* isLongRunning= */ true); + } + + /** {@code echoTool} requiring confirmation: the human-in-the-loop shape. */ + public static FunctionTool confirmingEchoFunctionTool() { + return FunctionTool.create(Tools.class, "echoTool", /* requireConfirmation= */ true); + } + + /** An {@link LlmAgent} named {@code name} whose model returns {@code texts}, one per call. */ + public static LlmAgent textAgent(String name, String... texts) { + return createTestAgentBuilder( + createTestLlm( + stream(texts).map(TestUtils::createTextLlmResponse).toArray(LlmResponse[]::new))) + .name(name) + .build(); + } + + /** An {@link LlmAgent} named {@code name} carrying the long-running {@code pendingTool}. */ + public static LlmAgent pausingAgent(String name, TestLlm llm) { + return createTestAgentBuilder(llm).name(name).tools(pendingFunctionTool()).build(); + } + + /** A model script that calls {@code pendingTool} once, then says {@code then}. */ + public static TestLlm pauseThenSay(String callId, String then) { + return createTestLlm( + TestUtils.createFunctionCallLlmResponse( + callId, "pendingTool", ImmutableMap.of("message", "hello")), + createTextLlmResponse(then)); + } + + /** A {@link Runner} over {@code rootAgent} with durable resumability enabled. */ + @SuppressWarnings("deprecation") // ResumabilityConfig is @Experimental, not deprecated, here. + public static Runner resumableRunner(BaseAgent rootAgent, BasePlugin... plugins) { + return Runner.builder() + .app( + App.builder() + .name("test") + .rootAgent(rootAgent) + .plugins(ImmutableList.copyOf(plugins)) + .resumabilityConfig(ResumabilityConfig.builder().resumable(true).build()) + .build()) + .build(); + } + + /** A {@link Runner} over {@code rootAgent} with only the deprecated plain-text shim enabled. */ + @SuppressWarnings("deprecation") // The plain-text continuation shim is deprecated by design. + public static Runner shimRunner(BaseAgent rootAgent) { + return Runner.builder() + .app( + App.builder() + .name("test") + .rootAgent(rootAgent) + .resumabilityConfig( + ResumabilityConfig.builder().plainTextContinuationAutoResume(true).build()) + .build()) + .build(); + } + + /** A fresh session on {@code runner}'s session service. */ + public static Session newSession(Runner runner) { + return runner.sessionService().createSession("test", "user").blockingGet(); + } + + /** Runs one new plain-text turn and returns its events. */ + @CanIgnoreReturnValue + public static ImmutableList runTurn(Runner runner, Session session, String text) { + return ImmutableList.copyOf( + runner + .runAsync("user", session.id(), Content.fromParts(Part.fromText(text))) + .toList() + .blockingGet()); + } + + /** Content answering a pending function call, as a user turn carries it. */ + public static Content functionResponseContent( + String callId, String toolName, Map response) { + return Content.fromParts( + Part.builder() + .functionResponse( + FunctionResponse.builder().id(callId).name(toolName).response(response)) + .build()); + } + + /** Resumes by answering the pending call {@code callId}; the runner resolves the invocation. */ + @CanIgnoreReturnValue + public static ImmutableList answerCall( + Runner runner, + Session session, + String callId, + String toolName, + Map response) { + return ImmutableList.copyOf( + runner + .runAsync("user", session.id(), functionResponseContent(callId, toolName, response)) + .toList() + .blockingGet()); + } + + /** The un-subscribed resume stream, for tests asserting on the error rather than the events. */ + public static Flowable resumeFlowable( + Runner runner, + Session session, + @Nullable String invocationId, + @Nullable Content newMessage, + @Nullable Map stateDelta) { + return runner.runAsync( + "user", session.id(), invocationId, newMessage, RunConfig.builder().build(), stateDelta); + } + + /** The six-arg resume overload, for tests that vary the invocation id or the state delta. */ + @CanIgnoreReturnValue + public static ImmutableList resume( + Runner runner, + Session session, + @Nullable String invocationId, + @Nullable Content newMessage, + @Nullable Map stateDelta) { + return ImmutableList.copyOf( + resumeFlowable(runner, session, invocationId, newMessage, stateDelta) + .toList() + .blockingGet()); + } + + /** Resumes {@code invocationId} with no new message. */ + @CanIgnoreReturnValue + public static ImmutableList resumeById( + Runner runner, Session session, @Nullable String invocationId) { + return resume(runner, session, invocationId, /* newMessage= */ null, /* stateDelta= */ null); + } + + /** Reloads {@code session} from the runner's session service, to assert on persisted state. */ + public static Session reloadSession(Runner runner, Session session) { + return runner + .sessionService() + .getSession("test", "user", session.id(), Optional.empty()) + .blockingGet(); + } + + /** + * A model turn issuing two long-running calls at once: the shape that pauses until both answer. + */ + public static Content twoPendingCalls(String firstId, String secondId) { + return Content.builder() + .role("model") + .parts(pendingCallPart(firstId, "a"), pendingCallPart(secondId, "b")) + .build(); + } + + private static Part pendingCallPart(String callId, String message) { + return Part.builder() + .functionCall( + FunctionCall.builder() + .id(callId) + .name("pendingTool") + .args(ImmutableMap.of("message", message)) + .build()) + .build(); + } + + /** Asserts {@code author} emitted an end-of-agent checkpoint. */ + public static void assertEndOfAgent(List events, String author) { + assertWithMessage( + "end-of-agent checkpoint for %s in %s", author, simplifyResumableEvents(events)) + .that( + events.stream() + .anyMatch(event -> Objects.equals(event.author(), author) && endsAgent(event))) + .isTrue(); + } + + /** Asserts {@code author} emitted no end-of-agent checkpoint. */ + public static void assertNoEndOfAgent(List events, String author) { + assertWithMessage( + "unexpected end-of-agent checkpoint for %s in %s", + author, simplifyResumableEvents(events)) + .that( + events.stream() + .anyMatch(event -> Objects.equals(event.author(), author) && endsAgent(event))) + .isFalse(); + } + + /** Asserts whether {@code author} emitted an agent-state checkpoint. */ + public static void assertAgentStateCheckpoint( + List events, String author, boolean expected) { + assertWithMessage( + "agent-state checkpoint for %s in %s", author, simplifyResumableEvents(events)) + .that( + events.stream() + .anyMatch( + event -> + Objects.equals(event.author(), author) + && event.actions().agentState().isPresent())) + .isEqualTo(expected); + } + + /** Asserts {@code session} holds no resumability checkpoints at all. */ + public static void assertNoCheckpoints(Session session) { + assertThat( + session.events().stream() + .filter(event -> endsAgent(event) || event.actions().agentState().isPresent()) + .collect(toImmutableList())) + .isEmpty(); + } + + private static boolean endsAgent(Event event) { + return event.actions().endOfAgent(); + } + + private ResumabilityTestUtils() {} +} diff --git a/core/src/test/java/com/google/adk/testing/TestUtils.java b/core/src/test/java/com/google/adk/testing/TestUtils.java index daed8d2e4..577241e78 100644 --- a/core/src/test/java/com/google/adk/testing/TestUtils.java +++ b/core/src/test/java/com/google/adk/testing/TestUtils.java @@ -24,6 +24,7 @@ import com.google.adk.agents.InvocationContext; import com.google.adk.agents.LlmAgent; import com.google.adk.agents.RunConfig; +import com.google.adk.apps.ResumabilityConfig; import com.google.adk.artifacts.InMemoryArtifactService; import com.google.adk.events.Event; import com.google.adk.events.EventActions; @@ -71,6 +72,21 @@ public static InvocationContext createInvocationContext(BaseAgent agent) { return createInvocationContext(agent, RunConfig.builder().build()); } + /** Like {@link #createInvocationContext(BaseAgent)} but with resumability enabled. */ + public static InvocationContext createResumableInvocationContext(BaseAgent agent) { + InMemorySessionService sessionService = new InMemorySessionService(); + return InvocationContext.builder() + .sessionService(sessionService) + .artifactService(new InMemoryArtifactService()) + .invocationId("invocationId") + .agent(agent) + .session(sessionService.createSession("test_app", "test-user").blockingGet()) + .userContent(Content.fromParts(Part.fromText("user content"))) + .runConfig(RunConfig.builder().build()) + .resumabilityConfig(ResumabilityConfig.builder().resumable(true).build()) + .build(); + } + public static InvocationContext createInvocationContext( BaseAgent agent, BaseSessionService sessionService, Session session) { return InvocationContext.builder() @@ -105,6 +121,30 @@ public static ImmutableList simplifyEvents(List events) { .collect(toImmutableList()); } + /** Marker rendered for an end-of-agent checkpoint event by {@link #simplifyResumableEvents}. */ + public static final String END_OF_AGENT = "end_of_agent"; + + /** + * Like {@link #simplifyEvents} but renders resumability checkpoint events distinctly: an + * end-of-agent event as {@link #END_OF_AGENT} and an agent-state event as {@code + * agent_state=...}. + */ + public static ImmutableList simplifyResumableEvents(List events) { + return events.stream() + .map(event -> event.author() + ": " + formatResumableEvent(event)) + .collect(toImmutableList()); + } + + private static String formatResumableEvent(Event event) { + if (event.actions().endOfAgent()) { + return END_OF_AGENT; + } + if (event.actions().agentState().isPresent()) { + return "agent_state=" + event.actions().agentState().get(); + } + return formatEventContent(event); + } + private static String formatEventContent(Event event) { return formatContent( event @@ -208,12 +248,10 @@ public static TestBaseAgent createSubAgent(String name, Supplier return new TestBaseAgent(name, eventSupplier, /* subAgents= */ ImmutableList.of()); } - // TODO: b/414071046 Deprecate. public static LlmAgent createTestAgent(BaseLlm llm) { return createTestAgentBuilder(llm).build(); } - // TODO: b/414071046 Make this return TestAgent. It can be used with toBuilder(). public static LlmAgent.Builder createTestAgentBuilder(BaseLlm llm) { return LlmAgent.builder().name("test agent").description("test agent description").model(llm); } diff --git a/core/src/test/java/com/google/adk/tools/BaseToolTest.java b/core/src/test/java/com/google/adk/tools/BaseToolTest.java index 960e8aab3..16a842cd5 100644 --- a/core/src/test/java/com/google/adk/tools/BaseToolTest.java +++ b/core/src/test/java/com/google/adk/tools/BaseToolTest.java @@ -45,7 +45,6 @@ import org.junit.runner.RunWith; import org.junit.runners.JUnit4; -// TODO(b/410859954): Cover more of the behavior of the default processLlmRequest @RunWith(JUnit4.class) public final class BaseToolTest { @@ -183,7 +182,7 @@ public void processLlmRequestWithGoogleSearchToolAddsToolToConfig() { } @Test - public void processLlmRequestWithLatestAliasAddsToolToConfig() { + public void processLlmRequestWithGoogleSearchToolAndLatestAliasAddsToolToConfig() { final GoogleSearchTool googleSearchTool = new GoogleSearchTool(); LlmRequest.Builder builder = LlmRequest.builder().model("gemini-flash-latest").build().toBuilder(); @@ -194,19 +193,23 @@ public void processLlmRequestWithLatestAliasAddsToolToConfig() { } @Test - public void processLlmRequestWithUnsupportedModelReturnsError() { - final GoogleSearchTool googleSearchTool = new GoogleSearchTool(); - LlmRequest.Builder builder = LlmRequest.builder().model("text-bison-001").build().toBuilder(); - Completable result = googleSearchTool.processLlmRequest(builder, null); - result.test().assertError(IllegalArgumentException.class); + public void processLlmRequestWithGoogleSearchToolAndNonGeminiModelAddsToolToConfig() { + GoogleSearchTool googleSearchTool = new GoogleSearchTool(); + LlmRequest.Builder builder = LlmRequest.builder().model("text-bison-001"); + Completable result = googleSearchTool.processLlmRequest(builder, /* toolContext= */ null); + result.test().assertComplete(); + assertThat(builder.build().config().get().tools().get()) + .containsExactly(Tool.builder().googleSearch(GoogleSearch.builder().build()).build()); } @Test - public void processLlmRequest_WithNullModel_ReturnsError() { - final GoogleSearchTool googleSearchTool = new GoogleSearchTool(); - LlmRequest.Builder builder = LlmRequest.builder().build().toBuilder(); - Completable result = googleSearchTool.processLlmRequest(builder, null); - result.test().assertError(IllegalArgumentException.class); + public void processLlmRequestWithGoogleSearchToolAndNoModelAddsToolToConfig() { + GoogleSearchTool googleSearchTool = new GoogleSearchTool(); + LlmRequest.Builder builder = LlmRequest.builder(); + Completable result = googleSearchTool.processLlmRequest(builder, /* toolContext= */ null); + result.test().assertComplete(); + assertThat(builder.build().config().get().tools().get()) + .containsExactly(Tool.builder().googleSearch(GoogleSearch.builder().build()).build()); } @Test @@ -238,6 +241,48 @@ public void processLlmRequestWithUrlContextToolAddsToolToConfig() { Tool.builder().urlContext(UrlContext.builder().build()).build()); } + @Test + public void processLlmRequestWithUrlContextToolAndLatestAliasAddsToolToConfig() { + UrlContextTool urlContextTool = new UrlContextTool(); + LlmRequest.Builder builder = LlmRequest.builder().model("gemini-flash-latest"); + Completable result = urlContextTool.processLlmRequest(builder, /* toolContext= */ null); + result.test().assertComplete(); + assertThat(builder.build().config().get().tools().get()) + .containsExactly(Tool.builder().urlContext(UrlContext.builder().build()).build()); + } + + @Test + public void processLlmRequestWithUrlContextToolAndResourcePathModelAddsToolToConfig() { + UrlContextTool urlContextTool = new UrlContextTool(); + LlmRequest.Builder builder = + LlmRequest.builder() + .model("projects/p/locations/us-central1/publishers/google/models/gemini-2.5-flash"); + Completable result = urlContextTool.processLlmRequest(builder, /* toolContext= */ null); + result.test().assertComplete(); + assertThat(builder.build().config().get().tools().get()) + .containsExactly(Tool.builder().urlContext(UrlContext.builder().build()).build()); + } + + @Test + public void processLlmRequestWithUrlContextToolAndNonGeminiModelAddsToolToConfig() { + UrlContextTool urlContextTool = new UrlContextTool(); + LlmRequest.Builder builder = LlmRequest.builder().model("text-bison-001"); + Completable result = urlContextTool.processLlmRequest(builder, /* toolContext= */ null); + result.test().assertComplete(); + assertThat(builder.build().config().get().tools().get()) + .containsExactly(Tool.builder().urlContext(UrlContext.builder().build()).build()); + } + + @Test + public void processLlmRequestWithUrlContextToolAndNoModelAddsToolToConfig() { + UrlContextTool urlContextTool = new UrlContextTool(); + LlmRequest.Builder builder = LlmRequest.builder(); + Completable result = urlContextTool.processLlmRequest(builder, /* toolContext= */ null); + result.test().assertComplete(); + assertThat(builder.build().config().get().tools().get()) + .containsExactly(Tool.builder().urlContext(UrlContext.builder().build()).build()); + } + private static InvocationContext.Builder testInvocationContext() { InvocationContext.Builder builder = InvocationContext.builder(); builder.agent(testAgent().build()); @@ -399,16 +444,6 @@ public void runAsync_withClassAndObjectMapper_convertsArguments() throws Excepti testObserver.assertValue(expected); } - @Test - public void testProcessLlmRequest_WithNoModel_DoesNotThrowsException() { - GoogleSearchTool tool = GoogleSearchTool.INSTANCE; - LlmRequest.Builder requestBuilder = LlmRequest.builder(); - - tool.processLlmRequest(requestBuilder, null); - - assertNotNull(requestBuilder); - } - public record TestToolArgs(int i, String s) {} @Test diff --git a/dev/INTENRAL_TODOS.md b/dev/INTENRAL_TODOS.md deleted file mode 100644 index ec18f572c..000000000 --- a/dev/INTENRAL_TODOS.md +++ /dev/null @@ -1,22 +0,0 @@ -# Dev TODOs - -This file contains TODOs for the dev ADK module based on -[Recommendations for making ADK Java more idiomatic](http://go/idiomatic-adk-java). - -## Dev UI - -- [ ] **Conditional UI**: Add a configuration property (e.g., - `adk.web.ui.enabled`) to conditionally enable/disable serving Dev UI static - assets (in `AdkWebServer`). -- [ ] **Integration Tests**: Add E2E tests (Selenium/Playwright/HtmlUnit) for - Dev UI to verify interaction between frontend assets and Spring Boot - backend. -- [ ] **Integration Tests**: Test critical paths like loading UI, WebSocket - connection, sending/receiving messages, and rich content handling (images). - -## Production Readiness - -- [ ] **Actuators**: Enable and configure Spring Boot Actuator endpoints for - monitoring and management. -- [ ] **Actuators**: Configure startup and readiness probes for production - environments. diff --git a/dev/README.md b/dev/README.md index 3f4e9a3d3..97c49ea36 100644 --- a/dev/README.md +++ b/dev/README.md @@ -6,3 +6,36 @@ The UI and its assets are served under `/dev-ui/`, and both `/` and `/dev-ui` redirect there, keeping the query string. The assets are not served from the origin root: `/adk_favicon.svg` and the like return 404, and only the `/dev-ui/` form resolves. + +## Behind a reverse proxy + +If the server runs behind a reverse proxy that strips a path prefix (for +example, forwarding `https://gateway.example.com/my-app/` to `/`), set +`adk.web.backend-url` to the external URL that browsers use to reach the server: + +```properties +adk.web.backend-url=https://gateway.example.com/my-app +``` + +Setting a usable `adk.web.backend-url`: + +- Sets `backendUrl` in `/dev-ui/assets/config/runtime-config.json` so the UI + sends API and WebSocket requests through the proxy URL. +- Prepends the URL's path (`/my-app`) to the `/` and `/dev-ui` entry redirects + (`/my-app/dev-ui/`). + +The value must be an absolute `http://` or `https://` URL without credentials, +query parameters, or a fragment. A value that does not meet those constraints is +ignored with a warning: it is neither served nor used as a redirect prefix, so +the bundled `backendUrl` stands exactly as when the property is unset. The UI +parses values without an `http://` or `https://` scheme as a WebSocket host, so +a relative path like `/my-app` cannot be used. If the application also sets +`server.servlet.context-path`, include it in `adk.web.backend-url`; when +`adk.web.backend-url` has a path prefix, the redirect uses that path directly +instead of prepending the context path or `X-Forwarded-Prefix`. + +When `adk.web.backend-url` is unset, `/` and `/dev-ui` redirect to `/dev-ui/` +and the bundled `backendUrl` value is preserved. If you instead enable +`server.forward-headers-strategy=framework`, make sure your reverse proxy strips +or overwrites incoming `Forwarded` and `X-Forwarded-*` headers from untrusted +clients. diff --git a/dev/src/main/java/com/google/adk/plugins/recordings/LlmRecording.java b/dev/src/main/java/com/google/adk/plugins/recordings/LlmRecording.java index 701b1e7ae..5a8e591f7 100644 --- a/dev/src/main/java/com/google/adk/plugins/recordings/LlmRecording.java +++ b/dev/src/main/java/com/google/adk/plugins/recordings/LlmRecording.java @@ -22,7 +22,7 @@ import com.google.auto.value.AutoValue; import java.util.List; import java.util.Optional; -import javax.annotation.Nullable; +import org.jspecify.annotations.Nullable; /** Paired LLM request and response for replay. */ @AutoValue diff --git a/dev/src/main/java/com/google/adk/plugins/recordings/Recording.java b/dev/src/main/java/com/google/adk/plugins/recordings/Recording.java index 259ceaab4..3a164dfc5 100644 --- a/dev/src/main/java/com/google/adk/plugins/recordings/Recording.java +++ b/dev/src/main/java/com/google/adk/plugins/recordings/Recording.java @@ -19,7 +19,7 @@ import com.fasterxml.jackson.databind.annotation.JsonPOJOBuilder; import com.google.auto.value.AutoValue; import java.util.Optional; -import javax.annotation.Nullable; +import org.jspecify.annotations.Nullable; /** Single interaction recording, ordered by request timestamp. */ @AutoValue diff --git a/dev/src/main/java/com/google/adk/plugins/recordings/ToolRecording.java b/dev/src/main/java/com/google/adk/plugins/recordings/ToolRecording.java index a6b1cb26e..bd382adf8 100644 --- a/dev/src/main/java/com/google/adk/plugins/recordings/ToolRecording.java +++ b/dev/src/main/java/com/google/adk/plugins/recordings/ToolRecording.java @@ -21,7 +21,7 @@ import com.google.genai.types.FunctionCall; import com.google.genai.types.FunctionResponse; import java.util.Optional; -import javax.annotation.Nullable; +import org.jspecify.annotations.Nullable; /** Paired tool call and response for replay. */ @AutoValue diff --git a/dev/src/main/java/com/google/adk/web/AdkWebServer.java b/dev/src/main/java/com/google/adk/web/AdkWebServer.java index 0a6035923..ac1879a10 100644 --- a/dev/src/main/java/com/google/adk/web/AdkWebServer.java +++ b/dev/src/main/java/com/google/adk/web/AdkWebServer.java @@ -25,6 +25,7 @@ import com.google.adk.memory.InMemoryMemoryService; import com.google.adk.sessions.BaseSessionService; import com.google.adk.sessions.InMemorySessionService; +import com.google.adk.web.config.BackendUrl; import com.google.adk.web.config.DevUiAssets; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -52,6 +53,15 @@ public class AdkWebServer implements WebMvcConfigurer { @Value("${adk.web.ui.dir:#{null}}") private String webUiDir; + @Value("${adk.web.backend-url:}") + private String backendUrlProperty; + + /** One bean, so this and the runtime-config endpoint cannot read the property differently. */ + @Bean + public BackendUrl backendUrl() { + return BackendUrl.from(backendUrlProperty); + } + @Bean public BaseSessionService sessionService() { // TODO: Add logic to select service based on config (e.g., DB URL) @@ -125,17 +135,25 @@ public void addResourceHandlers(ResourceHandlerRegistry registry) { } /** - * Configures simple automated controllers: "/" and "/dev-ui" both redirect to "/dev-ui/", which - * forwards to the UI's index.html. The trailing slash is required: index.html declares a {@code - * }, so served from "/dev-ui" the app resolves its own router path to "dev-ui" - * and matches none of its routes. The query string is carried across because the UI selects its - * agent from {@code ?app=} and the sample READMEs send users to the slashless "/dev-ui", so a - * redirect that dropped it would silently ignore the selection. + * Configures simple automated controllers: "/" and "/dev-ui" both redirect to the UI, at {@code + * adk.web.backend-url}'s path when a usable value is set, and it forwards to index.html. The + * trailing slash is required: index.html declares a {@code }, so served from + * "/dev-ui" the app resolves its own router path to "dev-ui" and matches none of its routes. The + * query string is carried across because the UI selects its agent from {@code ?app=}. */ @Override public void addViewControllers(ViewControllerRegistry registry) { - registry.addRedirectViewController("/", "/dev-ui/").setKeepQueryParams(true); - registry.addRedirectViewController("/dev-ui", "/dev-ui/").setKeepQueryParams(true); + String prefix = backendUrl().pathPrefix(); + // The configured value is the public base, so do not stack the context path on it. + boolean contextRelative = prefix.isEmpty(); + registry + .addRedirectViewController("/", prefix + "/dev-ui/") + .setKeepQueryParams(true) + .setContextRelative(contextRelative); + registry + .addRedirectViewController("/dev-ui", prefix + "/dev-ui/") + .setKeepQueryParams(true) + .setContextRelative(contextRelative); registry.addViewController("/dev-ui/").setViewName("forward:/dev-ui/index.html"); } @@ -152,7 +170,7 @@ public static void main(String[] args) { log.info("AdkWebServer application started successfully."); } - // TODO(vorburger): #later return Closeable, which can stop the server (and resets static) + // TODO: return a Closeable that stops the server and resets the static state. public static void start(BaseAgent... agents) { // Disable CompiledAgentLoader by setting property to prevent its creation System.setProperty("adk.agents.loader", "static"); diff --git a/dev/src/main/java/com/google/adk/web/config/BackendUrl.java b/dev/src/main/java/com/google/adk/web/config/BackendUrl.java new file mode 100644 index 000000000..35d671e58 --- /dev/null +++ b/dev/src/main/java/com/google/adk/web/config/BackendUrl.java @@ -0,0 +1,118 @@ +/* + * Copyright 2026 Google LLC + * + * Licensed 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 com.google.adk.web.config; + +import com.google.common.base.CharMatcher; +import com.google.common.base.Strings; +import java.net.URI; +import java.net.URISyntaxException; +import org.jspecify.annotations.Nullable; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +/** + * Parsed representation of {@code adk.web.backend-url}, shared by the runtime-config endpoint and + * the dev UI entry redirects. + */ +public final class BackendUrl { + + private static final Logger log = LoggerFactory.getLogger(BackendUrl.class); + + private static final BackendUrl UNSET = new BackendUrl("", ""); + + private final String value; + private final String pathPrefix; + + private BackendUrl(String value, String pathPrefix) { + this.value = value; + this.pathPrefix = pathPrefix; + } + + /** + * Parses {@code configured}. A URL the dev UI cannot use is ignored entirely, with a warning: it + * is neither served nor used as a prefix, so the bundled {@code backendUrl} stands. Serving it + * would publish whatever it contains, credentials included, to every browser. + */ + public static BackendUrl from(@Nullable String configured) { + String trimmed = Strings.nullToEmpty(configured).trim(); + if (trimmed.isEmpty()) { + return UNSET; + } + // Strip trailing slashes because the UI appends paths that already start with "/". + String normalized = CharMatcher.is('/').trimTrailingFrom(trimmed); + URI uri = parse(normalized); + String fault = faultIn(uri); + if (fault == null) { + return new BackendUrl(normalized, pathPrefixOf(uri)); + } + // The value is not echoed: it can carry credentials, and the operator set it themselves. + log.warn("adk.web.backend-url is ignored because {}.", fault); + return UNSET; + } + + /** Returns the URL to report in {@code runtime-config.json}, or empty when unset or ignored. */ + public String value() { + return value; + } + + /** + * Returns the URL path prefix to prepend to the entry redirect, or empty if unset or ignored. + * Preserves raw percent-encoding for use in a {@code Location} header. + */ + public String pathPrefix() { + return pathPrefix; + } + + private static @Nullable URI parse(String url) { + try { + return new URI(url); + } catch (URISyntaxException e) { + return null; + } + } + + private static @Nullable String faultIn(@Nullable URI uri) { + if (uri == null) { + return "it is not a valid URI"; + } + // The UI strips "http://" or "https://" case-sensitively when building its WebSocket URL. + if (!"http".equals(uri.getScheme()) && !"https".equals(uri.getScheme())) { + return "the scheme must be a lower-case http or https"; + } + if (uri.getHost() == null) { + return "it names no host"; + } + if (uri.getRawUserInfo() != null) { + return "it carries credentials, which would be served to every client"; + } + // The UI appends request paths directly onto backendUrl, so a query or fragment would corrupt + // the resulting URL. + if (uri.getRawQuery() != null || uri.getRawFragment() != null) { + return "it carries a query or fragment"; + } + return null; + } + + private static String pathPrefixOf(URI uri) { + String raw = uri.getRawPath(); + if (raw == null || raw.isEmpty()) { + return ""; + } + // Collapse leading slashes so "//segment" cannot be interpreted as a protocol-relative host. + return CharMatcher.is('/').trimTrailingFrom(raw.replaceAll("^/+", "/")); + } +} diff --git a/dev/src/main/java/com/google/adk/web/config/DevUiAssets.java b/dev/src/main/java/com/google/adk/web/config/DevUiAssets.java index 48c742c8b..2ef291c48 100644 --- a/dev/src/main/java/com/google/adk/web/config/DevUiAssets.java +++ b/dev/src/main/java/com/google/adk/web/config/DevUiAssets.java @@ -20,11 +20,15 @@ import org.springframework.core.io.ResourceLoader; /** - * Where the dev UI's static assets live. Normalizes {@code adk.web.ui.dir} into a resource - * location, so callers that need it do not each do it differently. + * Where the dev UI's static assets live. Shared so the resource handler and the runtime-config + * endpoint resolve the same location; normalizing {@code adk.web.ui.dir} separately in each would + * diverge silently. */ public final class DevUiAssets { + /** The runtime config, relative to the asset root. */ + public static final String RUNTIME_CONFIG_PATH = "assets/config/runtime-config.json"; + private static final String CLASSPATH_ROOT = ResourceLoader.CLASSPATH_URL_PREFIX + "/browser/"; /** @@ -42,5 +46,10 @@ public static String assetRoot(@Nullable String webUiDir) { return location.endsWith("/") ? location : location + "/"; } + /** The location of a single asset, given relative to the asset root. */ + public static String assetLocation(@Nullable String webUiDir, String relativePath) { + return assetRoot(webUiDir) + relativePath; + } + private DevUiAssets() {} } diff --git a/dev/src/main/java/com/google/adk/web/controller/RuntimeConfigController.java b/dev/src/main/java/com/google/adk/web/controller/RuntimeConfigController.java new file mode 100644 index 000000000..7f538e2b6 --- /dev/null +++ b/dev/src/main/java/com/google/adk/web/controller/RuntimeConfigController.java @@ -0,0 +1,104 @@ +/* + * Copyright 2026 Google LLC + * + * Licensed 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 com.google.adk.web.controller; + +import com.fasterxml.jackson.core.type.TypeReference; +import com.fasterxml.jackson.databind.ObjectMapper; +import com.google.adk.web.config.BackendUrl; +import com.google.adk.web.config.DevUiAssets; +import java.io.IOException; +import java.io.InputStream; +import java.util.LinkedHashMap; +import java.util.Map; +import org.jspecify.annotations.Nullable; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; +import org.springframework.beans.factory.annotation.Autowired; +import org.springframework.beans.factory.annotation.Value; +import org.springframework.core.io.Resource; +import org.springframework.core.io.ResourceLoader; +import org.springframework.http.CacheControl; +import org.springframework.http.MediaType; +import org.springframework.http.ResponseEntity; +import org.springframework.web.bind.annotation.GetMapping; +import org.springframework.web.bind.annotation.RestController; + +/** + * Serves the dev UI's runtime configuration, shadowing the copy bundled in the static assets so + * {@code adk.web.backend-url} can point the UI at the address browsers reach this server on. The + * bundled document is merged rather than replaced, so keys the UI gains in a later bundle survive. + */ +@RestController +public class RuntimeConfigController { + + private static final Logger log = LoggerFactory.getLogger(RuntimeConfigController.class); + + private final ResourceLoader resourceLoader; + private final ObjectMapper objectMapper; + private final @Nullable String webUiDir; + private final String backendUrl; + + /** Reads the bundled config through {@code resourceLoader}, or from {@code webUiDir} if set. */ + @Autowired + public RuntimeConfigController( + ResourceLoader resourceLoader, + ObjectMapper objectMapper, + @Value("${adk.web.ui.dir:#{null}}") @Nullable String webUiDir, + BackendUrl backendUrl) { + this.resourceLoader = resourceLoader; + this.objectMapper = objectMapper; + this.webUiDir = webUiDir; + this.backendUrl = backendUrl.value(); + } + + /** Serves the bundled config with {@code backendUrl} taken from configuration when set. */ + @GetMapping( + value = "/dev-ui/" + DevUiAssets.RUNTIME_CONFIG_PATH, + produces = MediaType.APPLICATION_JSON_VALUE) + public ResponseEntity> runtimeConfig() { + Map config = readBundledConfig(); + // Unset leaves a value the bundled document already carries, which is what used to be served. + if (backendUrl.isEmpty()) { + config.putIfAbsent("backendUrl", ""); + } else { + config.put("backendUrl", backendUrl); + } + // The bundled document can change on disk under adk.web.ui.dir, so do not let it be cached. + return ResponseEntity.ok().cacheControl(CacheControl.noStore()).body(config); + } + + /** + * Reads the bundled config, returning an empty map if the resource is missing or cannot be parsed + * so the endpoint still serves {@code backendUrl}. + */ + private Map readBundledConfig() { + Resource resource = + resourceLoader.getResource( + DevUiAssets.assetLocation(webUiDir, DevUiAssets.RUNTIME_CONFIG_PATH)); + if (!resource.exists()) { + log.debug("No bundled dev UI runtime config at {}; serving backendUrl only.", resource); + return new LinkedHashMap<>(); + } + try (InputStream in = resource.getInputStream()) { + Map parsed = objectMapper.readValue(in, new TypeReference<>() {}); + return parsed == null ? new LinkedHashMap<>() : new LinkedHashMap<>(parsed); + } catch (IOException e) { + log.warn("Could not read the bundled dev UI runtime config at {}.", resource, e); + return new LinkedHashMap<>(); + } + } +} diff --git a/dev/src/main/java/com/google/adk/web/controller/VersionController.java b/dev/src/main/java/com/google/adk/web/controller/VersionController.java new file mode 100644 index 000000000..841f26b71 --- /dev/null +++ b/dev/src/main/java/com/google/adk/web/controller/VersionController.java @@ -0,0 +1,49 @@ +/* + * Copyright 2026 Google LLC + * + * Licensed 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 com.google.adk.web.controller; + +import com.google.adk.Version; +import java.util.Map; +import org.springframework.http.MediaType; +import org.springframework.web.bind.annotation.GetMapping; +import org.springframework.web.bind.annotation.RestController; + +/** + * Reports the ADK and language versions, which the dev UI requests on startup, plus a liveness + * endpoint. + */ +@RestController +public class VersionController { + + /** Returns the ADK version, the implementation language, and the running JVM's version. */ + @GetMapping(value = "/version", produces = MediaType.APPLICATION_JSON_VALUE) + public Map version() { + return Map.of( + "version", + Version.JAVA_ADK_VERSION, + "language", + "java", + "language_version", + System.getProperty("java.version", "unknown")); + } + + /** Returns a fixed OK status, so a load balancer can tell the server is up. */ + @GetMapping(value = "/health", produces = MediaType.APPLICATION_JSON_VALUE) + public Map health() { + return Map.of("status", "ok"); + } +} diff --git a/dev/src/main/java/com/google/adk/web/dto/AgentRunRequest.java b/dev/src/main/java/com/google/adk/web/dto/AgentRunRequest.java index 025aa64eb..03640ce43 100644 --- a/dev/src/main/java/com/google/adk/web/dto/AgentRunRequest.java +++ b/dev/src/main/java/com/google/adk/web/dto/AgentRunRequest.java @@ -20,7 +20,7 @@ import com.google.adk.JsonBaseModel; import com.google.genai.types.Content; import java.util.Map; -import javax.annotation.Nullable; +import org.jspecify.annotations.Nullable; /** * Data Transfer Object (DTO) for POST /run and POST /run-sse requests. Contains information needed @@ -48,8 +48,7 @@ public class AgentRunRequest { * replay mode settings) without modifying the stored session. */ @JsonProperty("stateDelta") - @Nullable - public Map stateDelta; + public @Nullable Map stateDelta; public AgentRunRequest() {} @@ -83,8 +82,7 @@ public boolean getStreaming() { return streaming; } - @Nullable - public Map getStateDelta() { + public @Nullable Map getStateDelta() { return stateDelta; } } diff --git a/dev/src/test/java/com/google/adk/web/AdkWebServerProxyRedirectTest.java b/dev/src/test/java/com/google/adk/web/AdkWebServerProxyRedirectTest.java index 2f4e3bfd5..3a323d431 100644 --- a/dev/src/test/java/com/google/adk/web/AdkWebServerProxyRedirectTest.java +++ b/dev/src/test/java/com/google/adk/web/AdkWebServerProxyRedirectTest.java @@ -29,14 +29,10 @@ import org.springframework.test.web.servlet.MockMvc; /** - * Spring's own forwarded-header support keeps working: with {@code - * server.forward-headers-strategy=framework} the operator opts in, {@code ForwardedHeaderFilter} - * turns the forwarded prefix into the request's context path, and the redirect picks it up. This - * server adds nothing here and reads no header itself; the test exists so enabling that Spring - * feature keeps behaving as it did. - * - *

The filter honours the standard {@code Forwarded} header as well as {@code X-Forwarded-*}, so - * both are covered. + * Verifies that when {@code adk.web.backend-url} is unset and {@code + * server.forward-headers-strategy=framework} is enabled, Spring's {@code ForwardedHeaderFilter} + * applies {@code Forwarded} and {@code X-Forwarded-*} headers to the entry redirects. When {@code + * adk.web.backend-url} is set, its path prefix takes precedence over {@code X-Forwarded-Prefix}. */ @SpringBootTest(properties = "server.forward-headers-strategy=framework") @AutoConfigureMockMvc diff --git a/dev/src/test/java/com/google/adk/web/BackendUrlRedirectTest.java b/dev/src/test/java/com/google/adk/web/BackendUrlRedirectTest.java new file mode 100644 index 000000000..ce6d61381 --- /dev/null +++ b/dev/src/test/java/com/google/adk/web/BackendUrlRedirectTest.java @@ -0,0 +1,169 @@ +/* + * Copyright 2026 Google LLC + * + * Licensed 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 com.google.adk.web; + +import static org.springframework.test.web.servlet.request.MockMvcRequestBuilders.get; +import static org.springframework.test.web.servlet.result.MockMvcResultMatchers.redirectedUrl; +import static org.springframework.test.web.servlet.result.MockMvcResultMatchers.status; + +import org.junit.jupiter.api.Nested; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.ValueSource; +import org.springframework.beans.factory.annotation.Autowired; +import org.springframework.boot.test.context.SpringBootTest; +import org.springframework.boot.webmvc.test.autoconfigure.AutoConfigureMockMvc; +import org.springframework.test.web.servlet.MockMvc; + +/** + * With {@code adk.web.backend-url} set, the entry redirect carries the gateway's path prefix on its + * own, so a deployment behind a path-stripping proxy needs nothing forwarded from the proxy. The + * property governs the path only: host and scheme still come from forwarded headers wherever the + * operator has enabled {@code server.forward-headers-strategy}. + */ +public class BackendUrlRedirectTest { + + @Nested + @SpringBootTest( + properties = { + "adk.web.backend-url=https://gw.example.com/my-app", + // Enable ForwardedHeaderFilter so X-Forwarded-* headers reach the servlet request. + "server.forward-headers-strategy=framework" + }) + @AutoConfigureMockMvc + class ConfiguredWithTheFilter { + + @ParameterizedTest + @ValueSource(strings = {"/", "/dev-ui"}) + public void devUiEntryPoints_shouldCarryTheConfiguredPrefix( + String path, @Autowired MockMvc mockMvc) throws Exception { + mockMvc + .perform(get(path)) + .andExpect(status().is3xxRedirection()) + .andExpect(redirectedUrl("/my-app/dev-ui/")); + } + + @ParameterizedTest + @ValueSource(strings = {"/", "/dev-ui"}) + public void devUiEntryPoints_shouldNotStackAForwardedPrefix( + String path, @Autowired MockMvc mockMvc) throws Exception { + // The configured prefix replaces the context path created by ForwardedHeaderFilter. + mockMvc + .perform(get(path).header("X-Forwarded-Prefix", "/evil")) + .andExpect(status().is3xxRedirection()) + .andExpect(redirectedUrl("http://localhost/my-app/dev-ui/")); + } + + @Test + public void devUiEntryPoint_forwardedHostAndProto_stillSetTheRedirectsOrigin( + @Autowired MockMvc mockMvc) throws Exception { + // The configured value supplies the path; host and scheme still come from the headers. + mockMvc + .perform( + get("/") + .header("X-Forwarded-Host", "gw.example.com") + .header("X-Forwarded-Proto", "https")) + .andExpect(status().is3xxRedirection()) + .andExpect(redirectedUrl("https://gw.example.com/my-app/dev-ui/")); + } + } + + @Nested + @SpringBootTest(properties = "adk.web.backend-url=https://gw.example.com/my-app") + @AutoConfigureMockMvc + class ConfiguredWithoutTheFilter { + + @Test + public void devUiEntryPoint_forwardedHeaders_reachNothingUnderTheShippedStrategy( + @Autowired MockMvc mockMvc) throws Exception { + // With forward-headers-strategy unset (default), forwarded headers are ignored. + mockMvc + .perform( + get("/") + .header("X-Forwarded-Prefix", "/evil") + .header("X-Forwarded-Host", "evil.example.com") + .header("Forwarded", "host=evil.example.com;proto=https")) + .andExpect(status().is3xxRedirection()) + .andExpect(redirectedUrl("/my-app/dev-ui/")); + } + } + + @Nested + @SpringBootTest(properties = "adk.web.backend-url=https://gw.example.com") + @AutoConfigureMockMvc + class ConfiguredWithNoPath { + + @ParameterizedTest + @ValueSource(strings = {"/", "/dev-ui"}) + public void devUiEntryPoints_shouldRedirectWithoutAPrefix( + String path, @Autowired MockMvc mockMvc) throws Exception { + // A host-only URL has an empty pathPrefix(), so the redirect stays context-relative. + mockMvc + .perform(get(path)) + .andExpect(status().is3xxRedirection()) + .andExpect(redirectedUrl("/dev-ui/")); + } + } + + @Nested + @SpringBootTest(properties = "adk.web.backend-url=/my-app") + @AutoConfigureMockMvc + class ConfiguredUnusably { + + @ParameterizedTest + @ValueSource(strings = {"/", "/dev-ui"}) + public void devUiEntryPoints_shouldRedirectWithoutAPrefix( + String path, @Autowired MockMvc mockMvc) throws Exception { + // A scheme-less value is unusable by the UI, so pathPrefix() is empty. + mockMvc + .perform(get(path)) + .andExpect(status().is3xxRedirection()) + .andExpect(redirectedUrl("/dev-ui/")); + } + } + + @Nested + @SpringBootTest + @AutoConfigureMockMvc + class UnconfiguredWithoutTheFilter { + + @ParameterizedTest + @ValueSource(strings = {"/", "/dev-ui"}) + public void devUiEntryPoints_shouldRedirectWithoutAPrefix( + String path, @Autowired MockMvc mockMvc) throws Exception { + mockMvc + .perform(get(path)) + .andExpect(status().is3xxRedirection()) + .andExpect(redirectedUrl("/dev-ui/")); + } + + @ParameterizedTest + @ValueSource(strings = {"/", "/dev-ui"}) + public void devUiEntryPoints_withForwardedHeaders_shouldStillRedirectWithoutAPrefix( + String path, @Autowired MockMvc mockMvc) throws Exception { + // By default forward-headers-strategy is unset, so forwarded headers are ignored. + mockMvc + .perform( + get(path) + .header("X-Forwarded-Prefix", "/evil") + .header("X-Forwarded-Host", "evil.example.com") + .header("Forwarded", "host=evil.example.com;proto=https")) + .andExpect(status().is3xxRedirection()) + .andExpect(redirectedUrl("/dev-ui/")); + } + } +} diff --git a/dev/src/test/java/com/google/adk/web/config/BackendUrlTest.java b/dev/src/test/java/com/google/adk/web/config/BackendUrlTest.java new file mode 100644 index 000000000..d870d0169 --- /dev/null +++ b/dev/src/test/java/com/google/adk/web/config/BackendUrlTest.java @@ -0,0 +1,203 @@ +/* + * Copyright 2026 Google LLC + * + * Licensed 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 com.google.adk.web.config; + +import static com.google.common.truth.Truth.assertThat; + +import ch.qos.logback.classic.Level; +import ch.qos.logback.classic.Logger; +import ch.qos.logback.classic.spi.ILoggingEvent; +import ch.qos.logback.core.read.ListAppender; +import java.util.List; +import org.junit.jupiter.api.Test; +import org.slf4j.LoggerFactory; + +/** + * One reading of {@code adk.web.backend-url}. The value the dev UI is served and the prefix its + * entry redirect carries come from the same parse, so they cannot disagree. + */ +public class BackendUrlTest { + + private static final String CONFIGURED = "https://gw.example.com/my-app"; + + @Test + public void unset_isEmptyEverywhere() { + for (String in : new String[] {null, "", " "}) { + assertThat(BackendUrl.from(in).value()).isEmpty(); + assertThat(BackendUrl.from(in).pathPrefix()).isEmpty(); + } + } + + @Test + public void absoluteUrl_isServedAndSuppliesThePrefix() { + BackendUrl url = BackendUrl.from(CONFIGURED); + + assertThat(url.value()).isEqualTo(CONFIGURED); + assertThat(url.pathPrefix()).isEqualTo("/my-app"); + } + + @Test + public void trailingSlashes_areStripped() { + // The UI appends paths starting with a slash, so a trailing one would yield //run_live. + assertThat(BackendUrl.from(CONFIGURED + "/").value()).isEqualTo(CONFIGURED); + assertThat(BackendUrl.from(CONFIGURED + "///").pathPrefix()).isEqualTo("/my-app"); + assertThat(BackendUrl.from(" " + CONFIGURED + " ").value()).isEqualTo(CONFIGURED); + } + + @Test + public void hostWithoutPath_hasNoPrefix() { + assertThat(BackendUrl.from("https://gw.example.com").pathPrefix()).isEmpty(); + assertThat(BackendUrl.from("https://gw.example.com/").pathPrefix()).isEmpty(); + } + + @Test + public void percentEncoding_isKept() { + // This goes into a Location header, so decoding would re-encode wrongly and %2F would + // turn into a path separator. + assertThat(BackendUrl.from("https://gw.example.com/my%20app").pathPrefix()) + .isEqualTo("/my%20app"); + assertThat(BackendUrl.from("https://gw.example.com/a%2Fb").pathPrefix()).isEqualTo("/a%2Fb"); + } + + @Test + public void doubledSlash_doesNotBecomeAHost() { + // "//my-app/dev-ui/" is protocol-relative: a browser resolves it to the host "my-app". + assertThat(BackendUrl.from("https://gw.example.com//my-app").pathPrefix()).isEqualTo("/my-app"); + } + + @Test + public void unusableValue_suppliesNoPrefix() { + // Leave pathPrefix() empty when the UI cannot use the URL so redirect and UI behavior match. + for (String in : + new String[] { + "/my-app", + "HTTPS://gw.example.com/x", + "http://", + "gw.example.com", + "https://gw.example.com/a?q=1", + "https://user:pass@gw.example.com/my-app", + "https://gw.example.com/my app", + // No host at all, only a port: getRawAuthority() would call this usable. + "https://:8080/my-app", + // RFC-invalid, and unreachable anyway - Tomcat rejects such a Host header with a 400, + // and no CA will issue a certificate for one. + "https://gw_host.example.com/my-app" + }) { + assertThat(BackendUrl.from(in).pathPrefix()).isEmpty(); + } + } + + @Test + public void theValueIsNeverLogged() { + // Redacting userinfo was not enough: a query, or a "@" in the password, survived the mask. + for (String in : + new String[] { + "https://user:pass@gw.example.com/my-app", + "https://user:p@ss/word@gw.example.com/my-app", + "https://gw.example.com/my-app?token=sekrit" + }) { + assertThat(warningsFor(in)).hasSize(1); + String warning = warningsFor(in).get(0); + + assertThat(warning).doesNotContain("pass"); + assertThat(warning).doesNotContain("sekrit"); + assertThat(warning).doesNotContain(in); + // The old mask kept the host, so without this two of these three inputs pinned nothing. + assertThat(warning).doesNotContain("gw.example.com"); + } + } + + @Test + public void credentials_areNeverServed() { + // Served, this would put the password in runtime-config.json for every browser. + BackendUrl url = BackendUrl.from("https://user:pass@gw.example.com/my-app"); + + assertThat(url.value()).isEmpty(); + assertThat(url.pathPrefix()).isEmpty(); + } + + @Test + public void unusableValue_isIgnoredEntirelyButWarns() { + // Not served and not prefixed, so the bundled backendUrl stands. + for (String in : new String[] {"/my-app", "HTTPS://gw.example.com/x", "http://"}) { + assertThat(BackendUrl.from(in).value()).isEmpty(); + assertThat(warningsFor(in)).hasSize(1); + } + } + + @Test + public void slashOnly_isIgnored() { + // "/" trims to empty, and is not a usable backend URL either way. + assertThat(BackendUrl.from("/").value()).isEmpty(); + assertThat(BackendUrl.from("/").pathPrefix()).isEmpty(); + } + + @Test + public void queryFragmentOrUserinfo_yieldsNoPrefixAndWarns() { + // The UI appends onto the value, so these would land in the middle of every request. + for (String in : + new String[] { + "https://gw.example.com/a/b?q=1", + "https://gw.example.com/a/b#f", + "https://user:pass@gw.example.com/my-app" + }) { + assertThat(BackendUrl.from(in).pathPrefix()).isEmpty(); + assertThat(warningsFor(in)).hasSize(1); + } + } + + @Test + public void valueWithoutALeadingSlashPath_yieldsNoPrefix() { + // "gw.example.com/dev-ui/" in a Location header would be resolved as a relative path. + assertThat(BackendUrl.from("gw.example.com").pathPrefix()).isEmpty(); + assertThat(warningsFor("gw.example.com")).hasSize(1); + } + + @Test + public void unparseableValue_isIgnoredAndWarns() { + String malformed = "https://gw.example.com/my app"; + + assertThat(BackendUrl.from(malformed).value()).isEmpty(); + assertThat(BackendUrl.from(malformed).pathPrefix()).isEmpty(); + assertThat(warningsFor(malformed)).hasSize(1); + assertThat(warningsFor(malformed).get(0)).contains("not a valid URI"); + } + + @Test + public void usableValue_doesNotWarn() { + assertThat(warningsFor(CONFIGURED)).isEmpty(); + assertThat(warningsFor(CONFIGURED + "/")).isEmpty(); + assertThat(warningsFor("")).isEmpty(); + } + + /** The WARN messages logged while interpreting {@code configured}. */ + private static List warningsFor(String configured) { + Logger logger = (Logger) LoggerFactory.getLogger(BackendUrl.class); + ListAppender appender = new ListAppender<>(); + appender.start(); + logger.addAppender(appender); + try { + BackendUrl unused = BackendUrl.from(configured); + } finally { + logger.detachAppender(appender); + } + return appender.list.stream() + .filter(event -> event.getLevel() == Level.WARN) + .map(ILoggingEvent::getFormattedMessage) + .toList(); + } +} diff --git a/dev/src/test/java/com/google/adk/web/config/DevUiAssetsTest.java b/dev/src/test/java/com/google/adk/web/config/DevUiAssetsTest.java index 668f75e57..46ccaf97c 100644 --- a/dev/src/test/java/com/google/adk/web/config/DevUiAssetsTest.java +++ b/dev/src/test/java/com/google/adk/web/config/DevUiAssetsTest.java @@ -42,4 +42,10 @@ public void assetRoot_configuredDir_becomesAFileUrl(String webUiDir) { public void assetRoot_windowsSeparators_areNormalized() { assertThat(DevUiAssets.assetRoot("C:\\srv\\ui")).isEqualTo("file:C:/srv/ui/"); } + + @Test + public void assetLocation_appendsToTheRoot() { + assertThat(DevUiAssets.assetLocation(null, DevUiAssets.RUNTIME_CONFIG_PATH)) + .isEqualTo("classpath:/browser/assets/config/runtime-config.json"); + } } diff --git a/dev/src/test/java/com/google/adk/web/controller/RuntimeConfigControllerMergeTest.java b/dev/src/test/java/com/google/adk/web/controller/RuntimeConfigControllerMergeTest.java new file mode 100644 index 000000000..a01dd9462 --- /dev/null +++ b/dev/src/test/java/com/google/adk/web/controller/RuntimeConfigControllerMergeTest.java @@ -0,0 +1,123 @@ +/* + * Copyright 2026 Google LLC + * + * Licensed 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 com.google.adk.web.controller; + +import static com.google.common.truth.Truth.assertThat; + +import com.fasterxml.jackson.databind.ObjectMapper; +import com.google.adk.web.config.BackendUrl; +import java.nio.charset.StandardCharsets; +import java.util.Map; +import org.junit.jupiter.api.Test; +import org.springframework.core.io.ByteArrayResource; +import org.springframework.core.io.DescriptiveResource; +import org.springframework.core.io.Resource; +import org.springframework.core.io.ResourceLoader; + +/** + * The served runtime config merges into the bundled document rather than replacing it, so keys the + * dev UI bundle gains later are not silently dropped. {@code adk.web.backend-url} overrides one + * key; an unset value leaves the bundled document exactly as it was served before. + */ +public class RuntimeConfigControllerMergeTest { + + private static final String CONFIGURED = "https://gw.example.com/my-app"; + + @Test + public void runtimeConfig_shouldPreserveOtherBundledKeys() { + String bundled = "{\"backendUrl\":\"\",\"telemetry\":null,\"logo\":{\"text\":\"x\"}}"; + RuntimeConfigController controller = controllerFor(bundled, CONFIGURED); + + Map config = controller.runtimeConfig().getBody(); + + assertThat(config).containsEntry("backendUrl", CONFIGURED); + assertThat(config).containsKey("telemetry"); + assertThat(config).containsEntry("logo", Map.of("text", "x")); + } + + @Test + public void runtimeConfig_shouldOverrideTheBundledBackendUrl() { + RuntimeConfigController controller = + controllerFor("{\"backendUrl\":\"http://stale\"}", CONFIGURED); + + assertThat(controller.runtimeConfig().getBody()).containsEntry("backendUrl", CONFIGURED); + } + + @Test + public void runtimeConfig_propertyUnset_shouldKeepTheBundledBackendUrl() { + // The static handler served this file verbatim, so a hand-set value survived; it still does. + RuntimeConfigController controller = + controllerFor("{\"backendUrl\":\"http://elsewhere:9000\"}"); + + assertThat(controller.runtimeConfig().getBody()) + .containsExactly("backendUrl", "http://elsewhere:9000"); + } + + @Test + public void runtimeConfig_propertyUnsetAndNoBundledKey_shouldStillReportBackendUrl() { + RuntimeConfigController controller = controllerFor("{\"telemetry\":null}"); + + assertThat(controller.runtimeConfig().getBody()).containsEntry("backendUrl", ""); + } + + @Test + public void runtimeConfig_bundledFileMissing_shouldStillServe() { + RuntimeConfigController controller = controllerFor(null, CONFIGURED); + + assertThat(controller.runtimeConfig().getBody()).containsExactly("backendUrl", CONFIGURED); + } + + @Test + public void runtimeConfig_bundledFileMalformed_shouldStillServe() { + RuntimeConfigController controller = controllerFor("not json at all", CONFIGURED); + + assertThat(controller.runtimeConfig().getBody()).containsExactly("backendUrl", CONFIGURED); + } + + @Test + public void runtimeConfig_bundledFileNotAnObject_shouldStillServe() { + // A config that is not a JSON object is ignored rather than failing the request. + RuntimeConfigController controller = controllerFor("[1, 2, 3]", CONFIGURED); + + assertThat(controller.runtimeConfig().getBody()).containsExactly("backendUrl", CONFIGURED); + } + + /** A controller with {@code adk.web.backend-url} unset. */ + private static RuntimeConfigController controllerFor(String body) { + return controllerFor(body, ""); + } + + /** A controller whose bundled config is {@code body}, or absent when {@code body} is null. */ + private static RuntimeConfigController controllerFor(String body, String backendUrl) { + ResourceLoader loader = + new ResourceLoader() { + @Override + public Resource getResource(String location) { + return body == null + ? new DescriptiveResource("absent") + : new ByteArrayResource(body.getBytes(StandardCharsets.UTF_8)); + } + + @Override + public ClassLoader getClassLoader() { + return RuntimeConfigControllerMergeTest.class.getClassLoader(); + } + }; + return new RuntimeConfigController( + loader, new ObjectMapper(), null, BackendUrl.from(backendUrl)); + } +} diff --git a/dev/src/test/java/com/google/adk/web/controller/RuntimeConfigControllerTest.java b/dev/src/test/java/com/google/adk/web/controller/RuntimeConfigControllerTest.java new file mode 100644 index 000000000..1bbf58718 --- /dev/null +++ b/dev/src/test/java/com/google/adk/web/controller/RuntimeConfigControllerTest.java @@ -0,0 +1,140 @@ +/* + * Copyright 2026 Google LLC + * + * Licensed 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 com.google.adk.web.controller; + +import static org.hamcrest.Matchers.containsString; +import static org.hamcrest.Matchers.not; +import static org.springframework.test.web.servlet.request.MockMvcRequestBuilders.get; +import static org.springframework.test.web.servlet.result.MockMvcResultMatchers.content; +import static org.springframework.test.web.servlet.result.MockMvcResultMatchers.header; +import static org.springframework.test.web.servlet.result.MockMvcResultMatchers.jsonPath; +import static org.springframework.test.web.servlet.result.MockMvcResultMatchers.status; + +import org.junit.jupiter.api.Nested; +import org.junit.jupiter.api.Test; +import org.springframework.beans.factory.annotation.Autowired; +import org.springframework.boot.test.context.SpringBootTest; +import org.springframework.boot.webmvc.test.autoconfigure.AutoConfigureMockMvc; +import org.springframework.test.web.servlet.MockMvc; + +/** + * The dev UI learns where its backend lives from {@code adk.web.backend-url}, so its API calls + * reach a server published under a path prefix. Nothing is read from the request, so a client + * cannot influence the value it is served. + */ +public class RuntimeConfigControllerTest { + + private static final String CONFIG = "/dev-ui/assets/config/runtime-config.json"; + + @Nested + @SpringBootTest( + properties = { + "adk.web.backend-url=https://gw.example.com/my-app", + // Enable ForwardedHeaderFilter so X-Forwarded-* headers reach the servlet request. + "server.forward-headers-strategy=framework" + }) + @AutoConfigureMockMvc + class ConfiguredWithTheFilter { + + @Test + public void runtimeConfig_shouldReportTheConfiguredBackendUrl(@Autowired MockMvc mockMvc) + throws Exception { + mockMvc + .perform(get(CONFIG)) + .andExpect(status().isOk()) + .andExpect(jsonPath("$.backendUrl").value("https://gw.example.com/my-app")) + .andExpect(header().string("Cache-Control", "no-store")); + } + + @Test + public void runtimeConfig_forwardedHeaders_shouldNotChangeTheValue(@Autowired MockMvc mockMvc) + throws Exception { + mockMvc + .perform( + get(CONFIG) + .header("X-Forwarded-Prefix", "/evil") + .header("X-Forwarded-Host", "evil.example.com") + .header("X-Forwarded-Proto", "https")) + .andExpect(status().isOk()) + .andExpect(jsonPath("$.backendUrl").value("https://gw.example.com/my-app")); + } + } + + @Nested + @SpringBootTest(properties = "adk.web.backend-url=https://user:hunter2@gw.example.com/my-app") + @AutoConfigureMockMvc + class ConfiguredWithCredentials { + + @Test + public void runtimeConfig_shouldNotServeTheCredentials(@Autowired MockMvc mockMvc) + throws Exception { + // This document goes to every browser that loads the UI. + mockMvc + .perform(get(CONFIG)) + .andExpect(status().isOk()) + .andExpect(content().string(not(containsString("hunter2")))) + .andExpect(jsonPath("$.backendUrl").value("")); + } + } + + @Nested + @SpringBootTest + @AutoConfigureMockMvc + class UnconfiguredWithoutTheFilter { + + @Test + public void runtimeConfig_forwardedHeaders_areNotReadFromTheRequest(@Autowired MockMvc mockMvc) + throws Exception { + mockMvc + .perform( + get(CONFIG) + .header("X-Forwarded-Prefix", "/evil") + .header("X-Forwarded-Host", "evil.example.com") + .header("Forwarded", "host=evil.example.com;proto=https")) + .andExpect(status().isOk()) + .andExpect(jsonPath("$.backendUrl").value("")); + } + } + + @Nested + @SpringBootTest(properties = "server.forward-headers-strategy=framework") + @AutoConfigureMockMvc + class UnconfiguredWithTheFilter { + + @Test + public void runtimeConfig_shouldServeWhatTheBundledFileSaid(@Autowired MockMvc mockMvc) + throws Exception { + mockMvc + .perform(get(CONFIG)) + .andExpect(status().isOk()) + // Compact JSON proves the controller answered rather than the pretty-printed static file. + .andExpect(content().string("{\"backendUrl\":\"\"}")); + } + + @Test + public void runtimeConfig_forwardedHeaders_shouldStillReportEmpty(@Autowired MockMvc mockMvc) + throws Exception { + mockMvc + .perform( + get(CONFIG) + .header("X-Forwarded-Prefix", "/evil") + .header("X-Forwarded-Host", "evil.example.com")) + .andExpect(status().isOk()) + .andExpect(jsonPath("$.backendUrl").value("")); + } + } +} diff --git a/dev/src/test/java/com/google/adk/web/controller/VersionControllerTest.java b/dev/src/test/java/com/google/adk/web/controller/VersionControllerTest.java new file mode 100644 index 000000000..390b3ecc1 --- /dev/null +++ b/dev/src/test/java/com/google/adk/web/controller/VersionControllerTest.java @@ -0,0 +1,58 @@ +/* + * Copyright 2026 Google LLC + * + * Licensed 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 com.google.adk.web.controller; + +import static org.springframework.test.web.servlet.request.MockMvcRequestBuilders.get; +import static org.springframework.test.web.servlet.result.MockMvcResultMatchers.content; +import static org.springframework.test.web.servlet.result.MockMvcResultMatchers.jsonPath; +import static org.springframework.test.web.servlet.result.MockMvcResultMatchers.status; + +import com.google.adk.Version; +import org.junit.jupiter.api.Test; +import org.springframework.beans.factory.annotation.Autowired; +import org.springframework.boot.test.context.SpringBootTest; +import org.springframework.boot.webmvc.test.autoconfigure.AutoConfigureMockMvc; +import org.springframework.http.MediaType; +import org.springframework.test.web.servlet.MockMvc; + +/** The dev UI requests /version on startup; it used to 404. */ +@SpringBootTest +@AutoConfigureMockMvc +public class VersionControllerTest { + + @Autowired private MockMvc mockMvc; + + @Test + public void version_shouldReportAdkAndLanguageVersions() throws Exception { + mockMvc + .perform(get("/version")) + .andExpect(status().isOk()) + .andExpect(content().contentType(MediaType.APPLICATION_JSON)) + .andExpect(jsonPath("$.version").value(Version.JAVA_ADK_VERSION)) + .andExpect(jsonPath("$.language").value("java")) + .andExpect(jsonPath("$.language_version").value(System.getProperty("java.version"))); + } + + @Test + public void health_shouldReportOk() throws Exception { + mockMvc + .perform(get("/health")) + .andExpect(status().isOk()) + .andExpect(content().contentType(MediaType.APPLICATION_JSON)) + .andExpect(content().json("{\"status\":\"ok\"}")); + } +} diff --git a/tokt/src/main/kotlin/com/google/adk/tokt/JavaAdkToKt.kt b/tokt/src/main/kotlin/com/google/adk/tokt/JavaAdkToKt.kt index e75ca84e4..c1959bb86 100644 --- a/tokt/src/main/kotlin/com/google/adk/tokt/JavaAdkToKt.kt +++ b/tokt/src/main/kotlin/com/google/adk/tokt/JavaAdkToKt.kt @@ -41,32 +41,32 @@ import kotlinx.coroutines.CoroutineDispatcher /** * Forward interop entry point: adapts ADK Java tools, toolsets, plugins, services, and models so - * they can run on the ADK Kotlin engine. Wrap the adapted pieces in a Kotlin `LlmAgent`; this does - * not convert a whole Java agent. + * they can run on the ADK Kotlin engine. Assemble the adapted pieces into a Kotlin `LlmAgent`, + * `App`, and `Runner`; this does not convert a whole Java agent. * * An adapted component behaves as it does on ADK Java. It sees the session as it currently stands, * including events and state written earlier in the same turn, and its state, artifact and - * control-flow writes reach the engine. Blocking work is fine: calls are dispatched off the thread - * driving the agent, onto the optional `dispatcher` (default `Dispatchers.IO`) each conversion - * accepts. That dispatcher must be able to run nested bridged calls concurrently -- a bridged tool - * or plugin that itself makes a blocking bridged call (e.g. one that blocks on a bridged service) - * holds its thread until that call returns, so a single-threaded or tightly bounded dispatcher can - * deadlock; the default `Dispatchers.IO` grows its pool and avoids this. + * control-flow writes reach the engine. Blocking work is fine because calls run on the optional + * `dispatcher` (default `Dispatchers.IO`) that each conversion accepts. That dispatcher must be + * able to run nested bridged calls in parallel: a bridged tool or plugin that blocks on another + * bridged call, such as a service call, holds its thread until that call returns, so a + * single-threaded or tightly bounded dispatcher can deadlock. `Dispatchers.IO` deadlocks only if + * all of its threads (at least 64 by default) block at once. * * A bridged plugin's error callbacks fire: `onRunErrorCallback` is notification-only -- the engine * re-raises the run's error to the caller afterwards regardless, so it cannot recover the run (it * is for logging, telemetry, or cleanup) -- while the `onModelErrorCallback` and * `onToolErrorCallback` recovery hooks fire and can recover. * - * The interop surfaces a signal the engine cannot honor rather than silently dropping it: - * - Setting `branch` on a bridged context throws. The branch is the engine's to set. - * - A bridged tool's or plugin's `requestedAuthConfigs` or `deletedArtifactIds` write throws - the - * engine's event actions have no equivalent. Its `skipSummarization`, - * `requestedToolConfirmations` and `agentState` writes do cross, from a tool and a plugin alike. - * - Behind an adapted Java session service ([asKtSessionService]), a resumable workflow's engine - * state (`EventActions.agentState`) crosses and is restored, so it resumes rather than restarts. - * A rewind request (`rewindBeforeInvocationId`) still does not cross - ADK Java's `EventActions` - * has no such field. + * When a bridged tool call or plugin callback writes one of the following signals to its context, + * the interop throws rather than silently dropping it: + * - Setting `branch` on the invocation context throws. The branch is the engine's to set. + * - If Java code writes `requestedAuthConfigs` or `deletedArtifactIds` to the context's + * `EventActions`, the adapter throws once the Java call returns - the engine's event actions have + * no equivalent. + * + * Writes to the context's `skipSummarization`, `requestedToolConfirmations`, `agentState`, and + * `rewindBeforeInvocationId` do cross to the engine. */ object JavaAdkToKt { @@ -79,7 +79,7 @@ object JavaAdkToKt { ): KtBaseTool = JavaToolToKt(javaTool, dispatcher) /** - * Adapts a whole collection of ADK Java tools (e.g. an `LlmAgent`'s `tools`), each on + * Adapts a whole collection of ADK Java tools (such as an `LlmAgent`'s `tools`), each on * `dispatcher`. Kept alongside [asKtTool] for Java callers, who would otherwise write * `stream().map(...).toList()`. */ @@ -115,7 +115,7 @@ object JavaAdkToKt { ): KtPlugin = JavaPluginToKt(javaPlugin, dispatcher) /** - * Adapts a whole collection of ADK Java plugins (e.g. a `Runner`'s `plugins`), each on + * Adapts a whole collection of ADK Java plugins (such as an `App`'s `plugins()`), each on * `dispatcher`. */ @JvmStatic @@ -137,9 +137,11 @@ object JavaAdkToKt { ): KtModel = JavaModelToKt(javaLlm, dispatcher) /** - * Adapts an ADK Java session service for the Kotlin engine, unwrapping a round-tripped Kotlin one - * rather than stacking a second adapter. Its calls run on `dispatcher`. A - * `rewindBeforeInvocationId` does not survive, since ADK Java has no such field. + * Adapts an ADK Java session service for the Kotlin engine, running its calls on `dispatcher`. A + * Java view of a Kotlin service is unwrapped to that Kotlin service, with no `dispatcher` hop. + * Resumption works across the adapter (`EventActions.agentState` crosses); + * `rewindBeforeInvocationId` crosses too, but a rewind drops the rewound turns only if the Java + * service stores it, which ADK Java's `VertexAiSessionService` does not. */ @JvmStatic @JvmOverloads @@ -149,9 +151,9 @@ object JavaAdkToKt { ): KtSessionService = javaSessionServiceAsKt(service, dispatcher) /** - * Adapts an ADK Java artifact service for the Kotlin engine, unwrapping a round-tripped Kotlin - * one rather than stacking adapters. Its calls run on `dispatcher`. An empty or unmapped artifact - * part is rejected outright. + * Adapts an ADK Java artifact service for the Kotlin engine, running its calls on `dispatcher`. A + * Java view of a Kotlin service is unwrapped to that Kotlin service, with no `dispatcher` hop. + * The adapter throws on an empty artifact part or one it cannot convert. */ @JvmStatic @JvmOverloads @@ -161,8 +163,8 @@ object JavaAdkToKt { ): KtArtifactService = javaArtifactServiceAsKt(service, dispatcher) /** - * Adapts an ADK Java memory service for the Kotlin engine, unwrapping a round-tripped Kotlin one - * rather than stacking adapters. Its calls run on `dispatcher`. + * Adapts an ADK Java memory service for the Kotlin engine, running its calls on `dispatcher`. A + * Java view of a Kotlin service is unwrapped to that Kotlin service, with no `dispatcher` hop. */ @JvmStatic @JvmOverloads diff --git a/tokt/src/main/kotlin/com/google/adk/tokt/KotlinAdkToJava.kt b/tokt/src/main/kotlin/com/google/adk/tokt/KotlinAdkToJava.kt index 9f0740924..b68e0342c 100644 --- a/tokt/src/main/kotlin/com/google/adk/tokt/KotlinAdkToJava.kt +++ b/tokt/src/main/kotlin/com/google/adk/tokt/KotlinAdkToJava.kt @@ -21,19 +21,19 @@ import com.google.adk.runner.Runner as JavaRunner import kotlinx.coroutines.CoroutineDispatcher /** - * Reverse interop entry point: exposes an ADK Kotlin-engine [KtRunner] through the ADK Java - * [JavaRunner] surface, so it can be injected into code written against the Java runner. The - * returned runner is a real [JavaRunner] whose `runAsync` streams `Event`s backed by the Kotlin - * engine; live mode is not bridged. The forward direction (Java components onto the Kotlin engine) - * lives in [com.google.adk.tokt.JavaAdkToKt]. + * Reverse interop entry point: [asJavaRunner] wraps an ADK Kotlin-engine `Runner` in a real ADK + * Java `Runner`, so it can be injected into code written against the Java runner. The wrapper's + * `runAsync` streams `Event`s backed by the Kotlin engine; live mode is not bridged. The forward + * direction (Java components onto the Kotlin engine) lives in [com.google.adk.tokt.JavaAdkToKt]. */ object KotlinAdkToJava { /** - * Exposes a Kotlin-engine [runner] as an ADK Java [JavaRunner]. Its reverse service adapters - * bridge the Java RxJava calls onto the Kotlin engine via `dispatcher` (default - * `Dispatchers.IO`), which must be able to run nested bridged calls concurrently, so a - * single-threaded or tightly bounded dispatcher can deadlock. + * Exposes a Kotlin-engine [runner] as an ADK Java `Runner`. Its reverse service adapters bridge + * RxJava calls onto the Kotlin engine via `dispatcher` (default `Dispatchers.IO`). That + * dispatcher must be able to run nested bridged calls in parallel: if code on one of its threads + * (such as a bridged call's RxJava callback) blocks on another bridged call, it holds that thread + * until the second call returns, so a single-threaded or tightly bounded dispatcher can deadlock. */ @JvmStatic @JvmOverloads diff --git a/tokt/src/main/kotlin/com/google/adk/tokt/KtRunnerToJava.kt b/tokt/src/main/kotlin/com/google/adk/tokt/KtRunnerToJava.kt index c141d184e..0467ffec1 100644 --- a/tokt/src/main/kotlin/com/google/adk/tokt/KtRunnerToJava.kt +++ b/tokt/src/main/kotlin/com/google/adk/tokt/KtRunnerToJava.kt @@ -18,10 +18,13 @@ package com.google.adk.tokt import com.google.adk.agents.LiveRequestQueue import com.google.adk.agents.RunConfig as JavaRunConfig +import com.google.adk.annotations.Experimental import com.google.adk.artifacts.BaseArtifactService as JavaArtifactService import com.google.adk.artifacts.InMemoryArtifactService as JavaInMemoryArtifactService import com.google.adk.events.Event as JavaEvent import com.google.adk.kt.runners.Runner as KtRunner +import com.google.adk.kt.types.Content as KtContent +import com.google.adk.kt.types.Role as KtRole import com.google.adk.plugins.PluginManager as JavaPluginManager import com.google.adk.runner.Runner as JavaRunner import com.google.adk.sessions.Session as JavaSession @@ -78,6 +81,29 @@ internal class KtRunnerToJava(private val ktRunner: KtRunner, dispatcher: Corout newMessage: GenaiContent, runConfig: JavaRunConfig, stateDelta: MutableMap?, + ): Flowable = + runAsync( + userId, + sessionId, + invocationId = null, + newMessage = newMessage, + runConfig = runConfig, + stateDelta = stateDelta, + ) + + /** + * Resumes on the Kotlin engine with its semantics instead of the Java runner's. A non-resumable + * app runs [newMessage] as a new invocation under [invocationId], as Python 2.x does, instead of + * throwing. ADK Kotlin 1.1.0 and earlier drop [stateDelta] when [newMessage] is null. + */ + @Experimental + override fun runAsync( + userId: String, + sessionId: String, + invocationId: String?, + newMessage: GenaiContent?, + runConfig: JavaRunConfig, + stateDelta: MutableMap?, ): Flowable { // Defer so a request-conversion or RunConfig rejection surfaces via onError, not at the call // site - matching the base Runner and this class's own runLive. @@ -87,7 +113,8 @@ internal class KtRunnerToJava(private val ktRunner: KtRunner, dispatcher: Corout .runAsync( userId = userId, sessionId = sessionId, - newMessage = ContentCodec.fromJava(newMessage), + invocationId = invocationId, + newMessage = newMessage?.let { withUserRole(ContentCodec.fromJava(it)) }, // Translate the Java REMOVED sentinel so a caller-passed deletion deletes the key // rather than storing it as a value (identity-matched by the engine). stateDelta = stateDelta?.let { stateDeltaFromJava(it) }, @@ -146,4 +173,11 @@ internal class KtRunnerToJava(private val ktRunner: KtRunner, dispatcher: Corout private fun liveUnsupported() = UnsupportedOperationException("Live mode is not supported when running on the Kotlin engine.") + + /** + * Defaults a role-less message to the user role, as Python does, since model requests omit + * role-less content. + */ + private fun withUserRole(content: KtContent): KtContent = + if (content.role.isNullOrEmpty()) content.copy(role = KtRole.USER) else content } diff --git a/tokt/src/main/kotlin/com/google/adk/tokt/adapters/ActionReconciler.kt b/tokt/src/main/kotlin/com/google/adk/tokt/adapters/ActionReconciler.kt index cf63f30c0..340a99529 100644 --- a/tokt/src/main/kotlin/com/google/adk/tokt/adapters/ActionReconciler.kt +++ b/tokt/src/main/kotlin/com/google/adk/tokt/adapters/ActionReconciler.kt @@ -26,8 +26,9 @@ import com.google.adk.tokt.codecs.reconcileRemovedSentinels /** * Copies what a bridged tool or plugin callback wrote on its [JavaEventActions] onto the live * [KtEventActions] the engine reads. Fields the live view already writes through are re-applied - * idempotently; skip-summarization, tool confirmations, agentState and compaction are carried here. - * A signal the engine cannot represent (requestedAuthConfigs, deletedArtifactIds) throws. + * idempotently; skip-summarization, tool confirmations, agentState, compaction and + * rewindBeforeInvocationId are carried here. A signal the engine cannot represent + * (requestedAuthConfigs, deletedArtifactIds) throws. */ internal fun reconcileActionsToKt(java: JavaEventActions, kt: KtEventActions) { require(java.requestedAuthConfigs().isNullOrEmpty()) { @@ -41,6 +42,7 @@ internal fun reconcileActionsToKt(java: JavaEventActions, kt: KtEventActions) { java.skipSummarization().ifPresent { kt.skipSummarization = it } java.agentState().ifPresent { kt.agentState = agentStateFromJava(it) } java.compaction().ifPresent { kt.compaction = EventCompactionCodec.fromJava(it) } + java.rewindBeforeInvocationId().ifPresent { kt.rewindBeforeInvocationId = it } kt.endOfAgent = kt.endOfAgent || java.endOfAgent() // Merge deltas only when the tool replaced the actions (else they are the same live Kotlin map). if (java.stateDelta() !== kt.stateDelta) kt.stateDelta.putAll(java.stateDelta()) diff --git a/tokt/src/main/kotlin/com/google/adk/tokt/codecs/EventCodec.kt b/tokt/src/main/kotlin/com/google/adk/tokt/codecs/EventCodec.kt index d49c5c166..d064821ff 100644 --- a/tokt/src/main/kotlin/com/google/adk/tokt/codecs/EventCodec.kt +++ b/tokt/src/main/kotlin/com/google/adk/tokt/codecs/EventCodec.kt @@ -31,20 +31,15 @@ import kotlin.jvm.optionals.getOrNull * * Carries the identity (`id`, `invocationId`, `author`, `branch`), the [content][ContentCodec], the * [actions][KtEventActions] (state/artifact deltas, the control-flow flags - transfer, escalate, - * skipSummarization, endOfAgent - the requested tool confirmations, and the context-compaction - * summary), the streaming/error signals (`partial`, `turnComplete`, `interrupted`, `errorMessage`, - * `errorCode`, `finishReason`), the token `usageMetadata` ([UsageMetadataCodec]) and - * `groundingMetadata` ([GroundingMetadataCodec]), long-running tool ids, and the timestamp. + * skipSummarization, endOfAgent - the requested tool confirmations, the context-compaction summary, + * the resumable-run `agentState`, and the rewind marker `rewindBeforeInvocationId`), the + * streaming/error signals (`partial`, `turnComplete`, `interrupted`, `errorMessage`, `errorCode`, + * `finishReason`), the token `usageMetadata` ([UsageMetadataCodec]) and `groundingMetadata` + * ([GroundingMetadataCodec]), long-running tool ids, and the timestamp. */ internal object EventCodec { - /** - * Returns the Java [JavaEvent] view of the Kotlin [event]. - * - * `EventActions.agentState` crosses (see [agentStateToJava]), so a resumable run's state survives - * a round trip through a Java session service. `rewindBeforeInvocationId` is still dropped: ADK - * Java's `EventActions` has no such field. - */ + /** Returns the Java [JavaEvent] view of the Kotlin [event]. */ fun toJava(event: KtEvent): JavaEvent { val builder = JavaEvent.builder() @@ -98,6 +93,7 @@ internal object EventCodec { compaction = javaActions.compaction().getOrNull()?.let { EventCompactionCodec.fromJava(it) }, agentState = agentStateFromJava(javaActions.agentState().getOrNull()), + rewindBeforeInvocationId = javaActions.rewindBeforeInvocationId().getOrNull(), ), branch = event.branch().getOrNull(), partial = event.partial().getOrNull() ?: false, diff --git a/tokt/src/main/kotlin/com/google/adk/tokt/codecs/KtToJavaCodecs.kt b/tokt/src/main/kotlin/com/google/adk/tokt/codecs/KtToJavaCodecs.kt index 6384fc8c1..7b89ec409 100644 --- a/tokt/src/main/kotlin/com/google/adk/tokt/codecs/KtToJavaCodecs.kt +++ b/tokt/src/main/kotlin/com/google/adk/tokt/codecs/KtToJavaCodecs.kt @@ -25,6 +25,7 @@ import com.google.adk.kt.sessions.SessionKey as KtSessionKey import com.google.adk.kt.sessions.State as KtState import com.google.adk.sessions.Session as JavaSession import com.google.adk.sessions.State as JavaState +import java.util.AbstractList import java.util.AbstractMap import java.util.Optional import java.util.concurrent.ConcurrentMap @@ -73,6 +74,13 @@ internal class KtEventActionsToJavaView(private val actions: KtEventActions) : J actions.compaction = compaction?.let { EventCompactionCodec.fromJava(it) } } + override fun rewindBeforeInvocationId(): Optional = + Optional.ofNullable(actions.rewindBeforeInvocationId) + + override fun setRewindBeforeInvocationId(rewindBeforeInvocationId: String?) { + actions.rewindBeforeInvocationId = rewindBeforeInvocationId + } + // Overridden so the write reaches the Kotlin side; the base setters write private fields the // overridden getters never read. @@ -104,12 +112,10 @@ internal class KtEventActionsToJavaView(private val actions: KtEventActions) : J /** * Builds a read-only Java [JavaEventActions] snapshot from a Kotlin [KtEventActions], carrying the * deltas, the control-flow signals (skip-summarization, transfer, escalate, end-of-agent), the - * requested tool confirmations, and the resumable-run `agentState` ([agentStateToJava]). Used when - * exposing an existing Kotlin event to a Java runner ([EventCodec.toJava]); unlike - * [KtEventActionsToJavaView] (a live write sink) nothing writes back through it. - * - * Does not carry `rewindBeforeInvocationId`: ADK Java's `EventActions` has no such field, so a - * rewind request does not survive a round trip through a Java session service. + * requested tool confirmations, the compaction summary, the resumable-run `agentState` + * ([agentStateToJava]), and the rewind marker (`rewindBeforeInvocationId`). Used when exposing an + * existing Kotlin event to a Java runner ([EventCodec.toJava]); unlike [KtEventActionsToJavaView] + * (a live write sink), nothing writes back through it. */ internal fun eventActionsToJava(actions: KtEventActions): JavaEventActions = JavaEventActions.builder() @@ -124,6 +130,7 @@ internal fun eventActionsToJava(actions: KtEventActions): JavaEventActions = .endOfAgent(actions.endOfAgent) .compaction(actions.compaction?.let { EventCompactionCodec.toJava(it) }) .agentState(agentStateToJava(actions.agentState)) + .rewindBeforeInvocationId(actions.rewindBeforeInvocationId) .build() /** Copies a Kotlin state delta to Java, mapping the Kotlin [KtState.REMOVED] deletion sentinel. */ @@ -232,9 +239,9 @@ private class KtStateAsJavaConcurrentMap(private val state: KtState) : Concurren /** * Like [ktSessionToJava] but backs the session with live views: `events()` and `state()` convert on * access, so an adapted Java flow always reads the current session, including changes made earlier - * this turn. `events()` is read-only (the Kotlin runner owns the list) and re-converts each element - * on access, since a Kotlin event's actions are mutable; copy the list once rather than indexing it - * in a loop. `lastUpdateTime` is a snapshot, not live. + * this turn. `events()` allows appending but no other mutation, and it re-converts each element on + * access because a Kotlin event's actions are mutable; copy the list once instead of indexing it in + * a loop. `lastUpdateTime` is a snapshot, not live. */ // eventsView is deprecated for application code, but an interop adapter is its intended caller. @Suppress("DEPRECATION") @@ -243,18 +250,22 @@ internal fun ktSessionToJavaLive(session: KtSession): JavaSession = .appName(session.key.appName) .userId(session.key.userId) .state(JavaState(KtStateAsJavaConcurrentMap(session.state))) - .eventsView(KtBackedEventsView(session)) + .eventsView(KtBackedEventsMutableView(session)) .lastUpdateTime(session.lastUpdateTime.toJavaInstant()) .build() /** - * The read-only, converting `events()` of [ktSessionToJavaLive]. Named rather than anonymous so a - * caller holding a Java [JavaSession] can tell a live view from a snapshot: this one is already - * backed by the Kotlin session, so it must not be mirrored into (and would throw if tried). + * The converting `events()` list for [ktSessionToJavaLive]. Its `add` method appends to the Kotlin + * [session] so in-place updates from an ADK Java session service reach the running session, while + * `KtSessionServiceToJava.appendEvent` unwraps this view to append there directly. Mutations other + * than appending throw [UnsupportedOperationException]. */ -internal class KtBackedEventsView(private val session: KtSession) : AbstractList() { +internal class KtBackedEventsMutableView(internal val session: KtSession) : + AbstractList() { override val size: Int get() = session.events.size override fun get(index: Int): JavaEvent = EventCodec.toJava(session.events[index]) + + override fun add(element: JavaEvent): Boolean = session.events.add(EventCodec.fromJava(element)) } diff --git a/tokt/src/main/kotlin/com/google/adk/tokt/codecs/SessionCodec.kt b/tokt/src/main/kotlin/com/google/adk/tokt/codecs/SessionCodec.kt index bca1ee27e..f0cc5bcac 100644 --- a/tokt/src/main/kotlin/com/google/adk/tokt/codecs/SessionCodec.kt +++ b/tokt/src/main/kotlin/com/google/adk/tokt/codecs/SessionCodec.kt @@ -20,6 +20,7 @@ import com.google.adk.kt.sessions.Session as KtSession import com.google.adk.kt.sessions.SessionKey import com.google.adk.kt.sessions.State import com.google.adk.sessions.Session as JavaSession +import java.util.concurrent.CopyOnWriteArrayList import kotlin.time.toKotlinInstant /** Converts a Java session (key, state, and event history) to the Kotlin [KtSession]. */ @@ -30,7 +31,8 @@ internal object SessionCodec { key = SessionKey(appName = session.appName(), userId = session.userId(), id = session.id()), // Translate the Java removal sentinel so no foreign sentinel leaks into the Kotlin state. state = State(initialState = session.state().mapValues { stateValueFromJava(it.value) }), - events = session.events().map { EventCodec.fromJava(it) }.toMutableList(), + // Thread-safe like a Kotlin Session's default: parallel tool calls may append during reads. + events = CopyOnWriteArrayList(session.events().map { EventCodec.fromJava(it) }), lastUpdateTime = (session.lastUpdateTime() ?: java.time.Instant.EPOCH).toKotlinInstant(), ) } diff --git a/tokt/src/main/kotlin/com/google/adk/tokt/services/KtSessionServiceToJava.kt b/tokt/src/main/kotlin/com/google/adk/tokt/services/KtSessionServiceToJava.kt index 9aff981cc..613a72f9b 100644 --- a/tokt/src/main/kotlin/com/google/adk/tokt/services/KtSessionServiceToJava.kt +++ b/tokt/src/main/kotlin/com/google/adk/tokt/services/KtSessionServiceToJava.kt @@ -25,8 +25,9 @@ import com.google.adk.sessions.GetSessionConfig as JavaGetSessionConfig import com.google.adk.sessions.ListEventsResponse as JavaListEventsResponse import com.google.adk.sessions.ListSessionsResponse as JavaListSessionsResponse import com.google.adk.sessions.Session as JavaSession +import com.google.adk.sessions.State as JavaState import com.google.adk.tokt.codecs.EventCodec -import com.google.adk.tokt.codecs.KtBackedEventsView +import com.google.adk.tokt.codecs.KtBackedEventsMutableView import com.google.adk.tokt.codecs.SessionCodec import com.google.adk.tokt.codecs.ktSessionToJava import io.reactivex.rxjava3.core.Completable @@ -103,39 +104,43 @@ internal class KtSessionServiceToJava( } /** - * Appends [event] to [session] on the Kotlin service, then mirrors its stored session (merged - * state, events, and last-update time) back into the caller's Java [session] so ADK Java's - * `Runner` keeps observing the appended state in place. + * Appends [event] on the Kotlin service. For a live view ([KtBackedEventsMutableView]), the + * service appends directly to the running Kotlin session, keeping that session's `lastUpdateTime` + * in step with the store; any other [session] is updated in place by + * [JavaBaseSessionService.appendEvent], plus the `temp:` state it skips. In both cases, + * [session]'s `lastUpdateTime` is refreshed from the Kotlin session the service updated. */ override fun appendEvent(session: JavaSession, event: JavaEvent): Single = rxSingle(dispatcher) { - val key = SessionKey(session.appName(), session.userId(), session.id()) - service.appendEvent(SessionCodec.fromJava(session), EventCodec.fromJava(event)) - service.getSession(key)?.let { stored -> - // Convert before touching the caller's session, then refill under the list's monitor so a - // concurrent reader never observes a transiently-empty event list. Session.events() is a - // Collections.synchronizedList, so its own monitor is the correct lock to hold here. - // Mirror only into a session we can actually write. The live view this module hands out - // (ktSessionToJavaLive, e.g. invocationContext.session()) is already backed by the Kotlin - // session: its events() is a read-only converting view that throws on clear/addAll, and - // its state() writes straight through, so it needs no mirroring and must not be mutated - // here. A plain snapshot session (ktSessionToJava) does. - if (session.events() is MutableList<*> && session.events() !is KtBackedEventsView) { - val storedEvents = stored.events.map { EventCodec.toJava(it) } - synchronized(session.events()) { - session.events().clear() - session.events().addAll(storedEvents) - } - // No lock: State is a ConcurrentMap, so readers never take this monitor. Drop stale keys - // and overwrite in place, leaving no transiently-empty window. - session.state().keys.retainAll(stored.state.keys) - session.state().putAll(stored.state) - } - session.lastUpdateTime(stored.lastUpdateTime.toJavaInstant()) + // ADK base session services ignore partial events, so the in-place update must too. + if (event.partial().getOrNull() == true) return@rxSingle event + val backing = (session.events() as? KtBackedEventsMutableView)?.session + val ktSession = backing ?: SessionCodec.fromJava(session) + service.appendEvent(ktSession, EventCodec.fromJava(event)) + if (backing == null) { + super.appendEvent(session, event) + applyTempState(session, event) } + session.lastUpdateTime(ktSession.lastUpdateTime.toJavaInstant()) event } + /** + * Applies [event]'s `temp:` state to [session] the way Kotlin's `State.applyTempDelta` does for a + * Kotlin session, since the Java base skips it. + */ + private fun applyTempState(session: JavaSession, event: JavaEvent) { + val delta = event.actions()?.stateDelta() ?: return + for ((key, value) in delta) { + if (!key.startsWith(JavaState.TEMP_PREFIX)) continue + if (value === JavaState.REMOVED) { + session.state().remove(key) + } else { + session.state()[key] = value + } + } + } + private fun JavaGetSessionConfig.toKotlin(): KtGetSessionConfig = KtGetSessionConfig( numRecentEvents = numRecentEvents().getOrNull(), diff --git a/tokt/src/test/kotlin/com/google/adk/tokt/KtRunnerInteropTest.kt b/tokt/src/test/kotlin/com/google/adk/tokt/KtRunnerInteropTest.kt index b1861703e..7729f9308 100644 --- a/tokt/src/test/kotlin/com/google/adk/tokt/KtRunnerInteropTest.kt +++ b/tokt/src/test/kotlin/com/google/adk/tokt/KtRunnerInteropTest.kt @@ -29,6 +29,7 @@ import com.google.adk.events.EventCompaction as JavaEventCompaction import com.google.adk.kt.agents.BaseAgent as KtBaseAgent import com.google.adk.kt.agents.InvocationContext as KtInvocationContext import com.google.adk.kt.agents.LlmAgent as KtLlmAgent +import com.google.adk.kt.agents.ResumabilityConfig as KtResumabilityConfig import com.google.adk.kt.agents.RunConfig as KtRunConfig import com.google.adk.kt.agents.StreamingMode as KtStreamingMode import com.google.adk.kt.agents.TypedData @@ -89,11 +90,13 @@ import com.google.adk.sessions.GetSessionConfig as JavaGetSessionConfig import com.google.adk.sessions.InMemorySessionService as JavaInMemorySessionService import com.google.adk.sessions.ListEventsResponse as JavaListEventsResponse import com.google.adk.sessions.Session as JavaSession +import com.google.adk.sessions.SessionKey as JavaSessionKey import com.google.adk.sessions.State as JavaState import com.google.adk.tokt.adapters.reconcileActionsToKt import com.google.adk.tokt.codecs.EventCodec import com.google.adk.tokt.codecs.FunctionDeclarationCodec import com.google.adk.tokt.codecs.GroundingMetadataCodec +import com.google.adk.tokt.codecs.KtBackedEventsMutableView import com.google.adk.tokt.codecs.KtEventActionsToJavaView import com.google.adk.tokt.codecs.PartCodec import com.google.adk.tokt.codecs.RunConfigCodec @@ -113,6 +116,7 @@ import com.google.genai.types.ExecutableCode as GenaiExecutableCode import com.google.genai.types.FinishReason as GenaiFinishReason import com.google.genai.types.FunctionCall as GenaiFunctionCall import com.google.genai.types.FunctionDeclaration as GenaiFunctionDeclaration +import com.google.genai.types.FunctionResponse as GenaiFunctionResponse import com.google.genai.types.GenerateContentResponseUsageMetadata as GenaiUsageMetadata import com.google.genai.types.GroundingChunk as GenaiGroundingChunk import com.google.genai.types.GroundingChunkMaps as GenaiGroundingChunkMaps @@ -136,6 +140,7 @@ import io.reactivex.rxjava3.core.Completable import io.reactivex.rxjava3.core.Flowable import io.reactivex.rxjava3.core.Maybe import io.reactivex.rxjava3.core.Single +import java.time.Instant import java.util.Optional import java.util.concurrent.ConcurrentHashMap import java.util.concurrent.ConcurrentMap @@ -145,6 +150,7 @@ import java.util.concurrent.atomic.AtomicInteger import java.util.concurrent.atomic.AtomicReference import kotlin.jvm.optionals.getOrNull import kotlin.test.Test +import kotlin.test.assertContains import kotlin.test.assertContentEquals import kotlin.test.assertEquals import kotlin.test.assertFailsWith @@ -278,6 +284,19 @@ class KtRunnerInteropTest { } } + /** A long-running Java tool that returns only a pending status. */ + private class JavaPendingTool : + JavaBaseTool("java_pending", "long-running", /* isLongRunning= */ true) { + override fun declaration(): Optional = + Optional.of(GenaiFunctionDeclaration.builder().name("java_pending").build()) + + @JvmSuppressWildcards + override fun runAsync( + args: Map, + toolContext: JavaToolContext, + ): Single> = Single.just(mapOf("status" to "pending")) + } + /** * A Java tool that mutates its [JavaToolContext.actions] in place (not via `setActions`) to set * every control-flow signal, exercising the live [KtEventActionsToJavaView] write-through. @@ -701,6 +720,29 @@ class KtRunnerInteropTest { } } + /** A Kotlin runner that records each runAsync call once its stream is collected. */ + private class RecordingKtRunner(private val delegate: KtRunner) : KtRunner by delegate { + data class RunAsyncCall( + val userId: String, + val sessionId: String, + val invocationId: String?, + val newMessage: KtContent?, + ) + + val runAsyncCalls = mutableListOf() + + override fun runAsync( + userId: String, + sessionId: String, + invocationId: String?, + newMessage: KtContent?, + stateDelta: Map?, + runConfig: KtRunConfig?, + ): Flow = flow { + runAsyncCalls.add(RunAsyncCall(userId, sessionId, invocationId, newMessage)) + } + } + @Test fun javaAdkToKt_convertsEntireCollections() { // Tools: order preserved, each Java tool wrapped as a Kotlin tool. @@ -881,6 +923,42 @@ class KtRunnerInteropTest { ) } + @Test + fun ktRunner_rewindThroughAnAdaptedJavaSessionService_dropsTheRewoundTurnFromTheNextPrompt() = + runBlocking { + // Without the marker surviving the Java store, the engine keeps sending the rewound turn. + val javaSessions = JavaInMemorySessionService() + val model = SequentialJavaModel(listOf(modelText("ok"))) + val runner = + KtInMemoryRunner( + app = + KtApp( + appName = "app", + rootAgent = KtLlmAgent(name = "a", model = JavaAdkToKt.asKtModel(model)), + ), + sessionService = JavaAdkToKt.asKtSessionService(javaSessions), + ) + runner.turn("keep me") + val rewound = assertNotNull(runner.turn("forget me").first().invocationId) + + runner.rewindAsync(userId = "u", sessionId = "s", rewindBeforeInvocationId = rewound) + runner.turn("after the rewind") + + val stored = + javaSessions.getSession("app", "u", "s", Optional.empty()).blockingGet() + ?: fail("the adapted Java service should hold the session") + assertTrue( + stored.events().any { it.actions().rewindBeforeInvocationId().getOrNull() == rewound }, + "the stored rewind event should keep its marker", + ) + val lastPrompt = + model.requests.last().contents().flatMap { content -> + content.parts().getOrNull().orEmpty().mapNotNull { it.text().getOrNull() } + } + assertContains(lastPrompt, "keep me", "the turn before the rewind should stay in the prompt") + assertFalse("forget me" in lastPrompt, "the rewound turn must not reach the model") + } + @Test fun ktRunner_pluginBeforeToolCallback_firesForNativeKotlinTool_andCanDenyIt() = runBlocking { // A native Kotlin tool has no Java form, yet a Java plugin's beforeToolCallback must still @@ -1159,6 +1237,307 @@ class KtRunnerInteropTest { ) } + @Test + fun asJavaRunner_resumeOverload_resumesThePausedInvocation() { + // A resumable app pauses on a confirmation; the resume overload continues that invocation. + val ktRunner = + KtInMemoryRunner( + app = + KtApp( + appName = "app", + rootAgent = + KtLlmAgent( + name = "a", + model = + JavaAdkToKt.asKtModel( + SequentialJavaModel( + listOf( + modelFunctionCall("java_confirm", emptyMap()), + modelText("confirmed done"), + ) + ) + ), + tools = listOf(JavaAdkToKt.asKtTool(JavaConfirmTool())), + ), + resumabilityConfig = KtResumabilityConfig(isResumable = true), + ) + ) + val javaRunner = KotlinAdkToJava.asJavaRunner(ktRunner) + val autoCreate = JavaRunConfig.builder().autoCreateSession(true).build() + + val turn1 = + javaRunner + .runAsync( + "u", + "s", + GenaiContent.builder().role("user").parts(GenaiPart.fromText("go")).build(), + autoCreate, + ) + .toList() + .blockingGet() + val (pausedEvent, confirmationCall) = + turn1 + .flatMap { event -> event.functionCalls().map { event to it } } + .first { (_, call) -> call.name().getOrNull() == "adk_request_confirmation" } + val pausedInvocationId = pausedEvent.invocationId() + val approval = + GenaiContent.builder() + .role("user") + .parts( + GenaiPart.builder() + .functionResponse( + GenaiFunctionResponse.builder() + .name("adk_request_confirmation") + .id(confirmationCall.id().get()) + .response(mapOf("confirmed" to true)) + .build() + ) + .build() + ) + .build() + + // The engine infers the paused invocation from the approval's function-response id. + val turn2 = + javaRunner + .runAsync("u", "s", /* invocationId= */ null, approval, autoCreate, /* stateDelta= */ null) + .toList() + .blockingGet() + + val (resultEvent, toolResponse) = + turn2 + .flatMap { event -> event.functionResponses().map { event to it } } + .first { (_, response) -> response.name().getOrNull() == "java_confirm" } + assertEquals( + "confirmed", + toolResponse.response().getOrNull()?.get("status"), + "the resumed run should deliver the approval to the tool", + ) + assertEquals( + pausedInvocationId, + resultEvent.invocationId(), + "turn 2 should continue the paused invocation, not start a new one", + ) + + // Resuming the finished invocation by id alone is a no-op rather than an error. + val turn3 = + javaRunner + .runAsync( + "u", + "s", + pausedInvocationId, + /* newMessage= */ null, + autoCreate, + /* stateDelta= */ null, + ) + .toList() + .blockingGet() + assertTrue(turn3.isEmpty(), "resuming a finished invocation should emit nothing") + } + + @Test + fun asJavaRunner_resumeOverload_roleLessResponse_reachesTheModelAsAUserTurn() { + val model = + SequentialJavaModel(listOf(modelFunctionCall("java_pending", emptyMap()), modelText("done"))) + val ktRunner = + KtInMemoryRunner( + app = + KtApp( + appName = "app", + rootAgent = + KtLlmAgent( + name = "a", + model = JavaAdkToKt.asKtModel(model), + tools = listOf(JavaAdkToKt.asKtTool(JavaPendingTool())), + ), + resumabilityConfig = KtResumabilityConfig(isResumable = true), + ) + ) + val javaRunner = KotlinAdkToJava.asJavaRunner(ktRunner) + val autoCreate = JavaRunConfig.builder().autoCreateSession(true).build() + val turn1 = + javaRunner + .runAsync( + "u", + "s", + GenaiContent.builder().role("user").parts(GenaiPart.fromText("go")).build(), + autoCreate, + ) + .toList() + .blockingGet() + val call = + turn1.flatMap { it.functionCalls() }.first { it.name().getOrNull() == "java_pending" } + val roleLessResponse = + GenaiContent.builder() + .parts( + GenaiPart.builder() + .functionResponse( + GenaiFunctionResponse.builder() + .name("java_pending") + .id(call.id().get()) + .response(mapOf("status" to "done")) + .build() + ) + .build() + ) + .build() + + val unusedResumeEvents = + javaRunner + .runAsync( + "u", + "s", + turn1.first().invocationId(), + roleLessResponse, + autoCreate, + /* stateDelta= */ null, + ) + .toList() + .blockingGet() + + val sentResponses = + model.requests + .last() + .contents() + .filter { it.role().getOrNull() == "user" } + .flatMap { it.parts().getOrNull().orEmpty() } + .mapNotNull { it.functionResponse().getOrNull()?.response()?.getOrNull() } + assertTrue( + mapOf("status" to "done") in sentResponses, + "the role-less response should reach the model as a user turn", + ) + } + + @Test + fun asJavaRunner_resumeOverload_nonResumableApp_startsNewInvocationWithCallerId() { + // The Java Runner throws for a non-resumable app; the Kotlin engine, like Python, runs it. + val javaRunner = + KotlinAdkToJava.asJavaRunner( + KtInMemoryRunner( + agent = + KtLlmAgent( + name = "a", + model = JavaAdkToKt.asKtModel(SequentialJavaModel(listOf(modelText("done")))), + ), + appName = "app", + ) + ) + + val events = + javaRunner + .runAsync( + "u", + "s", + "caller-id", + GenaiContent.builder().role("user").parts(GenaiPart.fromText("hi")).build(), + JavaRunConfig.builder().autoCreateSession(true).build(), + /* stateDelta= */ null, + ) + .toList() + .blockingGet() + + assertTrue(events.isNotEmpty(), "the run should start a new invocation instead of throwing") + assertTrue( + events.all { it.invocationId() == "caller-id" }, + "the new invocation should use the caller's invocation id", + ) + } + + @Test + fun asJavaRunner_resumeOverload_nonResumableApp_nullMessage_throws() { + val javaRunner = + KotlinAdkToJava.asJavaRunner( + KtInMemoryRunner( + agent = + KtLlmAgent(name = "a", model = JavaAdkToKt.asKtModel(SequentialJavaModel(emptyList()))), + appName = "app", + ) + ) + + val error = + assertFailsWith { + javaRunner + .runAsync( + "u", + "s", + "caller-id", + /* newMessage= */ null, + JavaRunConfig.builder().autoCreateSession(true).build(), + /* stateDelta= */ null, + ) + .toList() + .blockingGet() + } + assertContains(error.message.orEmpty(), "not resumable") + } + + @Test + fun asJavaRunner_everyRunAsyncOverload_routesToKotlinEngine() { + // Every public runAsync overload must reach the Kotlin engine, not the base Runner. + val recording = + RecordingKtRunner( + KtInMemoryRunner( + agent = + KtLlmAgent(name = "a", model = JavaAdkToKt.asKtModel(SequentialJavaModel(emptyList()))), + appName = "app", + ) + ) + val javaRunner = KotlinAdkToJava.asJavaRunner(recording) + val key = JavaSessionKey("app", "u", "s") + // Pre-create the session for overloads whose default RunConfig disables autoCreateSession. + javaRunner.sessionService().createSession(key).ignoreElement().blockingAwait() + + val msg = GenaiContent.builder().role("user").parts(GenaiPart.fromText("hi")).build() + val autoCreate = JavaRunConfig.builder().autoCreateSession(true).build() + + javaRunner.runAsync("u", "s", msg, autoCreate).ignoreElements().blockingAwait() + javaRunner + .runAsync("u", "s", msg, autoCreate, /* stateDelta= */ null) + .ignoreElements() + .blockingAwait() + javaRunner.runAsync("u", "s", msg).ignoreElements().blockingAwait() + javaRunner + .runAsync(key, msg, autoCreate, /* stateDelta= */ null) + .ignoreElements() + .blockingAwait() + javaRunner.runAsync(key, msg, autoCreate).ignoreElements().blockingAwait() + javaRunner.runAsync(key, msg).ignoreElements().blockingAwait() + javaRunner + .runAsync("u", "s", "invocation-xyz", msg, autoCreate, /* stateDelta= */ null) + .ignoreElements() + .blockingAwait() + javaRunner + .runAsync( + "u", + "s", + "invocation-xyz", + /* newMessage= */ null, + autoCreate, + /* stateDelta= */ null, + ) + .ignoreElements() + .blockingAwait() + + assertEquals( + 8, + recording.runAsyncCalls.size, + "each runAsync call should reach the Kotlin engine exactly once", + ) + assertTrue( + recording.runAsyncCalls.all { it.userId == "u" && it.sessionId == "s" }, + "every overload should pass userId and sessionId through", + ) + assertEquals( + listOf(null, null, null, null, null, null, "invocation-xyz", "invocation-xyz"), + recording.runAsyncCalls.map { it.invocationId }, + "only the resume overload should carry an invocationId, unchanged", + ) + assertEquals( + listOf(true, true, true, true, true, true, true, false), + recording.runAsyncCalls.map { it.newMessage != null }, + "a null message on the resume overload should reach the engine as null", + ) + } + @Test fun asJavaRunner_nonDefaultRunConfig_reachesTheKotlinEngine() { // A Java RunConfig set through the Java API must map onto the Kotlin engine; the plugin reads @@ -1209,25 +1588,28 @@ class KtRunnerInteropTest { appName = "app", ) val javaRunner = KotlinAdkToJava.asJavaRunner(ktRunner) + // Checks the adapter's own error, not the one the inspection-only agent view would raise. + fun assertLiveUnsupported(run: () -> Flowable) { + val error = assertFailsWith { run().toList().blockingGet() } + assertEquals("Live mode is not supported when running on the Kotlin engine.", error.message) + } // Live mode is not bridged; it must fail loudly through the stream (like the base Runner), not // throw eagerly at the call site. - assertFailsWith { - javaRunner - .runLive("u", "s", LiveRequestQueue(), JavaRunConfig.builder().build()) - .toList() - .blockingGet() + assertLiveUnsupported { + javaRunner.runLive("u", "s", LiveRequestQueue(), JavaRunConfig.builder().build()) } // agent() returns an inspection-only view of the Kotlin agent - readable, but not runnable. assertEquals("a", javaRunner.agent().name(), "agent() should expose the Kotlin agent's name") - // The session-based runLive overload must fail the same way, through the stream. - val session = javaRunner.sessionService().createSession("app", "u", null, "s").blockingGet() - assertFailsWith { - javaRunner - .runLive(session, LiveRequestQueue(), JavaRunConfig.builder().build()) - .toList() - .blockingGet() + // The session-based and SessionKey runLive overloads must fail the same way. + val key = JavaSessionKey("app", "u", "s") + val session = javaRunner.sessionService().createSession(key).blockingGet() + assertLiveUnsupported { + javaRunner.runLive(session, LiveRequestQueue(), JavaRunConfig.builder().build()) + } + assertLiveUnsupported { + javaRunner.runLive(key, LiveRequestQueue(), JavaRunConfig.builder().build()) } } @@ -2873,8 +3255,7 @@ class KtRunnerInteropTest { observed["keys"] = artifacts.listArtifactKeys("app", "u", "s").blockingGet().filenames() observed["versions"] = artifacts.listVersions("app", "u", "s", "note.txt").blockingGet() - // KtSessionServiceToJava: getSession, then appendEvent, whose store-mirroring is the - // most intricate method in the module. + // KtSessionServiceToJava: getSession, then appendEvent on that snapshot. val javaSession = sessions.getSession("app", "u", "s", Optional.empty()).blockingGet()!! val eventsBefore = javaSession.events().size val appended = @@ -2890,7 +3271,7 @@ class KtRunnerInteropTest { ) .blockingGet() observed["appended_author"] = appended.author() - // appendEvent mirrors the Kotlin store back into this very Java session object. + // appendEvent mutates the caller's Java session in place. observed["events_grew"] = javaSession.events().size > eventsBefore observed["sessions"] = sessions.listSessions("app", "u").blockingGet().sessions().size @@ -2922,15 +3303,400 @@ class KtRunnerInteropTest { assertEquals("v1", observed["loaded"], "load must return what save stored") assertEquals(listOf("note.txt"), observed["keys"], "listArtifactKeys") assertEquals(listOf(0), observed["versions"], "listVersions") - assertEquals( - true, - observed["events_grew"], - "appendEvent must mirror the store into the session", - ) + assertEquals(true, observed["events_grew"], "appendEvent must update the session in place") assertEquals(1, observed["sessions"], "listSessions") assertEquals(1, observed["memories"], "searchMemory should hit the appended text") } + /** + * Java tool that appends a note through `ic.sessionService()` with `ic.session()` and records + * what that live session reflects immediately after the append. + */ + private class AppendNoteJavaTool(name: String = "append_note") : + JavaBaseTool(name, "appends a note through the live session") { + val observed = ConcurrentHashMap() + + override fun declaration(): Optional = + Optional.of(GenaiFunctionDeclaration.builder().name(name()).build()) + + @JvmSuppressWildcards + override fun runAsync( + args: Map, + toolContext: JavaToolContext, + ): Single> { + val ic = toolContext.invocationContext() + val live = ic.session() + val before = live.events().size + // Pick a time after anything in the session, so a lastUpdateTime refresh is observable. + val eventTime = live.lastUpdateTime().plusSeconds(1) + val unused = + ic + .sessionService() + .appendEvent( + live, + JavaEvent.builder() + .id(JavaEvent.generateEventId()) + .invocationId(ic.invocationId()) + .author(name()) + .timestamp(eventTime.toEpochMilli()) + .content(GenaiContent.fromParts(GenaiPart.fromText("appended note"))) + .actions( + JavaEventActions.builder() + .stateDelta(mutableMapOf("note" to "kept", "temp:scratch" to "t")) + .build() + ) + .build(), + ) + .blockingGet() + observed["grew_by"] = live.events().size - before + observed["note"] = live.state()["note"] ?: "MISSING" + observed["temp"] = live.state()["temp:scratch"] ?: "MISSING" + observed["event_time"] = eventTime + observed["last_update_time"] = live.lastUpdateTime() + return Single.just(mapOf("ok" to true)) + } + } + + private fun appendNoteAgent(tool: AppendNoteJavaTool) = + KtLlmAgent( + name = "a", + model = + JavaAdkToKt.asKtModel( + SequentialJavaModel(listOf(modelFunctionCall(tool.name(), emptyMap()), modelText("done"))) + ), + tools = listOf(JavaAdkToKt.asKtTool(tool)), + ) + + @Test + fun ktRunner_javaToolAppendsThroughItsLiveSession_invocationSeesTheEvent() = runBlocking { + // The append must reach the running session itself, not a converted copy of it. + val tool = AppendNoteJavaTool() + val runner = KtInMemoryRunner(appendNoteAgent(tool), appName = "app") + + runner.turn() + + assertEquals(1, tool.observed["grew_by"], "the live session's events() must include the event") + assertEquals("kept", tool.observed["note"], "the live session's state() must include its delta") + assertEquals( + tool.observed["event_time"], + tool.observed["last_update_time"], + "the live session's lastUpdateTime must follow the append", + ) + // Older Kotlin releases drop temp: keys, so expect what a Kotlin caller gets. + val kotlinCaller = runner.sessionService.createSession(KtSessionKey("app", "u", "kotlin")) + val unused = + runner.sessionService.appendEvent( + kotlinCaller, + KtEvent( + author = "a", + actions = KtEventActions(stateDelta = mutableMapOf("temp:scratch" to "t")), + ), + ) + assertEquals( + kotlinCaller.state["temp:scratch"] ?: "MISSING", + tool.observed["temp"], + "the live session's state() must match what the Kotlin service gives a Kotlin caller", + ) + val stored = runner.sessionService.getSession(KtSessionKey("app", "u", "s"))!! + assertEquals(1, stored.events.count { it.author == "append_note" }, "stored exactly once") + } + + @Test + fun ktRunner_onAJavaSessionService_liveSessionAppend_reachesRunningSessionAndStore() = + runBlocking { + // A Java session service appends to the live view in place, so the view must support add. + val javaSessions = JavaInMemorySessionService() + val tool = AppendNoteJavaTool() + val runner = + KtInMemoryRunner( + appendNoteAgent(tool), + appName = "app", + sessionService = JavaAdkToKt.asKtSessionService(javaSessions), + ) + + runner.turn() + + assertEquals( + 1, + tool.observed["grew_by"], + "the live session's events() must include the event", + ) + assertEquals( + "kept", + tool.observed["note"], + "the live session's state() must include its delta", + ) + val stored = javaSessions.getSession("app", "u", "s", Optional.empty()).blockingGet()!! + assertEquals(1, stored.events().count { it.author() == "append_note" }, "stored exactly once") + assertEquals("kept", stored.state()["note"], "the delta is persisted") + } + + @Test + fun ktRunner_staleCheckingKtService_javaToolLiveAppend_runnerAppendsStillAccepted() = + runBlocking { + // Appending through the view alone would leave the running session's lastUpdateTime stale. + val delegate = KtInMemorySessionService() + val rejected = AtomicInteger() + val staleChecking = + object : KtSessionService by delegate { + override suspend fun appendEvent(session: KtSession, event: KtEvent): KtEvent { + if (!event.partial) { + val stored = delegate.getSession(session.key) + if (stored != null && stored.lastUpdateTime != session.lastUpdateTime) { + rejected.incrementAndGet() + error("stale session") + } + } + return delegate.appendEvent(session, event) + } + } + val tool = AppendNoteJavaTool() + val runner = + KtInMemoryRunner(appendNoteAgent(tool), appName = "app", sessionService = staleChecking) + + runner.turn() + + assertEquals(0, rejected.get(), "no append may be rejected as stale") + assertEquals( + 1, + tool.observed["grew_by"], + "the live session's events() must include the event", + ) + val stored = delegate.getSession(KtSessionKey("app", "u", "s"))!! + assertEquals(1, stored.events.count { it.author == "append_note" }, "stored exactly once") + } + + /** + * Runs two turns of parallel and repeated Java tool appends on [sessionService], then asserts + * that each append lands exactly once in the live session and the store, with no doubled events. + */ + private suspend fun assertJavaAppendsLandOnce(sessionService: KtSessionService) { + data class EventIdAndAuthor(val id: String, val author: String) + val liveEvents = CopyOnWriteArrayList>() + val probe = + object : JavaBaseTool("probe", "records the live session's events") { + override fun declaration(): Optional = + Optional.of(GenaiFunctionDeclaration.builder().name("probe").build()) + + @JvmSuppressWildcards + override fun runAsync( + args: Map, + toolContext: JavaToolContext, + ): Single> { + liveEvents.add( + toolContext.invocationContext().session().events().map { + EventIdAndAuthor(it.id(), it.author()) + } + ) + return Single.just(mapOf("ok" to true)) + } + } + val parallelNotes = + GenaiContent.builder() + .role("model") + .parts( + GenaiPart.fromFunctionCall("note_a", emptyMap()), + GenaiPart.fromFunctionCall("note_b", emptyMap()), + ) + .build() + val agent = + KtLlmAgent( + name = "a", + model = + JavaAdkToKt.asKtModel( + SequentialJavaModel( + listOf( + parallelNotes, + modelFunctionCall("note_a", emptyMap()), + modelFunctionCall("probe", emptyMap()), + modelText("done"), + modelFunctionCall("note_b", emptyMap()), + modelFunctionCall("probe", emptyMap()), + modelText("done"), + ) + ) + ), + tools = + JavaAdkToKt.asKtTools( + listOf(AppendNoteJavaTool("note_a"), AppendNoteJavaTool("note_b"), probe) + ), + ) + val runner = KtInMemoryRunner(agent, appName = "app", sessionService = sessionService) + + runner.turn() + runner.turn() + + fun assertEachNoteOnce(where: String, events: List, notes: Map) { + assertEquals(events.size, events.map { it.id }.toSet().size, "$where: no event is doubled") + assertEquals( + notes, + events.map { it.author }.filter { it.startsWith("note_") }.groupingBy { it }.eachCount(), + "$where: one note per append", + ) + } + assertEachNoteOnce("turn 1 live session", liveEvents[0], mapOf("note_a" to 2, "note_b" to 1)) + assertEachNoteOnce("turn 2 live session", liveEvents[1], mapOf("note_a" to 2, "note_b" to 2)) + val stored = sessionService.getSession(KtSessionKey("app", "u", "s"))!! + assertEachNoteOnce( + "store", + stored.events.map { EventIdAndAuthor(it.id, it.author) }, + mapOf("note_a" to 2, "note_b" to 2), + ) + } + + @Test + fun ktRunner_parallelAndRepeatedJavaAppends_eachLandOnce() = runBlocking { + // Parallel tool calls append through one shared live session; none may be lost or doubled. + assertJavaAppendsLandOnce(KtInMemorySessionService()) + } + + @Test + fun ktRunner_onAJavaSessionService_parallelAndRepeatedJavaAppends_eachLandOnce() = runBlocking { + // Tools append through the view's add and the runner through a copy; neither may double. + assertJavaAppendsLandOnce(JavaAdkToKt.asKtSessionService(JavaInMemorySessionService())) + } + + @Test + fun ktRunner_onADecoratedJavaSessionService_parallelAndRepeatedJavaAppends_eachLandOnce() = + runBlocking { + // A decorator blocks unwrapping the adapter, so tool appends cross both adapters. + val adapted = JavaAdkToKt.asKtSessionService(JavaInMemorySessionService()) + assertJavaAppendsLandOnce(object : KtSessionService by adapted {}) + } + + @Test + fun liveEventsView_add_appendsToKotlinSession() { + val session = KtSession(key = KtSessionKey("app", "u", "s")) + val event = JavaEvent.builder().id(JavaEvent.generateEventId()).author("a").build() + + assertTrue(KtBackedEventsMutableView(session).add(event), "add must report the change") + + assertEquals(listOf(event.id()), session.events.map { it.id }, "add must reach the session") + } + + @Test + fun liveEventsView_otherMutations_throwUnsupportedOperation() { + val session = KtSession(key = KtSessionKey("app", "u", "s")) + val event = JavaEvent.builder().id(JavaEvent.generateEventId()).author("a").build() + session.events.add(EventCodec.fromJava(event)) + val events = KtBackedEventsMutableView(session) + + assertFailsWith { events[0] = event } + assertFailsWith { events.removeAt(0) } + assertFailsWith { events.add(0, event) } + } + + @Test + fun sessionCodec_fromJava_eventListToleratesAppendDuringIteration() { + val javaEvent = JavaEvent.builder().id(JavaEvent.generateEventId()).author("a").build() + val events = + SessionCodec.fromJava( + JavaSession.builder("s").appName("app").userId("u").events(listOf(javaEvent)).build() + ) + .events + val iterator = events.iterator() + + events.add(KtEvent(author = "b")) + + assertEquals(javaEvent.id(), iterator.next().id, "an append must not break an open iterator") + } + + /** Returns a runner after one text-only turn so session "s" exists in its store. */ + private suspend fun ktRunnerAfterOneTurn(): KtInMemoryRunner { + val runner = + KtInMemoryRunner( + KtLlmAgent( + name = "a", + model = JavaAdkToKt.asKtModel(SequentialJavaModel(listOf(modelText("done")))), + ), + appName = "app", + ) + runner.turn() + return runner + } + + @Test + fun asJavaRunner_appendEventOnASnapshot_updatesItInPlace() = runBlocking { + // Snapshot appends follow ADK Java's in-place contract, window included, plus temp: state. + val ktRunner = ktRunnerAfterOneTurn() + val sessions = KotlinAdkToJava.asJavaRunner(ktRunner).sessionService() + val windowed = + sessions + .getSession( + "app", + "u", + "s", + Optional.of(JavaGetSessionConfig.builder().numRecentEvents(1).build()), + ) + .blockingGet()!! + assertEquals(1, windowed.events().size, "precondition: a one-event window") + windowed.state()["local_only"] = "kept" + windowed.state()["temp:stale"] = "x" + val event = + JavaEvent.builder() + .id(JavaEvent.generateEventId()) + .invocationId("app-code") + .author("app_code") + .timestamp(windowed.lastUpdateTime().plusSeconds(1).toEpochMilli()) + .content(GenaiContent.fromParts(GenaiPart.fromText("from app code"))) + .actions( + JavaEventActions.builder() + .stateDelta( + mutableMapOf( + "color" to "teal", + "temp:scratch" to "t", + "temp:stale" to JavaState.REMOVED, + ) + ) + .build() + ) + .build() + + val unused = sessions.appendEvent(windowed, event).blockingGet() + + assertEquals(2, windowed.events().size, "the window plus the appended event") + assertSame(event, windowed.events().last(), "the caller's own event instance is appended") + assertEquals("teal", windowed.state()["color"], "the delta is applied in place") + assertEquals("t", windowed.state()["temp:scratch"], "temp: state is applied in place too") + assertNull(windowed.state()["temp:stale"], "temp: removals are applied in place") + assertEquals("kept", windowed.state()["local_only"], "unstored state keys survive the append") + val stored = ktRunner.sessionService.getSession(KtSessionKey("app", "u", "s"))!! + assertEquals(3, stored.events.size, "the event is persisted once") + assertEquals("teal", stored.state["color"], "the delta is persisted") + assertNull(stored.state["local_only"], "only the event's delta is persisted") + assertEquals( + Instant.ofEpochMilli(event.timestamp()), + windowed.lastUpdateTime(), + "lastUpdateTime follows the append", + ) + } + + @Test + fun asJavaRunner_appendPartialEventOnASnapshot_leavesItUnchanged() = runBlocking { + val ktRunner = ktRunnerAfterOneTurn() + val sessions = KotlinAdkToJava.asJavaRunner(ktRunner).sessionService() + val snapshot = sessions.getSession("app", "u", "s", Optional.empty()).blockingGet()!! + val eventsBefore = snapshot.events().size + val partial = + JavaEvent.builder() + .id(JavaEvent.generateEventId()) + .author("app_code") + .partial(true) + .actions( + JavaEventActions.builder() + .stateDelta(mutableMapOf("temp:scratch" to "t")) + .build() + ) + .build() + + val unused = sessions.appendEvent(snapshot, partial).blockingGet() + + assertEquals(eventsBefore, snapshot.events().size, "a partial event is not appended") + assertNull(snapshot.state()["temp:scratch"], "a partial event's temp: state is not applied") + val stored = ktRunner.sessionService.getSession(KtSessionKey("app", "u", "s"))!! + assertEquals(eventsBefore, stored.events.size, "a partial event is not stored") + } + @Test fun ktRunner_javaToolReadsPluginManager_throughItsInvocationContext() = runBlocking { // A Java tool reads toolContext.invocationContext().pluginManager(): the one-way bridge exposes @@ -3147,6 +3913,19 @@ class KtRunnerInteropTest { ) } + @Test + fun eventCodec_rewindBeforeInvocationId_roundTripsThroughTheJavaEvent() { + // Wiring check: without both directions, a Java session service drops the rewind marker. + val ktEvent = + KtEvent(author = "user", actions = KtEventActions(rewindBeforeInvocationId = "inv1")) + + assertEquals( + "inv1", + EventCodec.fromJava(EventCodec.toJava(ktEvent)).actions.rewindBeforeInvocationId, + "rewindBeforeInvocationId must survive EventCodec.toJava -> fromJava", + ) + } + @Test fun schemaCodec_carriesEveryFacet_roundTrips() { // Every facet the Kotlin Schema models must survive Java -> Kotlin -> Java, or structured @@ -3573,6 +4352,31 @@ class KtRunnerInteropTest { assertTrue(kt.endOfAgent, "setEndInvocation should reach the Kotlin actions") } + @Test + fun ktEventActionsView_rewindBeforeInvocationId_readsAndWritesThrough() { + // A marker set in place by a Java tool must reach the engine's actions. + val kt = KtEventActions(rewindBeforeInvocationId = "inv1") + val view = KtEventActionsToJavaView(kt) + + assertEquals( + "inv1", + view.rewindBeforeInvocationId().getOrNull(), + "the view should read through", + ) + view.setRewindBeforeInvocationId("inv2") + assertEquals("inv2", kt.rewindBeforeInvocationId, "the view should write through") + } + + @Test + fun reconcileActions_carriesRewindBeforeInvocationId() { + // A Java tool that replaces its actions (setActions) must still carry the marker. + val kt = KtEventActions() + + reconcileActionsToKt(JavaEventActions.builder().rewindBeforeInvocationId("inv1").build(), kt) + + assertEquals("inv1", kt.rewindBeforeInvocationId) + } + @Test fun stateValueFromJava_nullMapsToRemovalSentinel() { // ADK Java treats a null state value as a removal; the bridge must not NPE on the runAsync