From 891cd8cee35bb2621155fc21993b472ed2cd0edc Mon Sep 17 00:00:00 2001 From: xavidop Date: Wed, 1 Jul 2026 08:38:38 +0200 Subject: [PATCH 1/2] chore(deps)(deps): bump com.puppycrawl.tools:checkstyle (#235) Bumps [com.puppycrawl.tools:checkstyle](https://github.com/checkstyle/checkstyle) from 13.6.0 to 13.7.0. - [Release notes](https://github.com/checkstyle/checkstyle/releases) - [Commits](https://github.com/checkstyle/checkstyle/compare/checkstyle-13.6.0...checkstyle-13.7.0) --- updated-dependencies: - dependency-name: com.puppycrawl.tools:checkstyle dependency-version: 13.7.0 dependency-type: direct:production update-type: version-update:semver-minor ... Signed-off-by: dependabot[bot] Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com> --- pom.xml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pom.xml b/pom.xml index 041d1e72b..ede78d90b 100644 --- a/pom.xml +++ b/pom.xml @@ -142,7 +142,7 @@ 5.4.0 2.29 1.35.0 - 13.6.0 + 13.7.0 From e5e8435a700c12780038a17a6dfca1ce0eeec965 Mon Sep 17 00:00:00 2001 From: xavidop Date: Sat, 4 Jul 2026 16:22:39 +0200 Subject: [PATCH 2/2] feat: genkit agents --- .gitignore | 3 +- README.md | 58 +- .../main/java/com/google/genkit/ai/Agent.java | 245 --- .../com/google/genkit/ai/AgentConfig.java | 336 ---- .../genkit/ai/AgentHandoffException.java | 80 - .../google/genkit/ai/ModelResponseChunk.java | 21 + .../main/java/com/google/genkit/ai/Tool.java | 3 - .../com/google/genkit/ai/agent/Agent.java | 324 ++++ .../genkit/ai/agent/AgentAbortRequest.java | 78 + .../genkit/ai/agent/AgentAbortResponse.java | 106 ++ .../com/google/genkit/ai/agent/AgentApi.java | 111 ++ .../com/google/genkit/ai/agent/AgentChat.java | 409 +++++ .../google/genkit/ai/agent/AgentChunk.java | 95 ++ .../genkit/ai/agent/AgentFinishReason.java | 90 ++ .../com/google/genkit/ai/agent/AgentFn.java | 42 + .../genkit/ai/agent/AgentFnContext.java | 131 ++ .../com/google/genkit/ai/agent/AgentInit.java | 143 ++ .../google/genkit/ai/agent/AgentInput.java | 150 ++ .../genkit/ai/agent/AgentInterrupt.java | 88 ++ .../google/genkit/ai/agent/AgentMetadata.java | 140 ++ .../google/genkit/ai/agent/AgentOutput.java | 258 +++ .../com/google/genkit/ai/agent/AgentRef.java | 60 + .../google/genkit/ai/agent/AgentResponse.java | 179 +++ .../google/genkit/ai/agent/AgentResult.java | 137 ++ .../genkit/ai/agent/AgentSessionContext.java | 98 ++ .../genkit/ai/agent/AgentStreamChunk.java | 169 ++ .../genkit/ai/agent/AgentTransport.java | 72 + .../com/google/genkit/ai/agent/Artifact.java | 139 ++ .../google/genkit/ai/agent/ArtifactStore.java | 45 + .../genkit/ai/agent/ClientTransform.java | 37 + .../genkit/ai/agent/CustomAgentConfig.java | 210 +++ .../genkit/ai/agent/FileSessionStore.java | 762 +++++++++ .../genkit/ai/agent/GetSnapshotOptions.java | 104 ++ .../genkit/ai/agent/GetSnapshotRequest.java | 106 ++ .../genkit/ai/agent/InMemorySessionStore.java | 313 ++++ .../google/genkit/ai/agent/RuntimeError.java | 134 ++ .../com/google/genkit/ai/agent/Session.java | 349 +++++ .../google/genkit/ai/agent/SessionRunner.java | 448 ++++++ .../genkit/ai/agent/SessionSnapshot.java | 345 ++++ .../google/genkit/ai/agent/SessionState.java | 174 ++ .../google/genkit/ai/agent/SessionStore.java | 32 + .../genkit/ai/agent/SessionStoreOptions.java | 66 + .../genkit/ai/agent/SnapshotMutator.java | 42 + .../genkit/ai/agent/SnapshotReader.java | 40 + .../genkit/ai/agent/SnapshotStatus.java | 92 ++ .../genkit/ai/agent/SnapshotSubscriber.java | 45 + .../genkit/ai/agent/SnapshotWriter.java | 50 + .../google/genkit/ai/agent/ToolResume.java | 109 ++ .../com/google/genkit/ai/agent/TurnBody.java | 43 + .../google/genkit/ai/agent/TurnContext.java | 71 + .../com/google/genkit/ai/agent/TurnEnd.java | 106 ++ .../ai/agent/internal/AbortAwareMutator.java | 55 + .../ai/agent/internal/AgentActions.java | 428 +++++ .../ai/agent/internal/DetachController.java | 292 ++++ .../ai/agent/internal/InProcessTransport.java | 122 ++ .../ai/agent/internal/LeafSelection.java | 148 ++ .../agent/internal/PendingAbortRegistry.java | 101 ++ .../genkit/ai/agent/internal/PointerDoc.java | 115 ++ .../ai/agent/internal/SessionResolver.java | 331 ++++ .../ai/agent/internal/SnapshotSharding.java | 188 +++ .../ai/agent/internal/StreamEmitter.java | 130 ++ .../google/genkit/ai/agent/package-info.java | 23 + .../com/google/genkit/ai/session/Chat.java | 986 ------------ .../google/genkit/ai/session/ChatOptions.java | 322 ---- .../ai/session/InMemorySessionStore.java | 83 - .../com/google/genkit/ai/session/Session.java | 322 ---- .../genkit/ai/session/SessionContext.java | 167 -- .../google/genkit/ai/session/SessionData.java | 242 --- .../genkit/ai/session/SessionOptions.java | 155 -- .../genkit/ai/session/SessionStore.java | 75 - .../genkit/ai/session/package-info.java | 73 - .../com/google/genkit/ai/AgentConfigTest.java | 174 -- .../java/com/google/genkit/ai/AgentTest.java | 185 --- .../google/genkit/ai/agent/AgentChatTest.java | 596 +++++++ .../google/genkit/ai/agent/CoreFixesTest.java | 412 +++++ .../ai/agent/DefineCustomAgentTest.java | 416 +++++ .../google/genkit/ai/agent/DetachTest.java | 318 ++++ .../google/genkit/ai/agent/EnumSerdeTest.java | 119 ++ .../genkit/ai/agent/FileSessionStoreTest.java | 401 +++++ .../ai/agent/InMemorySessionStoreTest.java | 565 +++++++ .../genkit/ai/agent/LeafSelectionTest.java | 164 ++ .../genkit/ai/agent/SessionResolverTest.java | 246 +++ .../genkit/ai/agent/SessionRunnerTest.java | 392 +++++ .../google/genkit/ai/agent/SessionTest.java | 398 +++++ .../genkit/ai/agent/StreamEmitterTest.java | 243 +++ .../genkit/ai/agent/WireTypesSerdeTest.java | 504 ++++++ .../agent/internal/SnapshotShardingTest.java | 192 +++ .../genkit/ai/session/ChatOptionsTest.java | 285 ---- .../ai/session/InMemorySessionStoreTest.java | 210 --- .../genkit/ai/session/SessionContextTest.java | 263 ---- .../genkit/ai/session/SessionDataTest.java | 216 --- .../genkit/ai/session/SessionOptionsTest.java | 156 -- .../com/google/genkit/core/ActionContext.java | 207 ++- .../com/google/genkit/core/ActionType.java | 11 +- .../com/google/genkit/core/BidiAction.java | 128 ++ .../google/genkit/core/BidiActionImpl.java | 523 ++++++ .../genkit/core/BufferedInputSource.java | 93 ++ .../com/google/genkit/core/InputSource.java | 46 + .../genkit/core/jsonpatch/JsonPatch.java | 458 ++++++ .../genkit/core/jsonpatch/package-info.java | 52 + .../google/genkit/core/ActionContextTest.java | 74 + .../google/genkit/core/ActionTypeTest.java | 28 + .../genkit/core/BidiActionImplTest.java | 339 ++++ .../genkit/core/jsonpatch/JsonPatchTest.java | 445 ++++++ docs/astro.config.mjs | 18 +- .../docs/agents/background-execution.md | 115 ++ .../docs/agents/custom-orchestration.md | 153 ++ docs/src/content/docs/agents/define-agents.md | 191 +++ .../src/content/docs/agents/error-handling.md | 100 ++ docs/src/content/docs/agents/interrupts.md | 131 ++ .../docs/agents/multi-agent-delegation.md | 102 ++ docs/src/content/docs/agents/overview.md | 179 +++ .../src/content/docs/agents/run-and-stream.md | 187 +++ .../content/docs/agents/serve-over-http.md | 204 +++ .../src/content/docs/agents/session-stores.md | 236 +++ docs/src/content/docs/agents/sessions.md | 114 ++ docs/src/content/docs/chat-sessions.md | 244 --- docs/src/content/docs/multi-agent.md | 155 -- docs/src/content/docs/plugins/aws-bedrock.md | 14 + .../src/content/docs/plugins/azure-foundry.md | 17 + docs/src/content/docs/plugins/firebase.md | 44 +- docs/src/content/docs/plugins/jetty.md | 14 + docs/src/content/docs/plugins/spring.md | 14 + docs/src/content/docs/samples.md | 3 +- genkit/pom.xml | 6 + .../main/java/com/google/genkit/Genkit.java | 381 ++--- .../java/com/google/genkit/GenkitBeta.java | 555 +++++++ .../java/com/google/genkit/GenkitOptions.java | 22 + .../com/google/genkit/ReflectionServer.java | 64 +- .../com/google/genkit/ReflectionServerV2.java | 177 ++- .../com/google/genkit/agent/AgentConfig.java | 374 +++++ .../genkit/client/HttpAgentTransport.java | 267 ++++ .../com/google/genkit/client/RemoteAgent.java | 52 + .../genkit/client/RemoteAgentOptions.java | 192 +++ .../google/genkit/ExecutionContextTest.java | 247 +++ .../com/google/genkit/GenkitBetaTest.java | 576 +++++++ .../genkit/ReflectionServerV2BidiTest.java | 220 +++ .../ReflectionServerV2MultiTurnTest.java | 365 +++++ .../agent/AgentConformanceTest.java | 784 +++++++++ .../genkit/conformance/agent/Fixtures.java | 334 ++++ .../conformance/agent/ProgrammableModel.java | 148 ++ .../src/test/resources/conformance/agent.yaml | 1396 +++++++++++++++++ genkit/src/test/resources/logback-test.xml | 18 + .../test/resources/prompts/topicAgent.prompt | 4 + plugins/aws-bedrock/pom.xml | 21 +- .../plugins/awsbedrock/AwsBedrockPlugin.java | 9 + .../session/DynamoDbSessionStore.java | 739 +++++++++ .../session/DynamoDbSessionStoreOptions.java | 243 +++ .../awsbedrock/session/package-info.java | 27 + .../DynamoDbSessionStoreOptionsTest.java | 76 + .../session/DynamoDbSessionStoreTest.java | 193 +++ plugins/azure-foundry/pom.xml | 24 + .../azurefoundry/AzureFoundryPlugin.java | 42 +- .../session/CosmosSessionStore.java | 713 +++++++++ .../session/CosmosSessionStoreOptions.java | 273 ++++ .../azurefoundry/session/package-info.java | 27 + .../CosmosSessionStoreOptionsTest.java | 79 + .../session/CosmosSessionStoreTest.java | 190 +++ .../session/FirestoreSessionStore.java | 634 ++++++++ .../session/FirestoreSessionStoreOptions.java | 192 +++ .../firebase/session/package-info.java | 29 + .../session/FirestoreSessionStoreTest.java | 292 ++++ .../genkit/plugins/jetty/JettyPlugin.java | 434 +++++ .../plugins/jetty/AgentAbortHttpTest.java | 244 +++ .../plugins/jetty/AgentDetachHttpTest.java | 221 +++ .../jetty/AgentHeaderPropagationTest.java | 205 +++ .../plugins/jetty/AgentHttpServingTest.java | 246 +++ .../genkit/plugins/jetty/RemoteAgentTest.java | 487 ++++++ plugins/middleware/pom.xml | 92 ++ .../genkit/plugins/middleware/Agents.java | 244 +++ .../plugins/middleware/AgentsOptions.java | 227 +++ .../plugins/middleware/ArtifactStrategy.java | 37 + .../genkit/plugins/middleware/Artifacts.java | 222 +++ .../plugins/middleware/ArtifactsOptions.java | 89 ++ .../middleware/internal/Delegation.java | 356 +++++ .../plugins/middleware/package-info.java | 40 + .../genkit/plugins/middleware/AgentsTest.java | 551 +++++++ .../plugins/middleware/ArtifactsTest.java | 182 +++ .../plugins/spring/GenkitAgentController.java | 569 +++++++ .../spring/GenkitSpringApplication.java | 41 + .../plugins/spring/AgentAbortHttpTest.java | 245 +++ .../plugins/spring/AgentDetachHttpTest.java | 222 +++ .../spring/AgentHeaderPropagationTest.java | 209 +++ .../plugins/spring/AgentHttpServingTest.java | 248 +++ .../plugins/spring/RemoteAgentTest.java | 492 ++++++ pom.xml | 45 +- samples/README.md | 57 + samples/agents-cosmos-session/README.md | 92 ++ samples/agents-cosmos-session/pom.xml | 82 + samples/agents-cosmos-session/run.sh | 7 + .../genkit/samples/CosmosSessionAgentApp.java | 223 +++ .../src/main/resources/logback.xml | 24 + samples/agents-dynamodb-session/README.md | 47 + .../pom.xml | 10 +- samples/agents-dynamodb-session/run.sh | 7 + .../samples/DynamoDbSessionAgentApp.java | 200 +++ .../src/main/resources/logback.xml | 24 + samples/agents-firestore-session/README.md | 47 + samples/agents-firestore-session/pom.xml | 87 + samples/agents-firestore-session/run.sh | 7 + .../samples/FirestoreSessionAgentApp.java | 196 +++ .../src/main/resources/logback.xml | 23 + samples/agents-human-in-the-loop/README.md | 32 + samples/agents-human-in-the-loop/pom.xml | 77 + samples/agents-human-in-the-loop/run.sh | 7 + .../samples/HumanInTheLoopAgentApp.java | 233 +++ .../src/main/resources/logback.xml | 26 + samples/agents-orchestrator/pom.xml | 87 + samples/agents-orchestrator/run.sh | 7 + .../genkit/samples/OrchestratorApp.java | 193 +++ .../src/main/resources/logback.xml | 24 + samples/agents-remote/pom.xml | 72 + samples/agents-remote/run.sh | 7 + .../genkit/samples/RemoteAgentClientApp.java | 128 ++ .../src/main/resources/logback.xml | 24 + samples/agents-stateless/pom.xml | 77 + samples/agents-stateless/run.sh | 7 + .../genkit/samples/StatelessAgentApp.java | 165 ++ .../src/main/java/resources/logback.xml | 24 + .../{multi-agent => agents-weather}/pom.xml | 8 +- samples/agents-weather/run.sh | 7 + .../genkit/samples/WeatherAgentApp.java | 274 ++++ .../src/main/resources/logback.xml | 26 + samples/aws-bedrock/README.md | 2 +- .../genkit/samples/AwsBedrockSample.java | 13 +- samples/chat-session/README.md | 91 -- samples/chat-session/run.sh | 4 - .../google/genkit/samples/ChatSessionApp.java | 451 ------ .../google/genkit/samples/InterruptsApp.java | 310 +--- samples/multi-agent/README.md | 204 --- samples/multi-agent/run.sh | 4 - .../google/genkit/samples/MultiAgentApp.java | 560 ------- .../src/main/resources/logback.xml | 13 - .../google/genkit/samples/SessionSample.java | 345 ---- 234 files changed, 36268 insertions(+), 7240 deletions(-) delete mode 100644 ai/src/main/java/com/google/genkit/ai/Agent.java delete mode 100644 ai/src/main/java/com/google/genkit/ai/AgentConfig.java delete mode 100644 ai/src/main/java/com/google/genkit/ai/AgentHandoffException.java create mode 100644 ai/src/main/java/com/google/genkit/ai/agent/Agent.java create mode 100644 ai/src/main/java/com/google/genkit/ai/agent/AgentAbortRequest.java create mode 100644 ai/src/main/java/com/google/genkit/ai/agent/AgentAbortResponse.java create mode 100644 ai/src/main/java/com/google/genkit/ai/agent/AgentApi.java create mode 100644 ai/src/main/java/com/google/genkit/ai/agent/AgentChat.java create mode 100644 ai/src/main/java/com/google/genkit/ai/agent/AgentChunk.java create mode 100644 ai/src/main/java/com/google/genkit/ai/agent/AgentFinishReason.java create mode 100644 ai/src/main/java/com/google/genkit/ai/agent/AgentFn.java create mode 100644 ai/src/main/java/com/google/genkit/ai/agent/AgentFnContext.java create mode 100644 ai/src/main/java/com/google/genkit/ai/agent/AgentInit.java create mode 100644 ai/src/main/java/com/google/genkit/ai/agent/AgentInput.java create mode 100644 ai/src/main/java/com/google/genkit/ai/agent/AgentInterrupt.java create mode 100644 ai/src/main/java/com/google/genkit/ai/agent/AgentMetadata.java create mode 100644 ai/src/main/java/com/google/genkit/ai/agent/AgentOutput.java create mode 100644 ai/src/main/java/com/google/genkit/ai/agent/AgentRef.java create mode 100644 ai/src/main/java/com/google/genkit/ai/agent/AgentResponse.java create mode 100644 ai/src/main/java/com/google/genkit/ai/agent/AgentResult.java create mode 100644 ai/src/main/java/com/google/genkit/ai/agent/AgentSessionContext.java create mode 100644 ai/src/main/java/com/google/genkit/ai/agent/AgentStreamChunk.java create mode 100644 ai/src/main/java/com/google/genkit/ai/agent/AgentTransport.java create mode 100644 ai/src/main/java/com/google/genkit/ai/agent/Artifact.java create mode 100644 ai/src/main/java/com/google/genkit/ai/agent/ArtifactStore.java create mode 100644 ai/src/main/java/com/google/genkit/ai/agent/ClientTransform.java create mode 100644 ai/src/main/java/com/google/genkit/ai/agent/CustomAgentConfig.java create mode 100644 ai/src/main/java/com/google/genkit/ai/agent/FileSessionStore.java create mode 100644 ai/src/main/java/com/google/genkit/ai/agent/GetSnapshotOptions.java create mode 100644 ai/src/main/java/com/google/genkit/ai/agent/GetSnapshotRequest.java create mode 100644 ai/src/main/java/com/google/genkit/ai/agent/InMemorySessionStore.java create mode 100644 ai/src/main/java/com/google/genkit/ai/agent/RuntimeError.java create mode 100644 ai/src/main/java/com/google/genkit/ai/agent/Session.java create mode 100644 ai/src/main/java/com/google/genkit/ai/agent/SessionRunner.java create mode 100644 ai/src/main/java/com/google/genkit/ai/agent/SessionSnapshot.java create mode 100644 ai/src/main/java/com/google/genkit/ai/agent/SessionState.java create mode 100644 ai/src/main/java/com/google/genkit/ai/agent/SessionStore.java create mode 100644 ai/src/main/java/com/google/genkit/ai/agent/SessionStoreOptions.java create mode 100644 ai/src/main/java/com/google/genkit/ai/agent/SnapshotMutator.java create mode 100644 ai/src/main/java/com/google/genkit/ai/agent/SnapshotReader.java create mode 100644 ai/src/main/java/com/google/genkit/ai/agent/SnapshotStatus.java create mode 100644 ai/src/main/java/com/google/genkit/ai/agent/SnapshotSubscriber.java create mode 100644 ai/src/main/java/com/google/genkit/ai/agent/SnapshotWriter.java create mode 100644 ai/src/main/java/com/google/genkit/ai/agent/ToolResume.java create mode 100644 ai/src/main/java/com/google/genkit/ai/agent/TurnBody.java create mode 100644 ai/src/main/java/com/google/genkit/ai/agent/TurnContext.java create mode 100644 ai/src/main/java/com/google/genkit/ai/agent/TurnEnd.java create mode 100644 ai/src/main/java/com/google/genkit/ai/agent/internal/AbortAwareMutator.java create mode 100644 ai/src/main/java/com/google/genkit/ai/agent/internal/AgentActions.java create mode 100644 ai/src/main/java/com/google/genkit/ai/agent/internal/DetachController.java create mode 100644 ai/src/main/java/com/google/genkit/ai/agent/internal/InProcessTransport.java create mode 100644 ai/src/main/java/com/google/genkit/ai/agent/internal/LeafSelection.java create mode 100644 ai/src/main/java/com/google/genkit/ai/agent/internal/PendingAbortRegistry.java create mode 100644 ai/src/main/java/com/google/genkit/ai/agent/internal/PointerDoc.java create mode 100644 ai/src/main/java/com/google/genkit/ai/agent/internal/SessionResolver.java create mode 100644 ai/src/main/java/com/google/genkit/ai/agent/internal/SnapshotSharding.java create mode 100644 ai/src/main/java/com/google/genkit/ai/agent/internal/StreamEmitter.java create mode 100644 ai/src/main/java/com/google/genkit/ai/agent/package-info.java delete mode 100644 ai/src/main/java/com/google/genkit/ai/session/Chat.java delete mode 100644 ai/src/main/java/com/google/genkit/ai/session/ChatOptions.java delete mode 100644 ai/src/main/java/com/google/genkit/ai/session/InMemorySessionStore.java delete mode 100644 ai/src/main/java/com/google/genkit/ai/session/Session.java delete mode 100644 ai/src/main/java/com/google/genkit/ai/session/SessionContext.java delete mode 100644 ai/src/main/java/com/google/genkit/ai/session/SessionData.java delete mode 100644 ai/src/main/java/com/google/genkit/ai/session/SessionOptions.java delete mode 100644 ai/src/main/java/com/google/genkit/ai/session/SessionStore.java delete mode 100644 ai/src/main/java/com/google/genkit/ai/session/package-info.java delete mode 100644 ai/src/test/java/com/google/genkit/ai/AgentConfigTest.java delete mode 100644 ai/src/test/java/com/google/genkit/ai/AgentTest.java create mode 100644 ai/src/test/java/com/google/genkit/ai/agent/AgentChatTest.java create mode 100644 ai/src/test/java/com/google/genkit/ai/agent/CoreFixesTest.java create mode 100644 ai/src/test/java/com/google/genkit/ai/agent/DefineCustomAgentTest.java create mode 100644 ai/src/test/java/com/google/genkit/ai/agent/DetachTest.java create mode 100644 ai/src/test/java/com/google/genkit/ai/agent/EnumSerdeTest.java create mode 100644 ai/src/test/java/com/google/genkit/ai/agent/FileSessionStoreTest.java create mode 100644 ai/src/test/java/com/google/genkit/ai/agent/InMemorySessionStoreTest.java create mode 100644 ai/src/test/java/com/google/genkit/ai/agent/LeafSelectionTest.java create mode 100644 ai/src/test/java/com/google/genkit/ai/agent/SessionResolverTest.java create mode 100644 ai/src/test/java/com/google/genkit/ai/agent/SessionRunnerTest.java create mode 100644 ai/src/test/java/com/google/genkit/ai/agent/SessionTest.java create mode 100644 ai/src/test/java/com/google/genkit/ai/agent/StreamEmitterTest.java create mode 100644 ai/src/test/java/com/google/genkit/ai/agent/WireTypesSerdeTest.java create mode 100644 ai/src/test/java/com/google/genkit/ai/agent/internal/SnapshotShardingTest.java delete mode 100644 ai/src/test/java/com/google/genkit/ai/session/ChatOptionsTest.java delete mode 100644 ai/src/test/java/com/google/genkit/ai/session/InMemorySessionStoreTest.java delete mode 100644 ai/src/test/java/com/google/genkit/ai/session/SessionContextTest.java delete mode 100644 ai/src/test/java/com/google/genkit/ai/session/SessionDataTest.java delete mode 100644 ai/src/test/java/com/google/genkit/ai/session/SessionOptionsTest.java create mode 100644 core/src/main/java/com/google/genkit/core/BidiAction.java create mode 100644 core/src/main/java/com/google/genkit/core/BidiActionImpl.java create mode 100644 core/src/main/java/com/google/genkit/core/BufferedInputSource.java create mode 100644 core/src/main/java/com/google/genkit/core/InputSource.java create mode 100644 core/src/main/java/com/google/genkit/core/jsonpatch/JsonPatch.java create mode 100644 core/src/main/java/com/google/genkit/core/jsonpatch/package-info.java create mode 100644 core/src/test/java/com/google/genkit/core/BidiActionImplTest.java create mode 100644 core/src/test/java/com/google/genkit/core/jsonpatch/JsonPatchTest.java create mode 100644 docs/src/content/docs/agents/background-execution.md create mode 100644 docs/src/content/docs/agents/custom-orchestration.md create mode 100644 docs/src/content/docs/agents/define-agents.md create mode 100644 docs/src/content/docs/agents/error-handling.md create mode 100644 docs/src/content/docs/agents/interrupts.md create mode 100644 docs/src/content/docs/agents/multi-agent-delegation.md create mode 100644 docs/src/content/docs/agents/overview.md create mode 100644 docs/src/content/docs/agents/run-and-stream.md create mode 100644 docs/src/content/docs/agents/serve-over-http.md create mode 100644 docs/src/content/docs/agents/session-stores.md create mode 100644 docs/src/content/docs/agents/sessions.md delete mode 100644 docs/src/content/docs/chat-sessions.md delete mode 100644 docs/src/content/docs/multi-agent.md create mode 100644 genkit/src/main/java/com/google/genkit/GenkitBeta.java create mode 100644 genkit/src/main/java/com/google/genkit/agent/AgentConfig.java create mode 100644 genkit/src/main/java/com/google/genkit/client/HttpAgentTransport.java create mode 100644 genkit/src/main/java/com/google/genkit/client/RemoteAgent.java create mode 100644 genkit/src/main/java/com/google/genkit/client/RemoteAgentOptions.java create mode 100644 genkit/src/test/java/com/google/genkit/ExecutionContextTest.java create mode 100644 genkit/src/test/java/com/google/genkit/GenkitBetaTest.java create mode 100644 genkit/src/test/java/com/google/genkit/ReflectionServerV2BidiTest.java create mode 100644 genkit/src/test/java/com/google/genkit/ReflectionServerV2MultiTurnTest.java create mode 100644 genkit/src/test/java/com/google/genkit/conformance/agent/AgentConformanceTest.java create mode 100644 genkit/src/test/java/com/google/genkit/conformance/agent/Fixtures.java create mode 100644 genkit/src/test/java/com/google/genkit/conformance/agent/ProgrammableModel.java create mode 100644 genkit/src/test/resources/conformance/agent.yaml create mode 100644 genkit/src/test/resources/logback-test.xml create mode 100644 genkit/src/test/resources/prompts/topicAgent.prompt create mode 100644 plugins/aws-bedrock/src/main/java/com/google/genkit/plugins/awsbedrock/session/DynamoDbSessionStore.java create mode 100644 plugins/aws-bedrock/src/main/java/com/google/genkit/plugins/awsbedrock/session/DynamoDbSessionStoreOptions.java create mode 100644 plugins/aws-bedrock/src/main/java/com/google/genkit/plugins/awsbedrock/session/package-info.java create mode 100644 plugins/aws-bedrock/src/test/java/com/google/genkit/plugins/awsbedrock/session/DynamoDbSessionStoreOptionsTest.java create mode 100644 plugins/aws-bedrock/src/test/java/com/google/genkit/plugins/awsbedrock/session/DynamoDbSessionStoreTest.java create mode 100644 plugins/azure-foundry/src/main/java/com/google/genkit/plugins/azurefoundry/session/CosmosSessionStore.java create mode 100644 plugins/azure-foundry/src/main/java/com/google/genkit/plugins/azurefoundry/session/CosmosSessionStoreOptions.java create mode 100644 plugins/azure-foundry/src/main/java/com/google/genkit/plugins/azurefoundry/session/package-info.java create mode 100644 plugins/azure-foundry/src/test/java/com/google/genkit/plugins/azurefoundry/session/CosmosSessionStoreOptionsTest.java create mode 100644 plugins/azure-foundry/src/test/java/com/google/genkit/plugins/azurefoundry/session/CosmosSessionStoreTest.java create mode 100644 plugins/firebase/src/main/java/com/google/genkit/plugins/firebase/session/FirestoreSessionStore.java create mode 100644 plugins/firebase/src/main/java/com/google/genkit/plugins/firebase/session/FirestoreSessionStoreOptions.java create mode 100644 plugins/firebase/src/main/java/com/google/genkit/plugins/firebase/session/package-info.java create mode 100644 plugins/firebase/src/test/java/com/google/genkit/plugins/firebase/session/FirestoreSessionStoreTest.java create mode 100644 plugins/jetty/src/test/java/com/google/genkit/plugins/jetty/AgentAbortHttpTest.java create mode 100644 plugins/jetty/src/test/java/com/google/genkit/plugins/jetty/AgentDetachHttpTest.java create mode 100644 plugins/jetty/src/test/java/com/google/genkit/plugins/jetty/AgentHeaderPropagationTest.java create mode 100644 plugins/jetty/src/test/java/com/google/genkit/plugins/jetty/AgentHttpServingTest.java create mode 100644 plugins/jetty/src/test/java/com/google/genkit/plugins/jetty/RemoteAgentTest.java create mode 100644 plugins/middleware/pom.xml create mode 100644 plugins/middleware/src/main/java/com/google/genkit/plugins/middleware/Agents.java create mode 100644 plugins/middleware/src/main/java/com/google/genkit/plugins/middleware/AgentsOptions.java create mode 100644 plugins/middleware/src/main/java/com/google/genkit/plugins/middleware/ArtifactStrategy.java create mode 100644 plugins/middleware/src/main/java/com/google/genkit/plugins/middleware/Artifacts.java create mode 100644 plugins/middleware/src/main/java/com/google/genkit/plugins/middleware/ArtifactsOptions.java create mode 100644 plugins/middleware/src/main/java/com/google/genkit/plugins/middleware/internal/Delegation.java create mode 100644 plugins/middleware/src/main/java/com/google/genkit/plugins/middleware/package-info.java create mode 100644 plugins/middleware/src/test/java/com/google/genkit/plugins/middleware/AgentsTest.java create mode 100644 plugins/middleware/src/test/java/com/google/genkit/plugins/middleware/ArtifactsTest.java create mode 100644 plugins/spring/src/main/java/com/google/genkit/plugins/spring/GenkitAgentController.java create mode 100644 plugins/spring/src/test/java/com/google/genkit/plugins/spring/AgentAbortHttpTest.java create mode 100644 plugins/spring/src/test/java/com/google/genkit/plugins/spring/AgentDetachHttpTest.java create mode 100644 plugins/spring/src/test/java/com/google/genkit/plugins/spring/AgentHeaderPropagationTest.java create mode 100644 plugins/spring/src/test/java/com/google/genkit/plugins/spring/AgentHttpServingTest.java create mode 100644 plugins/spring/src/test/java/com/google/genkit/plugins/spring/RemoteAgentTest.java create mode 100644 samples/agents-cosmos-session/README.md create mode 100644 samples/agents-cosmos-session/pom.xml create mode 100755 samples/agents-cosmos-session/run.sh create mode 100644 samples/agents-cosmos-session/src/main/java/com/google/genkit/samples/CosmosSessionAgentApp.java create mode 100644 samples/agents-cosmos-session/src/main/resources/logback.xml create mode 100644 samples/agents-dynamodb-session/README.md rename samples/{chat-session => agents-dynamodb-session}/pom.xml (87%) create mode 100755 samples/agents-dynamodb-session/run.sh create mode 100644 samples/agents-dynamodb-session/src/main/java/com/google/genkit/samples/DynamoDbSessionAgentApp.java create mode 100644 samples/agents-dynamodb-session/src/main/resources/logback.xml create mode 100644 samples/agents-firestore-session/README.md create mode 100644 samples/agents-firestore-session/pom.xml create mode 100755 samples/agents-firestore-session/run.sh create mode 100644 samples/agents-firestore-session/src/main/java/com/google/genkit/samples/FirestoreSessionAgentApp.java create mode 100644 samples/agents-firestore-session/src/main/resources/logback.xml create mode 100644 samples/agents-human-in-the-loop/README.md create mode 100644 samples/agents-human-in-the-loop/pom.xml create mode 100755 samples/agents-human-in-the-loop/run.sh create mode 100644 samples/agents-human-in-the-loop/src/main/java/com/google/genkit/samples/HumanInTheLoopAgentApp.java create mode 100644 samples/agents-human-in-the-loop/src/main/resources/logback.xml create mode 100644 samples/agents-orchestrator/pom.xml create mode 100755 samples/agents-orchestrator/run.sh create mode 100644 samples/agents-orchestrator/src/main/java/com/google/genkit/samples/OrchestratorApp.java create mode 100644 samples/agents-orchestrator/src/main/resources/logback.xml create mode 100644 samples/agents-remote/pom.xml create mode 100755 samples/agents-remote/run.sh create mode 100644 samples/agents-remote/src/main/java/com/google/genkit/samples/RemoteAgentClientApp.java create mode 100644 samples/agents-remote/src/main/resources/logback.xml create mode 100644 samples/agents-stateless/pom.xml create mode 100755 samples/agents-stateless/run.sh create mode 100644 samples/agents-stateless/src/main/java/com/google/genkit/samples/StatelessAgentApp.java create mode 100644 samples/agents-stateless/src/main/java/resources/logback.xml rename samples/{multi-agent => agents-weather}/pom.xml (90%) create mode 100755 samples/agents-weather/run.sh create mode 100644 samples/agents-weather/src/main/java/com/google/genkit/samples/WeatherAgentApp.java create mode 100644 samples/agents-weather/src/main/resources/logback.xml delete mode 100644 samples/chat-session/README.md delete mode 100755 samples/chat-session/run.sh delete mode 100644 samples/chat-session/src/main/java/com/google/genkit/samples/ChatSessionApp.java delete mode 100644 samples/multi-agent/README.md delete mode 100755 samples/multi-agent/run.sh delete mode 100644 samples/multi-agent/src/main/java/com/google/genkit/samples/MultiAgentApp.java delete mode 100644 samples/multi-agent/src/main/resources/logback.xml delete mode 100644 samples/openai/src/main/java/com/google/genkit/samples/SessionSample.java diff --git a/.gitignore b/.gitignore index 5c2d3f4f6..13ca98480 100644 --- a/.gitignore +++ b/.gitignore @@ -30,4 +30,5 @@ node_modules *.log hs_err_pid* samples/google-genai/generated_media/ -.astro \ No newline at end of file +.astro +.snapshots \ No newline at end of file diff --git a/README.md b/README.md index 9247dfaa2..85b180b9d 100644 --- a/README.md +++ b/README.md @@ -27,6 +27,7 @@ Check the Docs: https://genkit-ai.github.io/genkit-java - [Evaluations](#evaluations) - [Pre-built Evaluators Plugin](#pre-built-evaluators-plugin) - [Streaming](#streaming) + - [Agents (Beta)](#agents-beta) - [Embeddings](#embeddings) - [Modules](#modules) - [Observability](#observability) @@ -568,6 +569,60 @@ ModelResponse response = genkit.generateStream( }); ``` +## Agents (Beta) + +Agents are stateful, multi-turn AI actors that carry conversation history across turns and can call tools. The Agents API requires opting in to experimental features. + +```java +import com.google.genkit.agent.AgentConfig; +import com.google.genkit.ai.agent.Agent; +import com.google.genkit.ai.agent.AgentChat; +import com.google.genkit.ai.agent.AgentResponse; +import com.google.genkit.ai.agent.FileSessionStore; +import com.google.genkit.core.ActionContext; + +// Enable experimental features +Genkit genkit = Genkit.builder() + .options(GenkitOptions.builder().experimental(true).build()) + .plugin(OpenAIPlugin.create()) + .build(); + +// Define a tool the agent can call +Tool getWeather = genkit.defineTool( + "getWeather", + "Returns current weather for a location", + (ctx, input) -> new WeatherOutput("Sunny and 22°C in " + input.getLocation()), + WeatherInput.class, + WeatherOutput.class); + +// Define a server-managed agent (state persisted to ./.snapshots) +Agent> weatherAgent = genkit.beta().defineAgent( + AgentConfig.>builder() + .name("weatherAgent") + .system("You are a helpful weather assistant. Use the getWeather tool.") + .tools(getWeather) + .model("openai/gpt-4o-mini") + .store(new FileSessionStore<>("./.snapshots")) // omit for client-managed (stateless) + .build()); + +// Multi-turn chat — history carries forward automatically +ActionContext ctx = new ActionContext(genkit.getRegistry()); +AgentChat> chat = weatherAgent.chat(ctx); + +AgentResponse> res1 = chat.send("What is the weather in London?"); +System.out.println(res1.text()); + +AgentResponse> res2 = chat.send("Now say that in French"); +System.out.println(res2.text()); + +// Streaming turn +AgentResponse> res3 = chat.sendStream( + "Summarise in one sentence", + chunk -> System.out.print(chunk.text())); +``` + +See [samples/agents-weather](samples/agents-weather) for a complete runnable example and the [Agents documentation](https://genkit-ai.github.io/genkit-java/agents/overview) for the full API reference. + ## Embeddings Generate vector embeddings for semantic search: @@ -705,12 +760,11 @@ The following samples are available in `java/samples/`. See the [samples README] | **dotprompt** | DotPrompt files with complex inputs/outputs, variants, and partials | | **structured-output** | Type-safe structured output generation | | **rag** | RAG application with local vector store | -| **chat-session** | Multi-turn chat with session persistence | +| **agents-weather** | Weather assistant demonstrating the beta Agents API (server-managed and client-managed sessions) | | **evaluations** | Custom evaluators and evaluation workflows | | **evaluators-plugin** | Pre-built RAGAS-style evaluators plugin demo | | **complex-io** | Complex nested types, arrays, maps in flow inputs/outputs | | **middleware** | Middleware patterns for logging, caching, rate limiting | -| **multi-agent** | Multi-agent orchestration patterns | | **interrupts** | Flow interrupts and human-in-the-loop patterns | | **mcp** | Model Context Protocol (MCP) integration | | **firebase** | Firebase integration with Firestore RAG and Cloud Functions | diff --git a/ai/src/main/java/com/google/genkit/ai/Agent.java b/ai/src/main/java/com/google/genkit/ai/Agent.java deleted file mode 100644 index 3a2b99f07..000000000 --- a/ai/src/main/java/com/google/genkit/ai/Agent.java +++ /dev/null @@ -1,245 +0,0 @@ -/* - * 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. - * - * SPDX-License-Identifier: Apache-2.0 - */ - -package com.google.genkit.ai; - -import java.util.ArrayList; -import java.util.HashMap; -import java.util.List; -import java.util.Map; - -/** - * Represents an agent that can be used as a tool in multi-agent systems. - * - *

An Agent wraps an AgentConfig and provides a Tool interface for delegation. When the model - * calls an agent as a tool, the agent's configuration (system prompt, model, tools) is applied to - * the conversation, effectively "transferring" control to the specialized agent. - * - *

Example usage: - * - * 

{@code
- * // Define a specialized agent
- * Agent reservationAgent = genkit.defineAgent(
- *     AgentConfig.builder()
- *         .name("reservationAgent")
- *         .description("Handles restaurant reservations")
- *         .system("You are a reservation specialist...")
- *         .tools(List.of(reservationTool))
- *         .build());
- *
- * // Use in a parent agent
- * Agent triageAgent = genkit.defineAgent(
- *     AgentConfig.builder()
- *         .name("triageAgent")
- *         .description("Routes requests")
- *         .system("Route customer requests to specialists")
- *         .agents(List.of(reservationAgent.getConfig()))
- *         .build());
- *
- * // Start chat with triage agent
- * Chat chat = genkit.chat(triageAgent);
- * }
- */ -public class Agent { - - private final AgentConfig config; - private final Tool, AgentTransferResult> asTool; - - /** - * Creates a new Agent. - * - * @param config the agent configuration - */ - @SuppressWarnings("unchecked") - public Agent(AgentConfig config) { - this.config = config; - this.asTool = createAgentTool(); - } - - /** - * Gets the agent configuration. - * - * @return the config - */ - public AgentConfig getConfig() { - return config; - } - - /** - * Gets the agent name. - * - * @return the name - */ - public String getName() { - return config.getName(); - } - - /** - * Gets the agent description. - * - * @return the description - */ - public String getDescription() { - return config.getDescription(); - } - - /** - * Gets the system prompt. - * - * @return the system prompt - */ - public String getSystem() { - return config.getSystem(); - } - - /** - * Gets the model name. - * - * @return the model name - */ - public String getModel() { - return config.getModel(); - } - - /** - * Gets the tools available to this agent. - * - * @return the tools - */ - public List> getTools() { - return config.getTools(); - } - - /** - * Gets the sub-agents. - * - * @return the sub-agents - */ - public List getAgents() { - return config.getAgents(); - } - - /** - * Gets all tools including sub-agent tools for handoff pattern. - * - *

This method collects all tools that should be available to the agent, including the agent's - * own tools and sub-agents as tools (for handoff). When a sub-agent tool is called, the Chat will - * handle the handoff by switching context to that agent. - * - * @param agentRegistry map of agent name to Agent instance - * @return combined list of all tools from this agent and sub-agents as tools - */ - public List> getAllTools(Map agentRegistry) { - List> allTools = new ArrayList<>(); - - // Add this agent's direct tools - if (config.getTools() != null) { - allTools.addAll(config.getTools()); - } - - // Add sub-agents as tools (for handoff pattern) - if (config.getAgents() != null) { - for (AgentConfig agentConfig : config.getAgents()) { - Agent agent = agentRegistry.get(agentConfig.getName()); - if (agent != null) { - // Add the sub-agent as a tool - when called, Chat will handle the handoff - allTools.add(agent.asTool()); - } - } - } - - return allTools; - } - - /** - * Returns this agent as a tool that can be used by other agents. - * - * @return the agent as a tool - */ - public Tool, AgentTransferResult> asTool() { - return asTool; - } - - /** - * Gets the tool definition for this agent. - * - * @return the tool definition - */ - public ToolDefinition getToolDefinition() { - return asTool.getDefinition(); - } - - /** Creates the agent-as-tool wrapper. */ - @SuppressWarnings("unchecked") - private Tool, AgentTransferResult> createAgentTool() { - // OpenAI requires "properties" field even if empty - Map inputSchema = new HashMap<>(); - inputSchema.put("type", "object"); - inputSchema.put("properties", new HashMap()); - inputSchema.put("additionalProperties", true); - - Map outputSchema = new HashMap<>(); - outputSchema.put("type", "object"); - outputSchema.put( - "properties", - Map.of( - "transferredTo", Map.of("type", "string"), "transferred", Map.of("type", "boolean"))); - - return new Tool<>( - config.getName(), - config.getDescription() != null - ? config.getDescription() - : "Transfer to " + config.getName(), - inputSchema, - outputSchema, - (Class>) (Class) Map.class, - (ctx, input) -> { - // Throw handoff exception to signal the chat to switch to this agent - throw new AgentHandoffException(config.getName(), config, input); - }); - } - - /** Result of an agent transfer. */ - public static class AgentTransferResult { - private final String transferredTo; - private final boolean transferred; - - public AgentTransferResult(String agentName) { - this.transferredTo = agentName; - this.transferred = true; - } - - public String getTransferredTo() { - return transferredTo; - } - - public boolean isTransferred() { - return transferred; - } - - @Override - public String toString() { - return "transferred to " + transferredTo; - } - } - - @Override - public String toString() { - return "Agent{" + "name='" + config.getName() + '\'' + '}'; - } -} diff --git a/ai/src/main/java/com/google/genkit/ai/AgentConfig.java b/ai/src/main/java/com/google/genkit/ai/AgentConfig.java deleted file mode 100644 index b0b534b3b..000000000 --- a/ai/src/main/java/com/google/genkit/ai/AgentConfig.java +++ /dev/null @@ -1,336 +0,0 @@ -/* - * 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. - * - * SPDX-License-Identifier: Apache-2.0 - */ - -package com.google.genkit.ai; - -import java.util.List; - -/** - * Configuration for defining an agent (prompt as tool). - * - *

An agent is a specialized prompt that can be used as a tool, enabling multi-agent systems - * where one agent can delegate tasks to other specialized agents. - * - *

Example usage: - * - * 

{@code
- * // Define a specialized agent
- * AgentConfig reservationAgent = AgentConfig.builder()
- *     .name("reservationAgent")
- *     .description("Handles restaurant reservations")
- *     .system("You are a reservation specialist. Help users make and manage reservations.")
- *     .model("openai/gpt-4o")
- *     .tools(List.of(reservationTool, cancelTool))
- *     .build();
- *
- * // Use as a tool in a triage agent
- * AgentConfig triageAgent = AgentConfig.builder()
- *     .name("triageAgent")
- *     .description("Routes customer requests to appropriate specialists")
- *     .system("You are a customer service triage agent...")
- *     .agents(List.of(reservationAgent, menuAgent)) // Sub-agents
- *     // as
- *     // tools
- *     .build();
- * }
- */ -public class AgentConfig { - - private String name; - private String description; - private String system; - private String model; - private List> tools; - private List agents; - private GenerationConfig config; - private OutputConfig output; - - /** Default constructor. */ - public AgentConfig() {} - - /** - * Gets the agent name. - * - * @return the name - */ - public String getName() { - return name; - } - - /** - * Sets the agent name. - * - * @param name the name - */ - public void setName(String name) { - this.name = name; - } - - /** - * Gets the description. - * - * @return the description - */ - public String getDescription() { - return description; - } - - /** - * Sets the description (used when agent is called as a tool). - * - * @param description the description - */ - public void setDescription(String description) { - this.description = description; - } - - /** - * Gets the system prompt. - * - * @return the system prompt - */ - public String getSystem() { - return system; - } - - /** - * Sets the system prompt. - * - * @param system the system prompt - */ - public void setSystem(String system) { - this.system = system; - } - - /** - * Gets the model name. - * - * @return the model name - */ - public String getModel() { - return model; - } - - /** - * Sets the model name. - * - * @param model the model name - */ - public void setModel(String model) { - this.model = model; - } - - /** - * Gets the tools available to this agent. - * - * @return the tools - */ - public List> getTools() { - return tools; - } - - /** - * Sets the tools available to this agent. - * - * @param tools the tools - */ - public void setTools(List> tools) { - this.tools = tools; - } - - /** - * Gets the sub-agents (agents that can be delegated to). - * - * @return the sub-agents - */ - public List getAgents() { - return agents; - } - - /** - * Sets the sub-agents. - * - * @param agents the sub-agents - */ - public void setAgents(List agents) { - this.agents = agents; - } - - /** - * Gets the generation config. - * - * @return the generation config - */ - public GenerationConfig getConfig() { - return config; - } - - /** - * Sets the generation config. - * - * @param config the generation config - */ - public void setConfig(GenerationConfig config) { - this.config = config; - } - - /** - * Gets the output config. - * - * @return the output config - */ - public OutputConfig getOutput() { - return output; - } - - /** - * Sets the output config. - * - * @param output the output config - */ - public void setOutput(OutputConfig output) { - this.output = output; - } - - /** - * Creates a new builder. - * - * @return a new builder - */ - public static Builder builder() { - return new Builder(); - } - - /** Builder for AgentConfig. */ - public static class Builder { - private String name; - private String description; - private String system; - private String model; - private List> tools; - private List agents; - private GenerationConfig config; - private OutputConfig output; - - /** - * Sets the agent name. - * - * @param name the name - * @return this builder - */ - public Builder name(String name) { - this.name = name; - return this; - } - - /** - * Sets the description. - * - * @param description the description - * @return this builder - */ - public Builder description(String description) { - this.description = description; - return this; - } - - /** - * Sets the system prompt. - * - * @param system the system prompt - * @return this builder - */ - public Builder system(String system) { - this.system = system; - return this; - } - - /** - * Sets the model name. - * - * @param model the model name - * @return this builder - */ - public Builder model(String model) { - this.model = model; - return this; - } - - /** - * Sets the tools available to this agent. - * - * @param tools the tools - * @return this builder - */ - public Builder tools(List> tools) { - this.tools = tools; - return this; - } - - /** - * Sets the sub-agents. - * - * @param agents the sub-agents - * @return this builder - */ - public Builder agents(List agents) { - this.agents = agents; - return this; - } - - /** - * Sets the generation config. - * - * @param config the generation config - * @return this builder - */ - public Builder config(GenerationConfig config) { - this.config = config; - return this; - } - - /** - * Sets the output config. - * - * @param output the output config - * @return this builder - */ - public Builder output(OutputConfig output) { - this.output = output; - return this; - } - - /** - * Builds the AgentConfig. - * - * @return the built config - */ - public AgentConfig build() { - AgentConfig agentConfig = new AgentConfig(); - agentConfig.setName(name); - agentConfig.setDescription(description); - agentConfig.setSystem(system); - agentConfig.setModel(model); - agentConfig.setTools(tools); - agentConfig.setAgents(agents); - agentConfig.setConfig(config); - agentConfig.setOutput(output); - return agentConfig; - } - } -} diff --git a/ai/src/main/java/com/google/genkit/ai/AgentHandoffException.java b/ai/src/main/java/com/google/genkit/ai/AgentHandoffException.java deleted file mode 100644 index 0cdd6d60a..000000000 --- a/ai/src/main/java/com/google/genkit/ai/AgentHandoffException.java +++ /dev/null @@ -1,80 +0,0 @@ -/* - * 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. - * - * SPDX-License-Identifier: Apache-2.0 - */ - -package com.google.genkit.ai; - -import java.util.Map; - -/** - * Exception thrown when an agent tool is called to signal a handoff. - * - *

When the model calls an agent-as-tool, this exception is thrown to signal that the chat should - * switch context to the target agent. The Chat class catches this exception and updates its system - * prompt, tools, and model to those of the target agent. - * - *

This enables the "handoff" pattern where conversations can be transferred between specialized - * agents. - */ -public class AgentHandoffException extends RuntimeException { - - private final String targetAgentName; - private final AgentConfig targetAgentConfig; - private final Map handoffInput; - - /** - * Creates a new AgentHandoffException. - * - * @param targetAgentName the name of the agent to hand off to - * @param targetAgentConfig the configuration of the target agent - * @param handoffInput the input passed to the agent tool (can be used for context) - */ - public AgentHandoffException( - String targetAgentName, AgentConfig targetAgentConfig, Map handoffInput) { - super("Handoff to agent: " + targetAgentName); - this.targetAgentName = targetAgentName; - this.targetAgentConfig = targetAgentConfig; - this.handoffInput = handoffInput; - } - - /** - * Gets the name of the target agent. - * - * @return the target agent name - */ - public String getTargetAgentName() { - return targetAgentName; - } - - /** - * Gets the configuration of the target agent. - * - * @return the target agent config - */ - public AgentConfig getTargetAgentConfig() { - return targetAgentConfig; - } - - /** - * Gets the input passed to the agent tool. - * - * @return the handoff input - */ - public Map getHandoffInput() { - return handoffInput; - } -} diff --git a/ai/src/main/java/com/google/genkit/ai/ModelResponseChunk.java b/ai/src/main/java/com/google/genkit/ai/ModelResponseChunk.java index 053b2ef5c..c94116df3 100644 --- a/ai/src/main/java/com/google/genkit/ai/ModelResponseChunk.java +++ b/ai/src/main/java/com/google/genkit/ai/ModelResponseChunk.java @@ -27,6 +27,9 @@ @JsonInclude(JsonInclude.Include.NON_NULL) public class ModelResponseChunk { + @JsonProperty("role") + private Role role; + @JsonProperty("content") private List content = new ArrayList<>(); @@ -77,6 +80,24 @@ public String getText() { // Getters and setters + /** + * Returns the role of the message this chunk contributes to. + * + * @return the role, or {@code null} if not set (defaults to {@code model} on the wire) + */ + public Role getRole() { + return role; + } + + /** + * Sets the role of the message this chunk contributes to. + * + * @param role the role + */ + public void setRole(Role role) { + this.role = role; + } + public List getContent() { return content; } diff --git a/ai/src/main/java/com/google/genkit/ai/Tool.java b/ai/src/main/java/com/google/genkit/ai/Tool.java index 0c1663050..d48085726 100644 --- a/ai/src/main/java/com/google/genkit/ai/Tool.java +++ b/ai/src/main/java/com/google/genkit/ai/Tool.java @@ -129,9 +129,6 @@ public O run(ActionContext ctx, I input) throws GenkitException { try { O result = handler.apply(ctx.withSpanContext(spanCtx), in); return result; - } catch (AgentHandoffException e) { - // Re-throw agent handoff exceptions for multi-agent pattern - throw e; } catch (ToolInterruptException e) { // Re-throw interrupt exceptions for human-in-the-loop pattern throw e; diff --git a/ai/src/main/java/com/google/genkit/ai/agent/Agent.java b/ai/src/main/java/com/google/genkit/ai/agent/Agent.java new file mode 100644 index 000000000..f280f1b23 --- /dev/null +++ b/ai/src/main/java/com/google/genkit/ai/agent/Agent.java @@ -0,0 +1,324 @@ +/* + * 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. + * + * SPDX-License-Identifier: Apache-2.0 + */ + +package com.google.genkit.ai.agent; + +import com.fasterxml.jackson.databind.JsonNode; +import com.google.genkit.ai.agent.internal.InProcessTransport; +import com.google.genkit.core.Action; +import com.google.genkit.core.ActionContext; +import com.google.genkit.core.ActionDesc; +import com.google.genkit.core.ActionRunResult; +import com.google.genkit.core.ActionType; +import com.google.genkit.core.BidiAction; +import com.google.genkit.core.BidiActionImpl; +import com.google.genkit.core.GenkitException; +import com.google.genkit.core.InputSource; +import com.google.genkit.core.Registry; +import java.util.Map; +import java.util.function.Consumer; + +/** + * Agent is a bidirectional streaming agent action that manages session state and provides typed + * facades for snapshot retrieval and abort. + * + *

Instances are created via {@code AgentActions.defineCustomAgent}. Do not instantiate directly. + * + * @param the type of custom session state + */ +public final class Agent + implements BidiAction, AgentStreamChunk, AgentInit>, AgentApi { + + private final BidiActionImpl, AgentStreamChunk, AgentInit> impl; + private final SessionStore store; + private final boolean serverManaged; + + /** Nullable: companion action for retrieving snapshots; only when store != null. */ + private final Action snapshotAction; + + /** Nullable: companion action for aborting snapshots; only when store is SnapshotSubscriber. */ + private final Action abortAction; + + private final ClientTransform clientTransform; + private final SessionStoreOptions opts; + private final String name; + private final String description; + private final Registry registry; + + /** + * Constructs an Agent. Called exclusively by {@code AgentActions}. + * + * @param impl the underlying BidiActionImpl + * @param store the session store, or {@code null} for client-managed mode + * @param serverManaged whether the agent is server-managed + * @param snapshotAction companion snapshot action; nullable + * @param abortAction companion abort action; nullable + * @param clientTransform state transform for client-managed mode; nullable + * @param opts store options; nullable + * @param name the agent name + * @param description the agent description; nullable + * @param registry the registry the agent is defined in; used to build a default {@link + * ActionContext} for the no-arg {@link #chat()} convenience overloads + */ + public Agent( + BidiActionImpl, AgentStreamChunk, AgentInit> impl, + SessionStore store, + boolean serverManaged, + Action snapshotAction, + Action abortAction, + ClientTransform clientTransform, + SessionStoreOptions opts, + String name, + String description, + Registry registry) { + this.impl = impl; + this.store = store; + this.serverManaged = serverManaged; + this.snapshotAction = snapshotAction; + this.abortAction = abortAction; + this.clientTransform = clientTransform; + this.opts = opts; + this.name = name; + this.description = description; + this.registry = registry; + } + + // ── BidiAction delegation ───────────────────────────────────────────────────── + + @Override + public String getName() { + return impl.getName(); + } + + @Override + public ActionType getType() { + return impl.getType(); + } + + @Override + public ActionDesc getDesc() { + return impl.getDesc(); + } + + @Override + public Map getInputSchema() { + return impl.getInputSchema(); + } + + @Override + public Map getOutputSchema() { + return impl.getOutputSchema(); + } + + @Override + public Map getMetadata() { + return impl.getMetadata(); + } + + @Override + public AgentOutput run(ActionContext ctx, AgentInput input) throws GenkitException { + return impl.run(ctx, input); + } + + @Override + public AgentOutput run(ActionContext ctx, AgentInput input, Consumer cb) + throws GenkitException { + return impl.run(ctx, input, cb); + } + + @Override + public JsonNode runJson(ActionContext ctx, JsonNode input, Consumer streamCallback) + throws GenkitException { + return impl.runJson(ctx, input, streamCallback); + } + + @Override + public ActionRunResult runJsonWithTelemetry( + ActionContext ctx, JsonNode input, Consumer streamCallback) throws GenkitException { + return impl.runJsonWithTelemetry(ctx, input, streamCallback); + } + + @Override + public AgentOutput runBidi( + ActionContext ctx, + AgentInit init, + InputSource inputs, + Consumer streamCallback) + throws GenkitException { + return impl.runBidi(ctx, init, inputs, streamCallback); + } + + @Override + public JsonNode runBidiJson( + ActionContext ctx, + JsonNode init, + InputSource inputs, + Consumer streamCallback) + throws GenkitException { + return impl.runBidiJson(ctx, init, inputs, streamCallback); + } + + @Override + public ActionRunResult runBidiJsonWithTelemetry( + ActionContext ctx, + JsonNode init, + InputSource inputs, + Consumer streamCallback) + throws GenkitException { + return impl.runBidiJsonWithTelemetry(ctx, init, inputs, streamCallback); + } + + @Override + public void register(Registry registry) { + impl.register(registry); + if (snapshotAction != null) { + snapshotAction.register(registry); + } + if (abortAction != null) { + abortAction.register(registry); + } + } + + // ── AgentApi typed facades ──────────────────────────────────────────────────── + + @Override + public SessionSnapshot getSnapshotData(GetSnapshotRequest req) { + if (store == null) { + return null; + } + GetSnapshotOptions.Builder optsBuilder = GetSnapshotOptions.builder(); + if (req.getSnapshotId() != null) { + optsBuilder.snapshotId(req.getSnapshotId()); + } + if (req.getSessionId() != null) { + optsBuilder.sessionId(req.getSessionId()); + } + return store.getSnapshot(optsBuilder.build()); + } + + @Override + public SnapshotStatus abort(String snapshotId) { + if (!(store instanceof SnapshotSubscriber)) { + return null; + } + final SnapshotStatus[] resultStatus = {null}; + store.saveSnapshot( + snapshotId, + existing -> { + if (existing == null) { + resultStatus[0] = null; + return null; + } + if (existing.getStatus() != SnapshotStatus.PENDING) { + resultStatus[0] = existing.getStatus(); + return existing; + } + existing.setStatus(SnapshotStatus.ABORTED); + resultStatus[0] = SnapshotStatus.ABORTED; + return existing; + }, + opts); + // In addition to the store-level status mutation above (the durable record of the abort, + // observable by anyone polling getSnapshot), also flip the live in-memory abort signal for a + // still-running DETACHED turn, if one is currently registered under this snapshot id. This is + // the only case where a turn's snapshot id is knowable while the turn is still running (see + // DetachController); a foreground turn has no resolvable id until after it returns, so there is + // no reachable window to signal it here. + com.google.genkit.ai.agent.internal.PendingAbortRegistry.signal(snapshotId); + return resultStatus[0]; + } + + @Override + public AgentRef ref() { + return new AgentRef(name, description); + } + + // ── Chat client ─────────────────────────────────────────────────────────────── + + @Override + public AgentChat chat() { + return chat(new ActionContext(registry), null); + } + + @Override + public AgentChat chat(AgentInit init) { + return chat(new ActionContext(registry), init); + } + + @Override + public AgentChat loadChat(GetSnapshotRequest lookup) { + return loadChat(new ActionContext(registry), lookup); + } + + @Override + public AgentChat chat(ActionContext ctx) { + return chat(ctx, null); + } + + @Override + public AgentChat chat(ActionContext ctx, AgentInit init) { + return new AgentChat<>(new InProcessTransport<>(this, ctx), init); + } + + @Override + public AgentChat loadChat(ActionContext ctx, GetSnapshotRequest lookup) { + InProcessTransport transport = new InProcessTransport<>(this, ctx); + AgentChat chat = new AgentChat<>(transport, null); + SessionSnapshot snap = getSnapshotData(lookup); + chat.loadSnapshot(snap); + return chat; + } + + // ── Accessor methods ────────────────────────────────────────────────────────── + + /** + * Returns whether this agent uses server-managed session state. + * + * @return true if server-managed, false if client-managed + */ + public boolean serverManaged() { + return serverManaged; + } + + /** + * Returns the session store, or {@code null} for client-managed agents. + * + * @return the session store, or null + */ + public SessionStore store() { + return store; + } + + /** + * Returns the companion snapshot action, or {@code null} if not applicable. + * + * @return the snapshot action, or null + */ + public Action getSnapshotDataAction() { + return snapshotAction; + } + + /** + * Returns the companion abort action, or {@code null} if not applicable. + * + * @return the abort action, or null + */ + public Action abortAgentAction() { + return abortAction; + } +} diff --git a/ai/src/main/java/com/google/genkit/ai/agent/AgentAbortRequest.java b/ai/src/main/java/com/google/genkit/ai/agent/AgentAbortRequest.java new file mode 100644 index 000000000..11a6354a4 --- /dev/null +++ b/ai/src/main/java/com/google/genkit/ai/agent/AgentAbortRequest.java @@ -0,0 +1,78 @@ +/* + * 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. + * + * SPDX-License-Identifier: Apache-2.0 + */ + +package com.google.genkit.ai.agent; + +import com.fasterxml.jackson.annotation.JsonInclude; +import com.fasterxml.jackson.annotation.JsonProperty; + +/** AgentAbortRequest is the request body for aborting an agent execution. */ +@JsonInclude(JsonInclude.Include.NON_NULL) +public class AgentAbortRequest { + + @JsonProperty("snapshotId") + private String snapshotId; + + /** Default constructor. */ + public AgentAbortRequest() {} + + private AgentAbortRequest(Builder builder) { + this.snapshotId = builder.snapshotId; + } + + /** + * Creates a builder for AgentAbortRequest. + * + * @return a new builder + */ + public static Builder builder() { + return new Builder(); + } + + /** + * Returns the snapshot ID to abort. + * + * @return the snapshot ID + */ + public String getSnapshotId() { + return snapshotId; + } + + /** + * Sets the snapshot ID to abort. + * + * @param snapshotId the snapshot ID + */ + public void setSnapshotId(String snapshotId) { + this.snapshotId = snapshotId; + } + + /** Builder for AgentAbortRequest. */ + public static class Builder { + private String snapshotId; + + public Builder snapshotId(String snapshotId) { + this.snapshotId = snapshotId; + return this; + } + + public AgentAbortRequest build() { + return new AgentAbortRequest(this); + } + } +} diff --git a/ai/src/main/java/com/google/genkit/ai/agent/AgentAbortResponse.java b/ai/src/main/java/com/google/genkit/ai/agent/AgentAbortResponse.java new file mode 100644 index 000000000..bc5ceed68 --- /dev/null +++ b/ai/src/main/java/com/google/genkit/ai/agent/AgentAbortResponse.java @@ -0,0 +1,106 @@ +/* + * 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. + * + * SPDX-License-Identifier: Apache-2.0 + */ + +package com.google.genkit.ai.agent; + +import com.fasterxml.jackson.annotation.JsonInclude; +import com.fasterxml.jackson.annotation.JsonProperty; + +/** AgentAbortResponse is the response body for an agent abort operation. */ +@JsonInclude(JsonInclude.Include.NON_NULL) +public class AgentAbortResponse { + + @JsonProperty("snapshotId") + private String snapshotId; + + @JsonProperty("status") + private SnapshotStatus status; + + /** Default constructor. */ + public AgentAbortResponse() {} + + private AgentAbortResponse(Builder builder) { + this.snapshotId = builder.snapshotId; + this.status = builder.status; + } + + /** + * Creates a builder for AgentAbortResponse. + * + * @return a new builder + */ + public static Builder builder() { + return new Builder(); + } + + /** + * Returns the snapshot ID. + * + * @return the snapshot ID + */ + public String getSnapshotId() { + return snapshotId; + } + + /** + * Sets the snapshot ID. + * + * @param snapshotId the snapshot ID + */ + public void setSnapshotId(String snapshotId) { + this.snapshotId = snapshotId; + } + + /** + * Returns the resulting snapshot status. + * + * @return the status + */ + public SnapshotStatus getStatus() { + return status; + } + + /** + * Sets the resulting snapshot status. + * + * @param status the status + */ + public void setStatus(SnapshotStatus status) { + this.status = status; + } + + /** Builder for AgentAbortResponse. */ + public static class Builder { + private String snapshotId; + private SnapshotStatus status; + + public Builder snapshotId(String snapshotId) { + this.snapshotId = snapshotId; + return this; + } + + public Builder status(SnapshotStatus status) { + this.status = status; + return this; + } + + public AgentAbortResponse build() { + return new AgentAbortResponse(this); + } + } +} diff --git a/ai/src/main/java/com/google/genkit/ai/agent/AgentApi.java b/ai/src/main/java/com/google/genkit/ai/agent/AgentApi.java new file mode 100644 index 000000000..1e13b97bc --- /dev/null +++ b/ai/src/main/java/com/google/genkit/ai/agent/AgentApi.java @@ -0,0 +1,111 @@ +/* + * 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. + * + * SPDX-License-Identifier: Apache-2.0 + */ + +package com.google.genkit.ai.agent; + +import com.google.genkit.core.ActionContext; + +/** + * AgentApi is the typed facade interface exposed by an agent, providing snapshot retrieval, abort, + * identification, and chat-client creation. + * + * @param the type of custom session state + */ +public interface AgentApi { + + /** + * Creates a fresh {@link AgentChat} for a new conversation, using a default action context built + * from the registry the agent was defined in. + * + *

This is the convenience overload for the common case. Use {@link #chat(ActionContext)} when + * you need to supply request-scoped context (e.g. auth / user info that tools should observe). + * + * @return a new chat + */ + AgentChat chat(); + + /** + * Creates an {@link AgentChat} seeded from {@code init} (snapshotId / sessionId / inline state), + * using a default action context built from the registry the agent was defined in. + * + * @param init the seed init (may be {@code null} for a fresh chat) + * @return a new chat + */ + AgentChat chat(AgentInit init); + + /** + * Loads a snapshot and hydrates an {@link AgentChat} whose next send resumes it, using a default + * action context built from the registry the agent was defined in. + * + * @param lookup which snapshot to load (by snapshotId or sessionId) + * @return a hydrated chat + */ + AgentChat loadChat(GetSnapshotRequest lookup); + + /** + * Creates a fresh {@link AgentChat} for a new conversation. + * + * @param ctx the action context used to run each turn + * @return a new chat + */ + AgentChat chat(ActionContext ctx); + + /** + * Creates an {@link AgentChat} seeded from {@code init} (snapshotId / sessionId / inline state). + * + * @param ctx the action context used to run each turn + * @param init the seed init (may be {@code null} for a fresh chat) + * @return a new chat + */ + AgentChat chat(ActionContext ctx, AgentInit init); + + /** + * Loads a snapshot and hydrates an {@link AgentChat} whose next send resumes it. + * + * @param ctx the action context used to run each turn + * @param lookup which snapshot to load (by snapshotId or sessionId) + * @return a hydrated chat + */ + AgentChat loadChat(ActionContext ctx, GetSnapshotRequest lookup); + + /** + * Retrieves a session snapshot by snapshot ID or session ID. + * + * @param req the request specifying which snapshot to retrieve + * @return the matching session snapshot, or {@code null} if not found + */ + SessionSnapshot getSnapshotData(GetSnapshotRequest req); + + /** + * Attempts to abort a snapshot that is currently in {@link SnapshotStatus#PENDING} state. + * + *

If the snapshot is already in a terminal state, returns that state unchanged. If the + * snapshot does not exist, returns {@code null}. + * + * @param snapshotId the ID of the snapshot to abort + * @return the resulting {@link SnapshotStatus}, or {@code null} if the snapshot was not found + */ + SnapshotStatus abort(String snapshotId); + + /** + * Returns a lightweight reference to this agent. + * + * @return an {@link AgentRef} carrying the agent's name and description + */ + AgentRef ref(); +} diff --git a/ai/src/main/java/com/google/genkit/ai/agent/AgentChat.java b/ai/src/main/java/com/google/genkit/ai/agent/AgentChat.java new file mode 100644 index 000000000..30704ac1e --- /dev/null +++ b/ai/src/main/java/com/google/genkit/ai/agent/AgentChat.java @@ -0,0 +1,409 @@ +/* + * 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. + * + * SPDX-License-Identifier: Apache-2.0 + */ + +package com.google.genkit.ai.agent; + +import com.fasterxml.jackson.databind.JsonNode; +import com.google.genkit.ai.Message; +import com.google.genkit.ai.Part; +import com.google.genkit.core.JsonUtils; +import com.google.genkit.core.jsonpatch.JsonPatch; +import java.util.ArrayList; +import java.util.Collections; +import java.util.List; +import java.util.function.Consumer; + +/** + * AgentChat is the ergonomic programmatic client for driving an agent across multiple turns while + * carrying session state automatically. + * + *

Each {@link #send} is exactly one invocation of the agent's bidi action (one-turn-per-send). + * The chat tracks the evolving {@code snapshotId} / {@code sessionId} / custom state / messages / + * artifacts so the next send resumes correctly. Resume is carried in the per-turn {@link + * AgentInit}: + * + *

    + *
  • Server-managed agents persist state server-side; the init carries the latest {@code + * snapshotId} (and {@code sessionId} once known) so the server resumes from there. + *
  • Client-managed agents do not persist; the init carries the full {@link SessionState} + * (sessionId, messages, custom, artifacts) so the agent rehydrates from it each turn. + *
+ * + *

Instances are created via {@link Agent#chat(com.google.genkit.core.ActionContext)} / {@link + * Agent#chat(com.google.genkit.core.ActionContext, AgentInit)} / {@link + * Agent#loadChat(com.google.genkit.core.ActionContext, GetSnapshotRequest)}. + * + *

Not thread-safe: a chat is a single conversation and is expected to be driven from one thread. + * + * @param the type of custom session state + */ +public final class AgentChat { + + private final AgentTransport transport; + private final boolean serverManaged; + + private String snapshotId; + private String sessionId; + private S state; + private final List messages = new ArrayList<>(); + private final List artifacts = new ArrayList<>(); + + /** + * Creates a chat over any transport. External code (e.g. {@code RemoteAgent}) uses this factory + * so they do not need to be in the same package. + * + * @param the type of custom session state + * @param transport the transport that runs each turn (must not be null) + * @param init optional seed init; may be {@code null} for a fresh chat + * @return a new {@link AgentChat} backed by {@code transport} + */ + public static AgentChat over(AgentTransport transport, AgentInit init) { + return new AgentChat<>(transport, init); + } + + /** + * Constructs a chat. Called by {@link Agent} and {@link #over}. + * + * @param transport the transport that runs each turn (must not be null) + * @param init optional seed init: snapshotId / sessionId / inline state to start from; may be + * {@code null} for a fresh chat + */ + AgentChat(AgentTransport transport, AgentInit init) { + this.transport = transport; + this.serverManaged = transport.serverManaged(); + if (init != null) { + this.snapshotId = init.getSnapshotId(); + this.sessionId = init.getSessionId(); + SessionState seed = init.getState(); + if (seed != null) { + hydrateFromState(seed); + } + } + } + + /** Hydrates this chat's tracked state from a {@link SessionState}. */ + private void hydrateFromState(SessionState seed) { + if (seed.getSessionId() != null) { + this.sessionId = seed.getSessionId(); + } + this.state = seed.getCustom(); + this.messages.clear(); + if (seed.getMessages() != null) { + this.messages.addAll(seed.getMessages()); + } + this.artifacts.clear(); + if (seed.getArtifacts() != null) { + this.artifacts.addAll(seed.getArtifacts()); + } + } + + /** + * Hydrates this chat from a snapshot (used by {@code loadChat}). + * + * @param snap the snapshot to load (may be {@code null}) + */ + void loadSnapshot(SessionSnapshot snap) { + if (snap == null) { + return; + } + this.snapshotId = snap.getSnapshotId(); + if (snap.getSessionId() != null) { + this.sessionId = snap.getSessionId(); + } + if (snap.getState() != null) { + hydrateFromState(snap.getState()); + } + } + + // ── send ─────────────────────────────────────────────────────────────────────── + + /** + * Sends a user text message and returns the turn's response. + * + * @param text the user message text + * @return the response + */ + public AgentResponse send(String text) { + return send(AgentInput.builder().message(Message.user(text)).build()); + } + + /** + * Sends a fully-formed input and returns the turn's response. + * + * @param input the turn input + * @return the response + */ + public AgentResponse send(AgentInput input) { + return sendStream(input, c -> {}); + } + + /** + * Sends a user text message, streaming chunks to {@code onChunk}, and returns the final response. + * + * @param text the user message text + * @param onChunk per-chunk callback (may be a no-op) + * @return the response + */ + public AgentResponse sendStream(String text, Consumer> onChunk) { + return sendStream(AgentInput.builder().message(Message.user(text)).build(), onChunk); + } + + /** + * Sends a fully-formed input, streaming chunks to {@code onChunk}, and returns the final + * response. + * + * @param input the turn input + * @param onChunk per-chunk callback (may be a no-op) + * @return the response + */ + public AgentResponse sendStream(AgentInput input, Consumer> onChunk) { + // Optimistically track the user message we send (server-managed agents do not echo full state). + Message userMessage = input != null ? input.getMessage() : null; + + // Running custom-state document for chunk.custom(): seeded from the chat's current custom + // state. The first customPatch of a turn is a whole-document replace, so this tracks the + // authoritative custom state for server-managed agents (which do not echo full state inline). + final JsonNode[] runningCustom = {JsonUtils.toJsonNode(state)}; + final boolean[] sawCustomPatch = {false}; + + AgentInit init = buildInit(); + AgentOutput output = + transport.runTurn( + input, + init, + raw -> { + S chunkCustom = null; + if (raw != null && raw.getCustomPatch() != null) { + JsonNode patched = JsonPatch.apply(runningCustom[0], raw.getCustomPatch()); + runningCustom[0] = patched; + sawCustomPatch[0] = true; + chunkCustom = deserializeCustom(patched); + } + if (onChunk != null) { + onChunk.accept(new AgentChunk<>(raw, chunkCustom)); + } + }); + + S streamedCustom = sawCustomPatch[0] ? deserializeCustom(runningCustom[0]) : null; + applyOutput(output, userMessage, streamedCustom, sawCustomPatch[0]); + return new AgentResponse<>(output, state); + } + + /** + * Resumes a paused turn by responding to its interrupts. + * + * @param respond the response parts + * @return the response + */ + public AgentResponse resume(List respond) { + ToolResume tr = ToolResume.builder().respond(respond).build(); + return send(AgentInput.builder().resume(tr).build()); + } + + /** + * Resumes a paused turn by restarting its interrupted tool requests. + * + *

Each restart part is a tool-request part (typically produced by {@code Tool.restart(...)}) + * carrying the resumed metadata; the tool is re-executed with that metadata so a restart-aware + * handler can observe its resumed status via {@code ActionContext.isResumed()}/{@code + * getResumed()}. Sibling to {@link #resume(List)}. + * + * @param restart the restart tool-request parts + * @return the response + */ + public AgentResponse restart(List restart) { + ToolResume tr = ToolResume.builder().restart(restart).build(); + return send(AgentInput.builder().resume(tr).build()); + } + + // ── lifecycle ──────────────────────────────────────────────────────────────── + + /** + * Aborts the latest pending snapshot for this chat. + * + * @return the resulting status, or {@code null} if there is no snapshot / abort is unsupported + */ + public SnapshotStatus abort() { + if (snapshotId == null) { + return null; + } + return transport.abort(snapshotId); + } + + // ── accessors ──────────────────────────────────────────────────────────────── + + /** + * Returns the latest snapshot ID (server-managed agents). + * + * @return the snapshot ID, or {@code null} before the first turn / for client-managed agents + */ + public String snapshotId() { + return snapshotId; + } + + /** + * Returns the session ID. + * + * @return the session ID, or {@code null} before it is known + */ + public String sessionId() { + return sessionId; + } + + /** + * Returns the current custom session state. + * + * @return the custom state, or {@code null} + */ + public S state() { + return state; + } + + /** + * Returns the accumulated conversation messages. + * + * @return an unmodifiable view of the messages + */ + public List messages() { + return Collections.unmodifiableList(messages); + } + + /** + * Returns the accumulated artifacts. + * + * @return an unmodifiable view of the artifacts + */ + public List artifacts() { + return Collections.unmodifiableList(artifacts); + } + + // ── internals ──────────────────────────────────────────────────────────────── + + /** + * Builds the {@link AgentInit} for the next turn from the chat's tracked state. + * + *

Server-managed: carry {@code snapshotId} (and {@code sessionId} when known) so the server + * resumes; first turn → empty init. Client-managed: carry the full current {@link SessionState}. + */ + private AgentInit buildInit() { + if (serverManaged) { + if (snapshotId == null && sessionId == null) { + return null; // fresh session + } + AgentInit.Builder b = AgentInit.builder(); + if (snapshotId != null) { + b.snapshotId(snapshotId); + } + if (sessionId != null) { + b.sessionId(sessionId); + } + return b.build(); + } + // client-managed: round-trip the full state + if (sessionId == null && messages.isEmpty() && state == null && artifacts.isEmpty()) { + return null; // fresh session + } + SessionState current = + SessionState.builder() + .sessionId(sessionId) + .messages(new ArrayList<>(messages)) + .custom(state) + .artifacts(new ArrayList<>(artifacts)) + .build(); + return AgentInit.builder().state(current).build(); + } + + /** + * Folds a turn's output back into the chat's tracked state so the next send resumes correctly. + * + * @param output the turn output + * @param userMessage the user message that was sent this turn (for server-managed history + * tracking); may be {@code null} + * @param streamedCustom custom state reconstructed from this turn's customPatch stream; used for + * server-managed agents which do not echo full state inline; may be {@code null} + * @param sawCustomPatch whether any customPatch was observed during the turn + */ + private void applyOutput( + AgentOutput output, Message userMessage, S streamedCustom, boolean sawCustomPatch) { + if (output == null) { + return; + } + if (output.getSnapshotId() != null) { + this.snapshotId = output.getSnapshotId(); + } + if (output.getSessionId() != null) { + this.sessionId = output.getSessionId(); + } + + SessionState outState = output.getState(); + if (outState != null) { + // Client-managed: the agent echoes the authoritative full state. + hydrateFromState(outState); + } else { + // Server-managed: no inline state. Track history locally by appending the user message we + // sent and the assistant message the agent returned, and merge any artifacts. + if (userMessage != null) { + this.messages.add(userMessage); + } + if (output.getMessage() != null) { + this.messages.add(output.getMessage()); + } + if (output.getArtifacts() != null) { + mergeArtifacts(output.getArtifacts()); + } + // The customPatch stream is authoritative for custom state in server-managed mode. + if (sawCustomPatch) { + this.state = streamedCustom; + } + } + } + + /** Merges turn artifacts into the tracked list, replacing same-named entries. */ + private void mergeArtifacts(List incoming) { + for (Artifact a : incoming) { + if (a == null) { + continue; + } + int idx = -1; + for (int i = 0; i < artifacts.size(); i++) { + if (a.getName() != null && a.getName().equals(artifacts.get(i).getName())) { + idx = i; + break; + } + } + if (idx >= 0) { + artifacts.set(idx, a); + } else { + artifacts.add(a); + } + } + } + + /** Deserializes a custom-state JsonNode into {@code S} (raw-type round-trip). */ + @SuppressWarnings("unchecked") + private S deserializeCustom(JsonNode node) { + if (node == null || node.isNull() || node.isMissingNode()) { + return null; + } + try { + return (S) JsonUtils.getObjectMapper().treeToValue(node, Object.class); + } catch (Exception e) { + return null; + } + } +} diff --git a/ai/src/main/java/com/google/genkit/ai/agent/AgentChunk.java b/ai/src/main/java/com/google/genkit/ai/agent/AgentChunk.java new file mode 100644 index 000000000..39f87ad82 --- /dev/null +++ b/ai/src/main/java/com/google/genkit/ai/agent/AgentChunk.java @@ -0,0 +1,95 @@ +/* + * 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. + * + * SPDX-License-Identifier: Apache-2.0 + */ + +package com.google.genkit.ai.agent; + +import com.google.genkit.ai.ModelResponseChunk; + +/** + * AgentChunk is the ergonomic wrapper {@link AgentChat#sendStream} hands to the per-chunk callback. + * + *

It exposes the common slices of an {@link AgentStreamChunk}: streamed model text, an artifact + * update, and the post-patch custom state (the running client-side custom state with this + * chunk's {@code customPatch} already applied). The raw chunk is available via {@link #raw()}. + * + * @param the type of custom session state + */ +public final class AgentChunk { + + private final AgentStreamChunk raw; + private final S custom; + + /** + * Constructs an AgentChunk. + * + * @param raw the underlying stream chunk + * @param custom the custom state after applying this chunk's patch, or {@code null} if this chunk + * carried no custom patch + */ + AgentChunk(AgentStreamChunk raw, S custom) { + this.raw = raw; + this.custom = custom; + } + + /** + * Returns the streamed model text for this chunk. + * + * @return the chunk text, or {@code null} if this chunk carried no model content + */ + public String text() { + ModelResponseChunk mc = raw != null ? raw.getModelChunk() : null; + return mc != null ? mc.getText() : null; + } + + /** + * Returns the model response chunk. + * + * @return the model chunk, or {@code null} + */ + public ModelResponseChunk modelChunk() { + return raw != null ? raw.getModelChunk() : null; + } + + /** + * Returns the artifact carried by this chunk. + * + * @return the artifact, or {@code null} + */ + public Artifact artifact() { + return raw != null ? raw.getArtifact() : null; + } + + /** + * Returns the custom state after this chunk's {@code customPatch} was applied to the running + * client-side state. + * + * @return the post-patch custom state, or {@code null} if this chunk carried no custom patch + */ + public S custom() { + return custom; + } + + /** + * Returns the underlying raw stream chunk. + * + * @return the raw chunk + */ + public AgentStreamChunk raw() { + return raw; + } +} diff --git a/ai/src/main/java/com/google/genkit/ai/agent/AgentFinishReason.java b/ai/src/main/java/com/google/genkit/ai/agent/AgentFinishReason.java new file mode 100644 index 000000000..a602f0c9b --- /dev/null +++ b/ai/src/main/java/com/google/genkit/ai/agent/AgentFinishReason.java @@ -0,0 +1,90 @@ +/* + * 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. + * + * SPDX-License-Identifier: Apache-2.0 + */ + +package com.google.genkit.ai.agent; + +import com.fasterxml.jackson.annotation.JsonCreator; +import com.fasterxml.jackson.annotation.JsonValue; + +/** AgentFinishReason indicates why an agent finished execution. */ +public enum AgentFinishReason { + /** The agent finished normally due to a stop condition. */ + STOP("stop"), + + /** The agent finished due to reaching token length limits. */ + LENGTH("length"), + + /** The agent was blocked from proceeding. */ + BLOCKED("blocked"), + + /** The agent execution was interrupted. */ + INTERRUPTED("interrupted"), + + /** The agent finished for some other reason. */ + OTHER("other"), + + /** The agent finish reason is unknown. */ + UNKNOWN("unknown"), + + /** The agent was aborted. */ + ABORTED("aborted"), + + /** The agent was detached. */ + DETACHED("detached"), + + /** The agent failed. */ + FAILED("failed"); + + private final String value; + + AgentFinishReason(String value) { + this.value = value; + } + + /** + * Returns the string value of the agent finish reason. + * + * @return the agent finish reason string value + */ + @JsonValue + public String getValue() { + return value; + } + + /** + * Creates an AgentFinishReason from a string value. + * + * @param value the string value + * @return the corresponding AgentFinishReason + * @throws IllegalArgumentException if the value doesn't match any AgentFinishReason + */ + @JsonCreator + public static AgentFinishReason fromValue(String value) { + for (AgentFinishReason reason : values()) { + if (reason.value.equals(value)) { + return reason; + } + } + throw new IllegalArgumentException("Unknown agent finish reason: " + value); + } + + @Override + public String toString() { + return value; + } +} diff --git a/ai/src/main/java/com/google/genkit/ai/agent/AgentFn.java b/ai/src/main/java/com/google/genkit/ai/agent/AgentFn.java new file mode 100644 index 000000000..73565a1b6 --- /dev/null +++ b/ai/src/main/java/com/google/genkit/ai/agent/AgentFn.java @@ -0,0 +1,42 @@ +/* + * 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. + * + * SPDX-License-Identifier: Apache-2.0 + */ + +package com.google.genkit.ai.agent; + +/** + * Custom agent handler: runs one invocation's turn loop body for ONE turn at a time via the runner. + * + *

Implementations receive the {@link SessionRunner} (giving access to session state, messages, + * and artifacts) and an {@link AgentFnContext} (chunk emitter, abort signal). Task 4.5 ({@code + * defineCustomAgent}) wraps this interface. + * + * @param the type of the custom session state + */ +@FunctionalInterface +public interface AgentFn { + + /** + * Runs one turn of agent logic. + * + * @param sess the session runner for reading/mutating session state + * @param ctx per-invocation context (streaming and abort support) + * @return the agent result for this invocation + * @throws Exception if the agent fails; callers should propagate or handle appropriately + */ + AgentResult run(SessionRunner sess, AgentFnContext ctx) throws Exception; +} diff --git a/ai/src/main/java/com/google/genkit/ai/agent/AgentFnContext.java b/ai/src/main/java/com/google/genkit/ai/agent/AgentFnContext.java new file mode 100644 index 000000000..620852d55 --- /dev/null +++ b/ai/src/main/java/com/google/genkit/ai/agent/AgentFnContext.java @@ -0,0 +1,131 @@ +/* + * 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. + * + * SPDX-License-Identifier: Apache-2.0 + */ + +package com.google.genkit.ai.agent; + +import com.google.genkit.core.ActionContext; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.function.Consumer; + +/** + * Context handed to an {@link AgentFn} for one invocation. + * + *

Provides access to a stream-chunk emitter and an abort signal that allows the agent to detect + * cancellation. Stream emission of custom-patch / artifacts is wired in Task 4.3; here the {@code + * sendChunk} consumer is available for the agent to emit model chunks. + * + *

It also carries the run's {@link ActionContext} (via {@link #context()}) so the agent function + * can read the request-scoped user context (e.g. {@code {"auth": {...}}}) and forward it into the + * generate call so tools observe it. + */ +public final class AgentFnContext { + + private final Consumer sendChunk; + private final AtomicBoolean abortSignal; + private final ActionContext context; + private final ToolResume resume; + + /** + * Constructs an AgentFnContext without an ActionContext. + * + * @param sendChunk consumer that emits a streaming chunk to the caller; may be a no-op + * @param abortSignal shared flag; set to {@code true} by the runtime to request cancellation + */ + public AgentFnContext(Consumer sendChunk, AtomicBoolean abortSignal) { + this(sendChunk, abortSignal, null); + } + + /** + * Constructs an AgentFnContext. + * + * @param sendChunk consumer that emits a streaming chunk to the caller; may be a no-op + * @param abortSignal shared flag; set to {@code true} by the runtime to request cancellation + * @param context the run's ActionContext (carries the user context); may be null + */ + public AgentFnContext( + Consumer sendChunk, AtomicBoolean abortSignal, ActionContext context) { + this(sendChunk, abortSignal, context, null); + } + + /** + * Constructs an AgentFnContext with resume data. + * + * @param sendChunk consumer that emits a streaming chunk to the caller; may be a no-op + * @param abortSignal shared flag; set to {@code true} by the runtime to request cancellation + * @param context the run's ActionContext (carries the user context); may be null + * @param resume the current turn's tool-resume data, or {@code null} if this is not a resume turn + */ + public AgentFnContext( + Consumer sendChunk, + AtomicBoolean abortSignal, + ActionContext context, + ToolResume resume) { + this.sendChunk = sendChunk != null ? sendChunk : chunk -> {}; + this.abortSignal = abortSignal != null ? abortSignal : new AtomicBoolean(false); + this.context = context; + this.resume = resume; + } + + /** + * Returns the stream-chunk emitter for this invocation. + * + * @return the {@link AgentStreamChunk} consumer (never null) + */ + public Consumer sendChunk() { + return sendChunk; + } + + /** + * Returns the abort signal for this invocation. + * + * @return an {@link AtomicBoolean} that is {@code true} when the caller has requested + * cancellation + */ + public AtomicBoolean abortSignal() { + return abortSignal; + } + + /** + * Convenience method to check whether the abort signal has been set. + * + * @return {@code true} if the caller has requested cancellation + */ + public boolean isAborted() { + return abortSignal.get(); + } + + /** + * Returns the run's {@link ActionContext}, which carries the request-scoped user context. + * + * @return the ActionContext for this invocation, or null if none was provided + */ + public ActionContext context() { + return context; + } + + /** + * Returns this turn's tool-resume data, i.e. the {@code respond}/{@code restart} parts passed to + * {@code AgentChat.resume(...)} (or the raw {@code AgentInput.resume} field for lower-level + * callers). + * + * @return the {@link ToolResume} for this turn, or {@code null} if this is not a resume turn + */ + public ToolResume resume() { + return resume; + } +} diff --git a/ai/src/main/java/com/google/genkit/ai/agent/AgentInit.java b/ai/src/main/java/com/google/genkit/ai/agent/AgentInit.java new file mode 100644 index 000000000..14e832f2d --- /dev/null +++ b/ai/src/main/java/com/google/genkit/ai/agent/AgentInit.java @@ -0,0 +1,143 @@ +/* + * 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. + * + * SPDX-License-Identifier: Apache-2.0 + */ + +package com.google.genkit.ai.agent; + +import com.fasterxml.jackson.annotation.JsonInclude; +import com.fasterxml.jackson.annotation.JsonProperty; + +/** + * AgentInit represents the initialization data for an agent session. + * + * @param the type of custom state + */ +@JsonInclude(JsonInclude.Include.NON_NULL) +public class AgentInit { + + @JsonProperty("snapshotId") + private String snapshotId; + + @JsonProperty("sessionId") + private String sessionId; + + @JsonProperty("state") + private SessionState state; + + /** Default constructor. */ + public AgentInit() {} + + private AgentInit(Builder builder) { + this.snapshotId = builder.snapshotId; + this.sessionId = builder.sessionId; + this.state = builder.state; + } + + /** + * Creates a builder for AgentInit. + * + * @param the type of custom state + * @return a new builder + */ + public static Builder builder() { + return new Builder<>(); + } + + /** + * Returns the snapshot ID. + * + * @return the snapshot ID + */ + public String getSnapshotId() { + return snapshotId; + } + + /** + * Sets the snapshot ID. + * + * @param snapshotId the snapshot ID + */ + public void setSnapshotId(String snapshotId) { + this.snapshotId = snapshotId; + } + + /** + * Returns the session ID. + * + * @return the session ID + */ + public String getSessionId() { + return sessionId; + } + + /** + * Sets the session ID. + * + * @param sessionId the session ID + */ + public void setSessionId(String sessionId) { + this.sessionId = sessionId; + } + + /** + * Returns the initial session state. + * + * @return the session state + */ + public SessionState getState() { + return state; + } + + /** + * Sets the initial session state. + * + * @param state the session state + */ + public void setState(SessionState state) { + this.state = state; + } + + /** + * Builder for AgentInit. + * + * @param the type of custom state + */ + public static class Builder { + private String snapshotId; + private String sessionId; + private SessionState state; + + public Builder snapshotId(String snapshotId) { + this.snapshotId = snapshotId; + return this; + } + + public Builder sessionId(String sessionId) { + this.sessionId = sessionId; + return this; + } + + public Builder state(SessionState state) { + this.state = state; + return this; + } + + public AgentInit build() { + return new AgentInit<>(this); + } + } +} diff --git a/ai/src/main/java/com/google/genkit/ai/agent/AgentInput.java b/ai/src/main/java/com/google/genkit/ai/agent/AgentInput.java new file mode 100644 index 000000000..8603b2c01 --- /dev/null +++ b/ai/src/main/java/com/google/genkit/ai/agent/AgentInput.java @@ -0,0 +1,150 @@ +/* + * 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. + * + * SPDX-License-Identifier: Apache-2.0 + */ + +package com.google.genkit.ai.agent; + +import com.fasterxml.jackson.annotation.JsonInclude; +import com.fasterxml.jackson.annotation.JsonProperty; +import com.google.genkit.ai.Message; + +/** + * AgentInput represents the input to an agent turn. + * + *

The {@code detach} field is omitted when false (default). + */ +@JsonInclude(JsonInclude.Include.NON_NULL) +public class AgentInput { + + @JsonProperty("message") + private Message message; + + @JsonProperty("resume") + private ToolResume resume; + + /** detach is omitted when false; use primitive boolean with NON_DEFAULT include. */ + @JsonProperty("detach") + @JsonInclude(JsonInclude.Include.NON_DEFAULT) + private boolean detach; + + /** Default constructor. */ + public AgentInput() {} + + private AgentInput(Builder builder) { + this.message = builder.message; + this.resume = builder.resume; + this.detach = builder.detach; + } + + /** + * Creates a builder for AgentInput. + * + * @return a new builder + */ + public static Builder builder() { + return new Builder(); + } + + /** + * Returns the input message. + * + * @return the message + */ + public Message getMessage() { + return message; + } + + /** + * Sets the input message. + * + * @param message the message + */ + public void setMessage(Message message) { + this.message = message; + } + + /** + * Returns the tool resume data. + * + * @return the resume, or null if not resuming + */ + public ToolResume getResume() { + return resume; + } + + /** + * Sets the tool resume data. + * + * @param resume the resume + */ + public void setResume(ToolResume resume) { + this.resume = resume; + } + + /** + * Returns whether the agent should detach. + * + * @return true if detaching, false otherwise + */ + public boolean isDetach() { + return detach; + } + + /** + * Gets the detach flag (alias for isDetach). + * + * @return true if detaching, false otherwise + */ + public boolean getDetach() { + return detach; + } + + /** + * Sets the detach flag. Omitted from JSON when false. + * + * @param detach true to detach + */ + public void setDetach(boolean detach) { + this.detach = detach; + } + + /** Builder for AgentInput. */ + public static class Builder { + private Message message; + private ToolResume resume; + private boolean detach; + + public Builder message(Message message) { + this.message = message; + return this; + } + + public Builder resume(ToolResume resume) { + this.resume = resume; + return this; + } + + public Builder detach(boolean detach) { + this.detach = detach; + return this; + } + + public AgentInput build() { + return new AgentInput(this); + } + } +} diff --git a/ai/src/main/java/com/google/genkit/ai/agent/AgentInterrupt.java b/ai/src/main/java/com/google/genkit/ai/agent/AgentInterrupt.java new file mode 100644 index 000000000..f89934b6a --- /dev/null +++ b/ai/src/main/java/com/google/genkit/ai/agent/AgentInterrupt.java @@ -0,0 +1,88 @@ +/* + * 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. + * + * SPDX-License-Identifier: Apache-2.0 + */ + +package com.google.genkit.ai.agent; + +import com.google.genkit.ai.Part; +import com.google.genkit.ai.ToolRequest; + +/** + * AgentInterrupt represents a tool request that paused agent execution awaiting a caller response. + * + *

Extracted from the final message's tool-request parts when a turn finishes with {@link + * AgentFinishReason#INTERRUPTED}. The caller resolves an interrupt by sending a {@link + * ToolResume#getRespond()} via {@link AgentChat#resume(java.util.List)}. + */ +public final class AgentInterrupt { + + private final String name; + private final Object input; + private final Part part; + + /** + * Constructs an AgentInterrupt. + * + * @param name the tool name + * @param input the tool input + * @param part the originating tool-request part (preserved for resume correlation) + */ + public AgentInterrupt(String name, Object input, Part part) { + this.name = name; + this.input = input; + this.part = part; + } + + /** + * Returns the interrupted tool's name. + * + * @return the tool name, or {@code null} if unknown + */ + public String name() { + return name; + } + + /** + * Returns the interrupted tool's input. + * + * @return the tool input, or {@code null} + */ + public Object input() { + return input; + } + + /** + * Returns the originating tool-request part. + * + * @return the part, or {@code null} + */ + public Part part() { + return part; + } + + /** + * Builds an interrupt from a tool-request part. + * + * @param part a part whose {@link Part#getToolRequest()} is non-null + * @return a new interrupt + */ + static AgentInterrupt fromPart(Part part) { + ToolRequest tr = part.getToolRequest(); + return new AgentInterrupt( + tr != null ? tr.getName() : null, tr != null ? tr.getInput() : null, part); + } +} diff --git a/ai/src/main/java/com/google/genkit/ai/agent/AgentMetadata.java b/ai/src/main/java/com/google/genkit/ai/agent/AgentMetadata.java new file mode 100644 index 000000000..b08ddbb21 --- /dev/null +++ b/ai/src/main/java/com/google/genkit/ai/agent/AgentMetadata.java @@ -0,0 +1,140 @@ +/* + * 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. + * + * SPDX-License-Identifier: Apache-2.0 + */ + +package com.google.genkit.ai.agent; + +import com.fasterxml.jackson.annotation.JsonInclude; +import com.fasterxml.jackson.annotation.JsonProperty; +import java.util.Map; + +/** + * AgentMetadata describes the capabilities and configuration of an agent endpoint. + * + *

{@code stateManagement} is the string {@code "server"} or {@code "client"}. + */ +@JsonInclude(JsonInclude.Include.NON_NULL) +public class AgentMetadata { + + @JsonProperty("stateManagement") + private String stateManagement; + + /** abortable is a primitive boolean — always serialized (not omitted when false). */ + @JsonProperty("abortable") + private boolean abortable; + + @JsonProperty("stateSchema") + private Map stateSchema; + + /** Default constructor. */ + public AgentMetadata() {} + + private AgentMetadata(Builder builder) { + this.stateManagement = builder.stateManagement; + this.abortable = builder.abortable; + this.stateSchema = builder.stateSchema; + } + + /** + * Creates a builder for AgentMetadata. + * + * @return a new builder + */ + public static Builder builder() { + return new Builder(); + } + + /** + * Returns the state management mode. + * + * @return {@code "server"} or {@code "client"} + */ + public String getStateManagement() { + return stateManagement; + } + + /** + * Sets the state management mode. + * + * @param stateManagement {@code "server"} or {@code "client"} + */ + public void setStateManagement(String stateManagement) { + this.stateManagement = stateManagement; + } + + /** + * Returns whether the agent supports abort. + * + * @return true if abortable + */ + public boolean isAbortable() { + return abortable; + } + + /** + * Sets whether the agent supports abort. + * + * @param abortable true if abortable + */ + public void setAbortable(boolean abortable) { + this.abortable = abortable; + } + + /** + * Returns the JSON schema for the agent state, if provided. + * + * @return the state schema, or null if not set + */ + public Map getStateSchema() { + return stateSchema; + } + + /** + * Sets the JSON schema for the agent state. + * + * @param stateSchema the state schema + */ + public void setStateSchema(Map stateSchema) { + this.stateSchema = stateSchema; + } + + /** Builder for AgentMetadata. */ + public static class Builder { + private String stateManagement; + private boolean abortable; + private Map stateSchema; + + public Builder stateManagement(String stateManagement) { + this.stateManagement = stateManagement; + return this; + } + + public Builder abortable(boolean abortable) { + this.abortable = abortable; + return this; + } + + public Builder stateSchema(Map stateSchema) { + this.stateSchema = stateSchema; + return this; + } + + public AgentMetadata build() { + return new AgentMetadata(this); + } + } +} diff --git a/ai/src/main/java/com/google/genkit/ai/agent/AgentOutput.java b/ai/src/main/java/com/google/genkit/ai/agent/AgentOutput.java new file mode 100644 index 000000000..3227b50b3 --- /dev/null +++ b/ai/src/main/java/com/google/genkit/ai/agent/AgentOutput.java @@ -0,0 +1,258 @@ +/* + * 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. + * + * SPDX-License-Identifier: Apache-2.0 + */ + +package com.google.genkit.ai.agent; + +import com.fasterxml.jackson.annotation.JsonInclude; +import com.fasterxml.jackson.annotation.JsonProperty; +import com.google.genkit.ai.Message; +import java.util.ArrayList; +import java.util.List; + +/** + * AgentOutput represents the output of an agent turn. + * + * @param the type of custom state + */ +@JsonInclude(JsonInclude.Include.NON_NULL) +public class AgentOutput { + + @JsonProperty("sessionId") + private String sessionId; + + @JsonProperty("snapshotId") + private String snapshotId; + + @JsonProperty("state") + private SessionState state; + + @JsonProperty("message") + private Message message; + + @JsonProperty("artifacts") + private List artifacts; + + @JsonProperty("finishReason") + private AgentFinishReason finishReason; + + @JsonProperty("error") + private RuntimeError error; + + /** Default constructor. */ + public AgentOutput() {} + + private AgentOutput(Builder builder) { + this.sessionId = builder.sessionId; + this.snapshotId = builder.snapshotId; + this.state = builder.state; + this.message = builder.message; + this.artifacts = builder.artifacts; + this.finishReason = builder.finishReason; + this.error = builder.error; + } + + /** + * Creates a builder for AgentOutput. + * + * @param the type of custom state + * @return a new builder + */ + public static Builder builder() { + return new Builder<>(); + } + + /** + * Returns the session ID. + * + * @return the session ID + */ + public String getSessionId() { + return sessionId; + } + + /** + * Sets the session ID. + * + * @param sessionId the session ID + */ + public void setSessionId(String sessionId) { + this.sessionId = sessionId; + } + + /** + * Returns the snapshot ID. + * + * @return the snapshot ID + */ + public String getSnapshotId() { + return snapshotId; + } + + /** + * Sets the snapshot ID. + * + * @param snapshotId the snapshot ID + */ + public void setSnapshotId(String snapshotId) { + this.snapshotId = snapshotId; + } + + /** + * Returns the session state. + * + * @return the state + */ + public SessionState getState() { + return state; + } + + /** + * Sets the session state. + * + * @param state the state + */ + public void setState(SessionState state) { + this.state = state; + } + + /** + * Returns the output message. + * + * @return the message + */ + public Message getMessage() { + return message; + } + + /** + * Sets the output message. + * + * @param message the message + */ + public void setMessage(Message message) { + this.message = message; + } + + /** + * Returns the output artifacts. + * + * @return the artifacts + */ + public List getArtifacts() { + return artifacts; + } + + /** + * Sets the output artifacts. + * + * @param artifacts the artifacts + */ + public void setArtifacts(List artifacts) { + this.artifacts = artifacts != null ? new ArrayList<>(artifacts) : null; + } + + /** + * Returns the finish reason. + * + * @return the finish reason + */ + public AgentFinishReason getFinishReason() { + return finishReason; + } + + /** + * Sets the finish reason. + * + * @param finishReason the finish reason + */ + public void setFinishReason(AgentFinishReason finishReason) { + this.finishReason = finishReason; + } + + /** + * Returns the runtime error, if any. + * + * @return the error, or null if no error + */ + public RuntimeError getError() { + return error; + } + + /** + * Sets the runtime error. + * + * @param error the error + */ + public void setError(RuntimeError error) { + this.error = error; + } + + /** + * Builder for AgentOutput. + * + * @param the type of custom state + */ + public static class Builder { + private String sessionId; + private String snapshotId; + private SessionState state; + private Message message; + private List artifacts; + private AgentFinishReason finishReason; + private RuntimeError error; + + public Builder sessionId(String sessionId) { + this.sessionId = sessionId; + return this; + } + + public Builder snapshotId(String snapshotId) { + this.snapshotId = snapshotId; + return this; + } + + public Builder state(SessionState state) { + this.state = state; + return this; + } + + public Builder message(Message message) { + this.message = message; + return this; + } + + public Builder artifacts(List artifacts) { + this.artifacts = artifacts; + return this; + } + + public Builder finishReason(AgentFinishReason finishReason) { + this.finishReason = finishReason; + return this; + } + + public Builder error(RuntimeError error) { + this.error = error; + return this; + } + + public AgentOutput build() { + return new AgentOutput<>(this); + } + } +} diff --git a/ai/src/main/java/com/google/genkit/ai/agent/AgentRef.java b/ai/src/main/java/com/google/genkit/ai/agent/AgentRef.java new file mode 100644 index 000000000..ddf4c8290 --- /dev/null +++ b/ai/src/main/java/com/google/genkit/ai/agent/AgentRef.java @@ -0,0 +1,60 @@ +/* + * 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. + * + * SPDX-License-Identifier: Apache-2.0 + */ + +package com.google.genkit.ai.agent; + +/** + * AgentRef is a lightweight reference to an agent, carrying only its name and description. + * + *

Returned by {@link AgentApi#ref()} so callers can identify the agent without holding the full + * implementation object. + */ +public final class AgentRef { + + private final String name; + private final String description; + + /** + * Constructs an AgentRef. + * + * @param name the agent's registered name + * @param description the agent's human-readable description; may be {@code null} + */ + public AgentRef(String name, String description) { + this.name = name; + this.description = description; + } + + /** + * Returns the agent name. + * + * @return the agent name + */ + public String getName() { + return name; + } + + /** + * Returns the agent description. + * + * @return the description, or {@code null} if not set + */ + public String getDescription() { + return description; + } +} diff --git a/ai/src/main/java/com/google/genkit/ai/agent/AgentResponse.java b/ai/src/main/java/com/google/genkit/ai/agent/AgentResponse.java new file mode 100644 index 000000000..1c75c29d9 --- /dev/null +++ b/ai/src/main/java/com/google/genkit/ai/agent/AgentResponse.java @@ -0,0 +1,179 @@ +/* + * 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. + * + * SPDX-License-Identifier: Apache-2.0 + */ + +package com.google.genkit.ai.agent; + +import com.google.genkit.ai.Message; +import com.google.genkit.ai.Part; +import java.util.ArrayList; +import java.util.Collections; +import java.util.List; + +/** + * AgentResponse is the ergonomic wrapper {@link AgentChat#send} returns for one completed turn. + * + *

It surfaces the assistant {@link Message} and convenience views derived from it (text, + * tool-request parts, interrupts) alongside the turn's identifiers, custom state, artifacts, finish + * reason, and the raw {@link AgentOutput}. + * + * @param the type of custom session state + */ +public final class AgentResponse { + + private final AgentOutput raw; + private final S state; + + /** + * Constructs an AgentResponse. + * + * @param raw the underlying agent output for the turn + * @param state the custom state after the turn (server-managed agents do not echo full state, so + * the chat supplies its tracked copy here); may be {@code null} + */ + AgentResponse(AgentOutput raw, S state) { + this.raw = raw; + this.state = state; + } + + /** + * Returns the assistant message produced by the turn. + * + * @return the message, or {@code null} if the turn produced none + */ + public Message message() { + return raw != null ? raw.getMessage() : null; + } + + /** + * Returns the concatenated text of the assistant message's text parts. + * + * @return the text (empty string if there is no message / no text parts) + */ + public String text() { + Message m = message(); + return m != null ? m.getText() : ""; + } + + /** + * Returns the tool-request parts of the assistant message. + * + * @return an unmodifiable list of tool-request parts (possibly empty) + */ + public List toolRequests() { + Message m = message(); + if (m == null || m.getContent() == null) { + return Collections.emptyList(); + } + List out = new ArrayList<>(); + for (Part p : m.getContent()) { + if (p != null && p.isToolRequest()) { + out.add(p); + } + } + return Collections.unmodifiableList(out); + } + + /** + * Returns the interrupts for this turn. + * + *

Best-effort: when the turn finished with {@link AgentFinishReason#INTERRUPTED}, every + * tool-request part of the assistant message is reported as an interrupt. The wire {@link Part} + * does not currently carry an explicit interrupt flag, so all pending tool requests are treated + * as interrupts. When the turn did not finish interrupted, this returns an empty list. + * + * @return an unmodifiable list of interrupts (possibly empty) + */ + public List interrupts() { + if (finishReason() != AgentFinishReason.INTERRUPTED) { + return Collections.emptyList(); + } + List out = new ArrayList<>(); + for (Part p : toolRequests()) { + out.add(AgentInterrupt.fromPart(p)); + } + return Collections.unmodifiableList(out); + } + + /** + * Returns the turn's finish reason. + * + * @return the finish reason, or {@code null} + */ + public AgentFinishReason finishReason() { + return raw != null ? raw.getFinishReason() : null; + } + + /** + * Returns the snapshot ID produced by the turn (server-managed agents). + * + * @return the snapshot ID, or {@code null} + */ + public String snapshotId() { + return raw != null ? raw.getSnapshotId() : null; + } + + /** + * Returns the session ID. + * + * @return the session ID, or {@code null} + */ + public String sessionId() { + return raw != null ? raw.getSessionId() : null; + } + + /** + * Returns the custom session state after the turn. + * + * @return the custom state, or {@code null} + */ + public S custom() { + return state; + } + + /** + * Returns the full session state at the end of the turn. + * + *

For client-managed agents this is the inline state echoed by the agent. For server-managed + * agents the agent does not echo full state inline, so this returns {@code null}; use {@link + * #snapshotId()} with {@link AgentChat} / {@code getSnapshotData} to read server-side state. + * + * @return the session state, or {@code null} + */ + public SessionState state() { + return raw != null ? raw.getState() : null; + } + + /** + * Returns the turn's artifacts. + * + * @return an unmodifiable list of artifacts (possibly empty) + */ + public List artifacts() { + List a = raw != null ? raw.getArtifacts() : null; + return a != null ? Collections.unmodifiableList(a) : Collections.emptyList(); + } + + /** + * Returns the underlying raw agent output. + * + * @return the raw output + */ + public AgentOutput raw() { + return raw; + } +} diff --git a/ai/src/main/java/com/google/genkit/ai/agent/AgentResult.java b/ai/src/main/java/com/google/genkit/ai/agent/AgentResult.java new file mode 100644 index 000000000..e61b69717 --- /dev/null +++ b/ai/src/main/java/com/google/genkit/ai/agent/AgentResult.java @@ -0,0 +1,137 @@ +/* + * 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. + * + * SPDX-License-Identifier: Apache-2.0 + */ + +package com.google.genkit.ai.agent; + +import com.fasterxml.jackson.annotation.JsonInclude; +import com.fasterxml.jackson.annotation.JsonProperty; +import com.google.genkit.ai.Message; +import java.util.ArrayList; +import java.util.List; + +/** AgentResult represents the final result of an agent execution. */ +@JsonInclude(JsonInclude.Include.NON_NULL) +public class AgentResult { + + @JsonProperty("message") + private Message message; + + @JsonProperty("artifacts") + private List artifacts; + + @JsonProperty("finishReason") + private AgentFinishReason finishReason; + + /** Default constructor. */ + public AgentResult() {} + + private AgentResult(Builder builder) { + this.message = builder.message; + this.artifacts = builder.artifacts; + this.finishReason = builder.finishReason; + } + + /** + * Creates a builder for AgentResult. + * + * @return a new builder + */ + public static Builder builder() { + return new Builder(); + } + + /** + * Returns the result message. + * + * @return the message + */ + public Message getMessage() { + return message; + } + + /** + * Sets the result message. + * + * @param message the message + */ + public void setMessage(Message message) { + this.message = message; + } + + /** + * Returns the result artifacts. + * + * @return the artifacts + */ + public List getArtifacts() { + return artifacts; + } + + /** + * Sets the result artifacts. + * + * @param artifacts the artifacts + */ + public void setArtifacts(List artifacts) { + this.artifacts = artifacts != null ? new ArrayList<>(artifacts) : null; + } + + /** + * Returns the finish reason. + * + * @return the finish reason + */ + public AgentFinishReason getFinishReason() { + return finishReason; + } + + /** + * Sets the finish reason. + * + * @param finishReason the finish reason + */ + public void setFinishReason(AgentFinishReason finishReason) { + this.finishReason = finishReason; + } + + /** Builder for AgentResult. */ + public static class Builder { + private Message message; + private List artifacts; + private AgentFinishReason finishReason; + + public Builder message(Message message) { + this.message = message; + return this; + } + + public Builder artifacts(List artifacts) { + this.artifacts = artifacts; + return this; + } + + public Builder finishReason(AgentFinishReason finishReason) { + this.finishReason = finishReason; + return this; + } + + public AgentResult build() { + return new AgentResult(this); + } + } +} diff --git a/ai/src/main/java/com/google/genkit/ai/agent/AgentSessionContext.java b/ai/src/main/java/com/google/genkit/ai/agent/AgentSessionContext.java new file mode 100644 index 000000000..65bd24b00 --- /dev/null +++ b/ai/src/main/java/com/google/genkit/ai/agent/AgentSessionContext.java @@ -0,0 +1,98 @@ +/* + * 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. + * + * SPDX-License-Identifier: Apache-2.0 + */ + +package com.google.genkit.ai.agent; + +import java.util.function.Supplier; + +/** + * AgentSessionContext binds a {@link Session} to the current thread so that prompts, middleware, + * and tools can access the active session without passing it through every call frame. + * + *

A {@link ThreadLocal} is used for binding; the context is always cleared after {@link #run} or + * {@link #call} returns (even if the body throws). + */ +public final class AgentSessionContext { + + private static final ThreadLocal> CURRENT = new ThreadLocal<>(); + + private AgentSessionContext() {} + + /** + * Executes {@code body} with {@code session} bound to the current thread context. Clears the + * binding when the body returns or throws. + * + * @param session the session to bind (must not be null) + * @param body the runnable to execute + */ + public static void run(Session session, Runnable body) { + Session prior = CURRENT.get(); + CURRENT.set(session); + try { + body.run(); + } finally { + if (prior == null) { + CURRENT.remove(); + } else { + CURRENT.set(prior); + } + } + } + + /** + * Executes {@code body} with {@code session} bound to the current thread context and returns the + * result. Clears the binding when the body returns or throws. + * + * @param the return type + * @param session the session to bind (must not be null) + * @param body the supplier to execute + * @return the value returned by {@code body} + */ + public static T call(Session session, Supplier body) { + Session prior = CURRENT.get(); + CURRENT.set(session); + try { + return body.get(); + } finally { + if (prior == null) { + CURRENT.remove(); + } else { + CURRENT.set(prior); + } + } + } + + /** + * Returns the {@link Session} currently bound to this thread, or {@code null} if none. + * + * @return the current session, or null + */ + public static Session current() { + return CURRENT.get(); + } + + /** + * Returns the current session viewed as an {@link ArtifactStore}, or {@code null} if no session + * is bound to this thread. + * + * @return the current session as an {@link ArtifactStore}, or null + */ + public static ArtifactStore currentArtifactStore() { + return CURRENT.get(); + } +} diff --git a/ai/src/main/java/com/google/genkit/ai/agent/AgentStreamChunk.java b/ai/src/main/java/com/google/genkit/ai/agent/AgentStreamChunk.java new file mode 100644 index 000000000..5486152eb --- /dev/null +++ b/ai/src/main/java/com/google/genkit/ai/agent/AgentStreamChunk.java @@ -0,0 +1,169 @@ +/* + * 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. + * + * SPDX-License-Identifier: Apache-2.0 + */ + +package com.google.genkit.ai.agent; + +import com.fasterxml.jackson.annotation.JsonInclude; +import com.fasterxml.jackson.annotation.JsonProperty; +import com.fasterxml.jackson.databind.JsonNode; +import com.google.genkit.ai.ModelResponseChunk; + +/** + * AgentStreamChunk represents a streaming chunk from an agent execution. + * + *

The {@code customPatch} field is a JSON array of JSON-patch operations (RFC 6902). + */ +@JsonInclude(JsonInclude.Include.NON_NULL) +public class AgentStreamChunk { + + @JsonProperty("modelChunk") + private ModelResponseChunk modelChunk; + + /** JSON-patch array of ops; typed as JsonNode to stay aligned with core.jsonpatch.JsonPatch. */ + @JsonProperty("customPatch") + private JsonNode customPatch; + + @JsonProperty("artifact") + private Artifact artifact; + + @JsonProperty("turnEnd") + private TurnEnd turnEnd; + + /** Default constructor. */ + public AgentStreamChunk() {} + + private AgentStreamChunk(Builder builder) { + this.modelChunk = builder.modelChunk; + this.customPatch = builder.customPatch; + this.artifact = builder.artifact; + this.turnEnd = builder.turnEnd; + } + + /** + * Creates a builder for AgentStreamChunk. + * + * @return a new builder + */ + public static Builder builder() { + return new Builder(); + } + + /** + * Returns the model response chunk. + * + * @return the model chunk + */ + public ModelResponseChunk getModelChunk() { + return modelChunk; + } + + /** + * Sets the model response chunk. + * + * @param modelChunk the model chunk + */ + public void setModelChunk(ModelResponseChunk modelChunk) { + this.modelChunk = modelChunk; + } + + /** + * Returns the custom-state patch as a JSON array of patch operations. + * + * @return the custom patch node (array) + */ + public JsonNode getCustomPatch() { + return customPatch; + } + + /** + * Sets the custom-state patch. + * + * @param customPatch a JSON array of patch operations + */ + public void setCustomPatch(JsonNode customPatch) { + this.customPatch = customPatch; + } + + /** + * Returns the artifact in this chunk. + * + * @return the artifact, or null if not present + */ + public Artifact getArtifact() { + return artifact; + } + + /** + * Sets the artifact. + * + * @param artifact the artifact + */ + public void setArtifact(Artifact artifact) { + this.artifact = artifact; + } + + /** + * Returns the turn-end signal. + * + * @return the turn end, or null if the turn has not ended + */ + public TurnEnd getTurnEnd() { + return turnEnd; + } + + /** + * Sets the turn-end signal. + * + * @param turnEnd the turn end + */ + public void setTurnEnd(TurnEnd turnEnd) { + this.turnEnd = turnEnd; + } + + /** Builder for AgentStreamChunk. */ + public static class Builder { + private ModelResponseChunk modelChunk; + private JsonNode customPatch; + private Artifact artifact; + private TurnEnd turnEnd; + + public Builder modelChunk(ModelResponseChunk modelChunk) { + this.modelChunk = modelChunk; + return this; + } + + public Builder customPatch(JsonNode customPatch) { + this.customPatch = customPatch; + return this; + } + + public Builder artifact(Artifact artifact) { + this.artifact = artifact; + return this; + } + + public Builder turnEnd(TurnEnd turnEnd) { + this.turnEnd = turnEnd; + return this; + } + + public AgentStreamChunk build() { + return new AgentStreamChunk(this); + } + } +} diff --git a/ai/src/main/java/com/google/genkit/ai/agent/AgentTransport.java b/ai/src/main/java/com/google/genkit/ai/agent/AgentTransport.java new file mode 100644 index 000000000..2674e669c --- /dev/null +++ b/ai/src/main/java/com/google/genkit/ai/agent/AgentTransport.java @@ -0,0 +1,72 @@ +/* + * 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. + * + * SPDX-License-Identifier: Apache-2.0 + */ + +package com.google.genkit.ai.agent; + +import java.util.function.Consumer; + +/** + * AgentTransport abstracts how a single agent turn is executed. The in-process transport ({@code + * internal.InProcessTransport}) drives a locally-defined {@link Agent}'s bidi action; a remote + * transport (future) would issue HTTP/RPC calls. {@link AgentChat} is written against this + * interface so it can drive either. + * + *

The transport models the agent's one-turn-per-request contract: {@link #runTurn} + * feeds exactly one {@link AgentInput} (with its {@link AgentInit}) and returns the turn's final + * {@link AgentOutput}, streaming {@link AgentStreamChunk}s along the way. + * + * @param the type of custom session state + */ +public interface AgentTransport { + + /** + * Runs ONE turn: feeds {@code input} with {@code init}, collects stream chunks via {@code + * onChunk}, and returns the final output. + * + * @param input the turn input (user message, resume, or detach) + * @param init the initialization carrying resume context (snapshotId / sessionId / inline state); + * may be {@code null} for the very first turn of a fresh session + * @param onChunk consumer invoked for each streamed chunk; may be a no-op + * @return the final {@link AgentOutput} for the turn + */ + AgentOutput runTurn(AgentInput input, AgentInit init, Consumer onChunk); + + /** + * Retrieves a session snapshot. + * + * @param req the request specifying which snapshot to retrieve + * @return the matching snapshot, or {@code null} if not found / not supported + */ + SessionSnapshot getSnapshot(GetSnapshotRequest req); + + /** + * Attempts to abort a pending snapshot. + * + * @param snapshotId the snapshot to abort + * @return the resulting status, or {@code null} if not found / not supported + */ + SnapshotStatus abort(String snapshotId); + + /** + * Returns whether the underlying agent is server-managed (state persisted server-side) versus + * client-managed (state round-tripped by the caller). + * + * @return {@code true} if server-managed + */ + boolean serverManaged(); +} diff --git a/ai/src/main/java/com/google/genkit/ai/agent/Artifact.java b/ai/src/main/java/com/google/genkit/ai/agent/Artifact.java new file mode 100644 index 000000000..30bcab7f4 --- /dev/null +++ b/ai/src/main/java/com/google/genkit/ai/agent/Artifact.java @@ -0,0 +1,139 @@ +/* + * 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. + * + * SPDX-License-Identifier: Apache-2.0 + */ + +package com.google.genkit.ai.agent; + +import com.fasterxml.jackson.annotation.JsonInclude; +import com.fasterxml.jackson.annotation.JsonProperty; +import com.google.genkit.ai.Part; +import java.util.ArrayList; +import java.util.List; +import java.util.Map; + +/** Artifact represents a named collection of content parts produced during agent execution. */ +@JsonInclude(JsonInclude.Include.NON_NULL) +public class Artifact { + + @JsonProperty("name") + private String name; + + /** parts is always present on the wire (required field). */ + @JsonProperty("parts") + private List parts = new ArrayList<>(); + + @JsonProperty("metadata") + private Map metadata; + + /** Default constructor. */ + public Artifact() {} + + private Artifact(Builder builder) { + this.name = builder.name; + this.parts = builder.parts != null ? new ArrayList<>(builder.parts) : new ArrayList<>(); + this.metadata = builder.metadata; + } + + /** + * Creates a builder for Artifact. + * + * @return a new builder + */ + public static Builder builder() { + return new Builder(); + } + + /** + * Returns the artifact name. + * + * @return the name + */ + public String getName() { + return name; + } + + /** + * Sets the artifact name. + * + * @param name the name + */ + public void setName(String name) { + this.name = name; + } + + /** + * Returns the artifact parts. + * + * @return the parts (never null) + */ + public List getParts() { + return parts; + } + + /** + * Sets the artifact parts. + * + * @param parts the parts + */ + public void setParts(List parts) { + this.parts = parts != null ? new ArrayList<>(parts) : new ArrayList<>(); + } + + /** + * Returns the artifact metadata. + * + * @return the metadata, or null if not set + */ + public Map getMetadata() { + return metadata; + } + + /** + * Sets the artifact metadata. + * + * @param metadata the metadata + */ + public void setMetadata(Map metadata) { + this.metadata = metadata; + } + + /** Builder for Artifact. */ + public static class Builder { + private String name; + private List parts; + private Map metadata; + + public Builder name(String name) { + this.name = name; + return this; + } + + public Builder parts(List parts) { + this.parts = parts; + return this; + } + + public Builder metadata(Map metadata) { + this.metadata = metadata; + return this; + } + + public Artifact build() { + return new Artifact(this); + } + } +} diff --git a/ai/src/main/java/com/google/genkit/ai/agent/ArtifactStore.java b/ai/src/main/java/com/google/genkit/ai/agent/ArtifactStore.java new file mode 100644 index 000000000..08a2b3cd5 --- /dev/null +++ b/ai/src/main/java/com/google/genkit/ai/agent/ArtifactStore.java @@ -0,0 +1,45 @@ +/* + * 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. + * + * SPDX-License-Identifier: Apache-2.0 + */ + +package com.google.genkit.ai.agent; + +import java.util.List; + +/** + * ArtifactStore is a state-agnostic view of artifact storage used by middleware and tools that do + * not need to know the custom state type {@code S} of the session. + * + *

{@link Session} implements this interface. + */ +public interface ArtifactStore { + + /** + * Returns a copy of the current list of artifacts. + * + * @return a copy of the artifacts (never null) + */ + List getArtifacts(); + + /** + * Adds artifacts, deduplicating by name. If an artifact with the same non-null name already + * exists, it is replaced in place. Artifacts with null names are always appended. + * + * @param artifacts the artifacts to add + */ + void addArtifacts(Artifact... artifacts); +} diff --git a/ai/src/main/java/com/google/genkit/ai/agent/ClientTransform.java b/ai/src/main/java/com/google/genkit/ai/agent/ClientTransform.java new file mode 100644 index 000000000..67f364470 --- /dev/null +++ b/ai/src/main/java/com/google/genkit/ai/agent/ClientTransform.java @@ -0,0 +1,37 @@ +/* + * 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. + * + * SPDX-License-Identifier: Apache-2.0 + */ + +package com.google.genkit.ai.agent; + +/** + * ClientTransform allows a client-managed agent to transform session state before returning it to + * the caller. + * + * @param the type of custom state + */ +@FunctionalInterface +public interface ClientTransform { + + /** + * Transforms the given session state before it is returned to the caller. + * + * @param state the current session state + * @return the transformed session state + */ + SessionState transformState(SessionState state); +} diff --git a/ai/src/main/java/com/google/genkit/ai/agent/CustomAgentConfig.java b/ai/src/main/java/com/google/genkit/ai/agent/CustomAgentConfig.java new file mode 100644 index 000000000..da9928da9 --- /dev/null +++ b/ai/src/main/java/com/google/genkit/ai/agent/CustomAgentConfig.java @@ -0,0 +1,210 @@ +/* + * 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. + * + * SPDX-License-Identifier: Apache-2.0 + */ + +package com.google.genkit.ai.agent; + +/** + * Configuration for defining a custom agent via {@code AgentActions.defineCustomAgent}. + * + *

Use the {@link Builder} to configure the agent. The only required field is {@link #getName()}. + * + * @param the type of custom session state + */ +public final class CustomAgentConfig { + + private final String name; + private final String description; + private final Class stateType; + private final SessionStore store; + private final ClientTransform clientTransform; + private final SessionStoreOptions storeOptions; + + private CustomAgentConfig(Builder builder) { + this.name = builder.name; + this.description = builder.description; + this.stateType = builder.stateType; + this.store = builder.store; + this.clientTransform = builder.clientTransform; + this.storeOptions = builder.storeOptions; + } + + /** + * Creates a builder for CustomAgentConfig. + * + * @param the type of custom session state + * @return a new builder + */ + public static Builder builder() { + return new Builder<>(); + } + + /** + * Returns the agent's registered name. + * + * @return the agent name (never null) + */ + public String getName() { + return name; + } + + /** + * Returns the agent's human-readable description. + * + * @return the description, or {@code null} if not set + */ + public String getDescription() { + return description; + } + + /** + * Returns the Java class for the agent's custom state type. + * + * @return the state type class, or {@code null} if not specified + */ + public Class getStateType() { + return stateType; + } + + /** + * Returns the session store for server-managed agents. + * + *

When {@code null}, the agent operates in client-managed mode. + * + * @return the session store, or {@code null} for client-managed mode + */ + public SessionStore getStore() { + return store; + } + + /** + * Returns the client-transform applied to session state before returning it to the caller in + * client-managed mode. + * + * @return the client transform, or {@code null} if not set + */ + public ClientTransform getClientTransform() { + return clientTransform; + } + + /** + * Returns the options forwarded to store operations. + * + * @return the store options, or {@code null} if not set + */ + public SessionStoreOptions getStoreOptions() { + return storeOptions; + } + + /** + * Builder for {@link CustomAgentConfig}. + * + * @param the type of custom session state + */ + public static final class Builder { + + private String name; + private String description; + private Class stateType; + private SessionStore store; + private ClientTransform clientTransform; + private SessionStoreOptions storeOptions; + + private Builder() {} + + /** + * Sets the agent's registered name (required). + * + * @param name the agent name + * @return this builder + */ + public Builder name(String name) { + this.name = name; + return this; + } + + /** + * Sets the agent's human-readable description. + * + * @param description the description + * @return this builder + */ + public Builder description(String description) { + this.description = description; + return this; + } + + /** + * Sets the Java class for the agent's custom state type. + * + * @param stateType the state type class + * @return this builder + */ + public Builder stateType(Class stateType) { + this.stateType = stateType; + return this; + } + + /** + * Sets the session store for server-managed mode. Pass {@code null} (or omit) for + * client-managed mode. + * + * @param store the session store + * @return this builder + */ + public Builder store(SessionStore store) { + this.store = store; + return this; + } + + /** + * Sets the client-transform applied to session state before returning it to the caller in + * client-managed mode. + * + * @param clientTransform the transform + * @return this builder + */ + public Builder clientTransform(ClientTransform clientTransform) { + this.clientTransform = clientTransform; + return this; + } + + /** + * Sets the options forwarded to store operations. + * + * @param storeOptions the store options + * @return this builder + */ + public Builder storeOptions(SessionStoreOptions storeOptions) { + this.storeOptions = storeOptions; + return this; + } + + /** + * Builds the {@link CustomAgentConfig}. + * + * @return a new {@link CustomAgentConfig} + * @throws IllegalStateException if {@code name} is null or blank + */ + public CustomAgentConfig build() { + if (name == null || name.isBlank()) { + throw new IllegalStateException("name is required"); + } + return new CustomAgentConfig<>(this); + } + } +} diff --git a/ai/src/main/java/com/google/genkit/ai/agent/FileSessionStore.java b/ai/src/main/java/com/google/genkit/ai/agent/FileSessionStore.java new file mode 100644 index 000000000..41df95df4 --- /dev/null +++ b/ai/src/main/java/com/google/genkit/ai/agent/FileSessionStore.java @@ -0,0 +1,762 @@ +/* + * 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. + * + * SPDX-License-Identifier: Apache-2.0 + */ + +package com.google.genkit.ai.agent; + +import com.fasterxml.jackson.databind.JsonNode; +import com.fasterxml.jackson.databind.ObjectMapper; +import com.google.genkit.ai.agent.internal.LeafSelection; +import com.google.genkit.ai.agent.internal.PointerDoc; +import com.google.genkit.core.GenkitException; +import com.google.genkit.core.JsonUtils; +import java.io.IOException; +import java.nio.file.Files; +import java.nio.file.Path; +import java.nio.file.Paths; +import java.nio.file.StandardCopyOption; +import java.time.Instant; +import java.util.ArrayList; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.UUID; +import java.util.concurrent.Executors; +import java.util.concurrent.ScheduledExecutorService; +import java.util.concurrent.ScheduledFuture; +import java.util.concurrent.TimeUnit; +import java.util.function.Consumer; +import java.util.stream.Stream; + +/** + * Disk-backed implementation of {@link SessionStore} and {@link SnapshotSubscriber}. + * + *

On-disk layout

+ * + * 
+ * <dir>/<prefix>/<snapshotId>.json              one snapshot per file
+ * <dir>/<prefix>/.pointers/<sessionId>.json     pointer: { currentSnapshotId, currentCreatedAt, updatedAt }
+ *
+ * + *

The default prefix is {@code "global"}. + * + *

Atomic writes

+ * + *

Each write goes to a {@code .tmp} file first, then is renamed to the target path using {@link + * StandardCopyOption#ATOMIC_MOVE} (falling back to {@link StandardCopyOption#REPLACE_EXISTING} if + * the filesystem does not support atomic moves). This ensures no partial JSON files are visible. + * + *

Subscriber / polling

+ * + *

Subscriptions use polling (a shared daemon {@link ScheduledExecutorService}) rather than + * {@code WatchService}. This is simpler and more reliable across platforms (particularly macOS + * where {@code WatchService} uses polling internally anyway). The poll interval defaults to 2000 ms + * and can be overridden via the builder for fast-iteration tests. + * + *

De-duplication is done by comparing the serialised content of the snapshot file on each poll + * tick; the callback fires only when the content has changed. + * + * @param the type of custom session state + */ +public final class FileSessionStore implements SessionStore, SnapshotSubscriber { + + private static final ObjectMapper MAPPER = JsonUtils.getObjectMapper(); + + /** Default prefix sub-directory used when none is configured. */ + static final String DEFAULT_PREFIX = "global"; + + /** Default polling interval in milliseconds. */ + static final long DEFAULT_POLL_INTERVAL_MS = 2000L; + + private final Path baseDir; + private final String prefix; + private final int maxPersistedChainLength; // 0 = unlimited + private final boolean rejectBranchingSessions; + private final long snapshotWatchPollIntervalMs; + + /** + * Shared daemon scheduler used for all poll-based subscriptions. One instance per store; shut + * down when the store itself is closed (or left to die as a daemon thread on JVM exit). + */ + private final ScheduledExecutorService scheduler; + + /** Guard for saveSnapshot / getSnapshot operations. */ + private final Object lock = new Object(); + + // ── constructors / builder ──────────────────────────────────────────────── + + /** + * Creates a new {@code FileSessionStore} with default options (prefix = {@code "global"}, no + * chain pruning, polling every 2000 ms). + * + * @param dir the root directory for persisted snapshots; created if absent + */ + public FileSessionStore(String dir) { + this(dir, DEFAULT_PREFIX, 0, false, DEFAULT_POLL_INTERVAL_MS); + } + + private FileSessionStore( + String dir, + String prefix, + int maxPersistedChainLength, + boolean rejectBranchingSessions, + long snapshotWatchPollIntervalMs) { + this.baseDir = Paths.get(dir); + this.prefix = prefix != null ? prefix : DEFAULT_PREFIX; + this.maxPersistedChainLength = maxPersistedChainLength; + this.rejectBranchingSessions = rejectBranchingSessions; + this.snapshotWatchPollIntervalMs = snapshotWatchPollIntervalMs; + this.scheduler = + Executors.newSingleThreadScheduledExecutor( + r -> { + Thread t = new Thread(r, "genkit-file-session-store-poll"); + t.setDaemon(true); + return t; + }); + + // Ensure base directory exists + try { + Files.createDirectories(this.baseDir); + } catch (IOException e) { + throw new GenkitException("Failed to create session store directory: " + dir, e); + } + } + + /** + * Creates a builder for {@code FileSessionStore}. + * + * @param the type of custom session state + * @param dir the root directory for persisted snapshots + * @return a new builder + */ + public static Builder builder(String dir) { + return new Builder<>(dir); + } + + // ── SnapshotReader ──────────────────────────────────────────────────────── + + /** + * {@inheritDoc} + * + *

When {@link GetSnapshotOptions#getSnapshotId()} is set, reads {@code /.json} and + * deserialises it (or returns {@code null} if absent). When {@link + * GetSnapshotOptions#getSessionId()} is set, tries the pointer first; falls back to scanning all + * {@code *.json} files in the prefix dir (skipping {@code .pointers/}), runs {@link + * LeafSelection}, rewrites the pointer, and returns the leaf. + */ + @Override + @SuppressWarnings({"unchecked", "rawtypes"}) + public SessionSnapshot getSnapshot(GetSnapshotOptions opts) { + if (opts == null) { + return null; + } + + if (opts.getSnapshotId() != null) { + checkPathSafety(opts.getSnapshotId()); + } + + synchronized (lock) { + Path prefixDir = prefixDir(); + + if (opts.getSnapshotId() != null) { + return readSnapshotFile(prefixDir.resolve(opts.getSnapshotId() + ".json")); + } + + if (opts.getSessionId() != null) { + String sessionId = opts.getSessionId(); + + // 1. Try pointer + Path pointerFile = pointerFile(prefixDir, sessionId); + if (Files.exists(pointerFile)) { + PointerDoc pointer = readPointerFile(pointerFile); + if (pointer != null && pointer.getCurrentSnapshotId() != null) { + Path snapshotFile = prefixDir.resolve(pointer.getCurrentSnapshotId() + ".json"); + if (Files.exists(snapshotFile)) { + SessionSnapshot snap = readSnapshotFile(snapshotFile); + if (snap != null && sessionId.equals(effectiveSessionId(snap))) { + return snap; + } + } + } + } + + // 2. Scan fallback + return scanAndSelectLeaf(prefixDir, sessionId, pointerFile); + } + + return null; + } + } + + // ── SnapshotWriter ──────────────────────────────────────────────────────── + + /** + * {@inheritDoc} + * + *

Implements the full saveSnapshot contract (mirrors {@code InMemorySessionStore}), but + * persists to disk: + * + *

    + *
  1. Validate path safety for {@code snapshotId} (if non-null). + *
  2. Under lock: read existing snapshot file. + *
  3. Apply mutator; if result is {@code null}, return {@code null} (no write). + *
  4. Determine final id; apply sessionId / status defaulting. + *
  5. Atomic write to {@code /.json}. + *
  6. Advance pointer if this is a new snapshot row. + *
  7. Prune chain if {@code maxPersistedChainLength > 0}. + *
  8. Notify subscribers if status changed. + *
+ */ + @Override + @SuppressWarnings({"unchecked", "rawtypes"}) + public String saveSnapshot( + String snapshotId, SnapshotMutator mutator, SessionStoreOptions options) { + + if (snapshotId != null) { + checkPathSafety(snapshotId); + } + + String finalId; + SessionSnapshot result; + boolean isNewRow; + SnapshotStatus existingStatus; + + synchronized (lock) { + Path prefixDir = prefixDir(); + ensureDirectoryExists(prefixDir); + + // Step 1: read existing + Path existingFile = snapshotId != null ? prefixDir.resolve(snapshotId + ".json") : null; + SessionSnapshot existing = + (existingFile != null && Files.exists(existingFile)) + ? (SessionSnapshot) readSnapshotFile(existingFile) + : null; + + // Step 2: apply mutator + result = mutator.apply(existing); + if (result == null) { + return null; + } + + // Step 3: determine final id + if (snapshotId != null) { + finalId = snapshotId; + } else if (result.getSnapshotId() != null && !result.getSnapshotId().isBlank()) { + finalId = result.getSnapshotId(); + } else { + finalId = UUID.randomUUID().toString(); + } + result.setSnapshotId(finalId); + + // Step 4: preserve sessionId from existing row + if (existing != null && existing.getSessionId() != null) { + result.setSessionId(existing.getSessionId()); + } + + // Step 4b: fall back to state.sessionId if top-level is null + if (result.getSessionId() == null && result.getState() != null) { + result.setSessionId(result.getState().getSessionId()); + } + + // Step 4c: validate sessionId + if (result.getSessionId() == null || result.getSessionId().isBlank()) { + throw GenkitException.builder() + .message("snapshot requires sessionId") + .errorCode("INVALID_ARGUMENT") + .build(); + } + + // Step 5: default null status to COMPLETED + if (result.getStatus() == null) { + result.setStatus(SnapshotStatus.COMPLETED); + } + + existingStatus = existing != null ? existing.getStatus() : null; + isNewRow = (existing == null); + + // Step 6: atomic write + Path targetFile = prefixDir.resolve(finalId + ".json"); + atomicWrite(targetFile, result); + + // Step 7: advance pointer (only for new rows with a sessionId) + if (isNewRow && result.getSessionId() != null && !result.getSessionId().isBlank()) { + advancePointer(prefixDir, result); + } + + // Step 8: chain pruning + if (maxPersistedChainLength > 0 && isNewRow) { + pruneChain(prefixDir, result); + } + } // lock released + + // Step 9: notify subscribers outside the lock + boolean statusChanged = existingStatus != result.getStatus(); + if (statusChanged) { + notifySubscribers(finalId, result); + } + + return finalId; + } + + // ── SnapshotSubscriber ──────────────────────────────────────────────────── + + /** + * {@inheritDoc} + * + *

Registers a polling subscription for the snapshot identified by {@code snapshotId}. If the + * snapshot file already exists, the callback is fired immediately with the current content. A + * shared daemon {@link ScheduledExecutorService} polls the file every {@code + * snapshotWatchPollIntervalMs} milliseconds; the callback fires when the serialised content + * changes. Calling {@link AutoCloseable#close()} on the returned handle cancels the poll. + */ + @Override + public AutoCloseable onSnapshotStateChange( + String snapshotId, Consumer> cb, SessionStoreOptions options) { + + checkPathSafety(snapshotId); + + Path prefixDir = prefixDir(); + Path snapshotFile = prefixDir.resolve(snapshotId + ".json"); + + // State for de-duplication: last known serialised content (null = not yet seen) + final String[] lastContent = {null}; + + // Fire immediately if file already exists + SessionSnapshot initial = readSnapshotFileUnsafe(snapshotFile); + if (initial != null) { + lastContent[0] = serializeQuietly(initial); + cb.accept(initial); + } + + // Schedule polling + ScheduledFuture future = + scheduler.scheduleAtFixedRate( + () -> { + try { + SessionSnapshot snap = readSnapshotFileUnsafe(snapshotFile); + if (snap == null) { + return; + } + String content = serializeQuietly(snap); + if (!content.equals(lastContent[0])) { + lastContent[0] = content; + cb.accept(snap); + } + } catch (Exception e) { + // Swallow poll errors — don't kill the scheduler thread + } + }, + snapshotWatchPollIntervalMs, + snapshotWatchPollIntervalMs, + TimeUnit.MILLISECONDS); + + return () -> future.cancel(false); + } + + // ── helpers ─────────────────────────────────────────────────────────────── + + /** Returns the prefix directory ({@code /}). */ + private Path prefixDir() { + return baseDir.resolve(prefix); + } + + /** Returns the pointer file path for the given session ID. */ + private Path pointerFile(Path prefixDir, String sessionId) { + return prefixDir.resolve(".pointers").resolve(sessionId + ".json"); + } + + /** Ensures the directory (and its parents) exist. */ + private static void ensureDirectoryExists(Path dir) { + try { + Files.createDirectories(dir); + } catch (IOException e) { + throw new GenkitException("Failed to create directory: " + dir, e); + } + } + + /** + * Validates that the given name component does not allow path traversal. + * + *

Rejects names that contain {@code /}, {@code \}, or NUL characters; equal to {@code .} or + * {@code ..}; or start with {@code .}. + * + * @throws IllegalArgumentException if the name is unsafe + */ + private static void checkPathSafety(String name) { + if (name == null) { + return; + } + if (name.contains("/") + || name.contains("\\") + || name.contains("\0") + || name.equals(".") + || name.equals("..") + || name.startsWith(".")) { + throw new IllegalArgumentException( + "Unsafe snapshotId/sessionId rejected (path traversal risk): " + name); + } + } + + /** Reads and deserialises a snapshot JSON file; returns {@code null} if the file is absent. */ + @SuppressWarnings({"unchecked", "rawtypes"}) + private SessionSnapshot readSnapshotFile(Path file) { + if (!Files.exists(file)) { + return null; + } + try { + byte[] bytes = Files.readAllBytes(file); + return (SessionSnapshot) MAPPER.readValue(bytes, SessionSnapshot.class); + } catch (IOException e) { + // Corrupt file — treat as absent + return null; + } + } + + /** Like {@link #readSnapshotFile(Path)} but without the type param (for subscriber use). */ + @SuppressWarnings("rawtypes") + private SessionSnapshot readSnapshotFileUnsafe(Path file) { + if (!Files.exists(file)) { + return null; + } + try { + byte[] bytes = Files.readAllBytes(file); + return MAPPER.readValue(bytes, SessionSnapshot.class); + } catch (IOException e) { + return null; + } + } + + /** Reads and deserialises a pointer JSON file; returns {@code null} if absent or corrupt. */ + private PointerDoc readPointerFile(Path file) { + if (!Files.exists(file)) { + return null; + } + try { + byte[] bytes = Files.readAllBytes(file); + return MAPPER.readValue(bytes, PointerDoc.class); + } catch (IOException e) { + return null; + } + } + + /** + * Atomically writes {@code value} as JSON to {@code target}. + * + *

Writes to a temp file ({@code ..tmp}) first, then renames to {@code target}. + */ + private static void atomicWrite(Path target, T value) { + ensureDirectoryExists(target.getParent()); + Path tmp = target.getParent().resolve(target.getFileName() + "." + UUID.randomUUID() + ".tmp"); + try { + byte[] bytes = MAPPER.writeValueAsBytes(value); + Files.write(tmp, bytes); + try { + Files.move(tmp, target, StandardCopyOption.ATOMIC_MOVE); + } catch (IOException atomicEx) { + // Fallback: REPLACE_EXISTING (non-atomic but safe enough) + Files.move(tmp, target, StandardCopyOption.REPLACE_EXISTING); + } + } catch (IOException e) { + // Clean up the temp file if we fail + try { + Files.deleteIfExists(tmp); + } catch (IOException ignored) { + // best effort + } + throw new GenkitException("Failed to write snapshot file: " + target, e); + } + } + + /** + * Advances the pointer for the session if this snapshot sorts ahead of the current pointer. + * + *

The pointer is only advanced when {@code snap} is a new row (existing == null). Comparison + * is by ({@code createdAt}, {@code snapshotId}). + */ + private void advancePointer(Path prefixDir, SessionSnapshot snap) { + if (snap.getSessionId() == null || snap.getSessionId().isBlank()) { + return; + } + Path pointersDir = prefixDir.resolve(".pointers"); + ensureDirectoryExists(pointersDir); + + Path pointerFile = pointersDir.resolve(snap.getSessionId() + ".json"); + PointerDoc existing = readPointerFile(pointerFile); + + boolean shouldAdvance = false; + if (existing == null || existing.getCurrentSnapshotId() == null) { + shouldAdvance = true; + } else { + // Compare by (createdAt, snapshotId) + int cmp = + compareCreatedAtThenId( + snap.getCreatedAt(), + snap.getSnapshotId(), + existing.getCurrentCreatedAt(), + existing.getCurrentSnapshotId()); + shouldAdvance = cmp > 0; + } + + if (shouldAdvance) { + String now = Instant.now().toString(); + PointerDoc pointer = new PointerDoc(snap.getSnapshotId(), snap.getCreatedAt(), now); + atomicWrite(pointerFile, pointer); + } + } + + /** + * Compares two (createdAt, snapshotId) pairs. Returns positive if (aCreatedAt, aId) sorts after + * (bCreatedAt, bId). + */ + private static int compareCreatedAtThenId( + String aCreatedAt, String aId, String bCreatedAt, String bId) { + Instant ia = parseInstant(aCreatedAt); + Instant ib = parseInstant(bCreatedAt); + int cmp = ia.compareTo(ib); + if (cmp != 0) { + return cmp; + } + String sa = aId != null ? aId : ""; + String sb = bId != null ? bId : ""; + return sa.compareTo(sb); + } + + /** Parses an RFC-3339 timestamp; returns {@link Instant#EPOCH} for null/unparseable. */ + private static Instant parseInstant(String rfc3339) { + if (rfc3339 == null || rfc3339.isEmpty()) { + return Instant.EPOCH; + } + try { + return Instant.parse(rfc3339); + } catch (Exception e) { + return Instant.EPOCH; + } + } + + /** + * Scans the prefix directory for all snapshot files belonging to {@code sessionId}, selects the + * leaf via {@link LeafSelection}, rewrites the pointer, and returns the leaf. + */ + @SuppressWarnings({"unchecked", "rawtypes"}) + private SessionSnapshot scanAndSelectLeaf(Path prefixDir, String sessionId, Path pointerFile) { + if (!Files.isDirectory(prefixDir)) { + return null; + } + + List> matching = new ArrayList<>(); + try (Stream files = Files.list(prefixDir)) { + files + .filter(p -> p.getFileName().toString().endsWith(".json")) + .filter(p -> !Files.isDirectory(p)) + .forEach( + p -> { + SessionSnapshot snap = readSnapshotFile(p); + if (snap != null && sessionId.equals(effectiveSessionId(snap))) { + matching.add(snap); + } + }); + } catch (IOException e) { + return null; + } + + if (matching.isEmpty()) { + return null; + } + + SessionSnapshot leaf = LeafSelection.selectLeaf(matching, rejectBranchingSessions); + if (leaf == null) { + return null; + } + + // Rewrite pointer + Path pointersDir = prefixDir.resolve(".pointers"); + ensureDirectoryExists(pointersDir); + String now = Instant.now().toString(); + PointerDoc pointer = new PointerDoc(leaf.getSnapshotId(), leaf.getCreatedAt(), now); + atomicWrite(pointerFile, pointer); + + return leaf; + } + + /** + * Prunes the parent chain rooted at {@code snap}, deleting snapshot files beyond the newest + * {@code maxPersistedChainLength}. + * + *

Walks the parent chain backwards from {@code snap}; collects all ancestors in order (newest + * first); deletes those beyond position {@code maxPersistedChainLength - 1}. + */ + private void pruneChain(Path prefixDir, SessionSnapshot snap) { + // Build the chain newest-first + List chain = new ArrayList<>(); + String current = snap.getSnapshotId(); + // Collect the chain by following parentId links (read from files) + // We need a map of snapshotId -> parentId to walk the chain. + // Build by collecting all snapshots for this session. + Map parentOf = new HashMap<>(); // snapshotId -> parentId + if (Files.isDirectory(prefixDir)) { + try (Stream files = Files.list(prefixDir)) { + files + .filter(p -> p.getFileName().toString().endsWith(".json")) + .filter(p -> !Files.isDirectory(p)) + .filter(p -> !p.getParent().getFileName().toString().equals(".pointers")) + .forEach( + p -> { + SessionSnapshot s = readSnapshotFileUnsafe(p); + if (s != null && s.getSnapshotId() != null) { + parentOf.put(s.getSnapshotId(), s.getParentId()); // parentId may be null + } + }); + } catch (IOException e) { + return; // best effort + } + } + + // Walk from snap backwards through parentId links + String node = snap.getSnapshotId(); + int maxWalk = parentOf.size() + 1; // guard against cycles + int walked = 0; + while (node != null && walked++ < maxWalk) { + chain.add(node); + node = parentOf.getOrDefault(node, null); + } + + // chain[0] is the newest; delete anything beyond position maxPersistedChainLength-1 + for (int i = maxPersistedChainLength; i < chain.size(); i++) { + Path toDelete = prefixDir.resolve(chain.get(i) + ".json"); + try { + Files.deleteIfExists(toDelete); + } catch (IOException e) { + // best effort + } + } + } + + /** Returns the effective sessionId for a snapshot (top-level, then state.sessionId). */ + private static String effectiveSessionId(SessionSnapshot snap) { + if (snap == null) { + return null; + } + if (snap.getSessionId() != null) { + return snap.getSessionId(); + } + if (snap.getState() != null) { + return snap.getState().getSessionId(); + } + return null; + } + + /** + * Serialises a snapshot to a JSON string for change-detection de-duplication; returns "" on + * error. + */ + private static String serializeQuietly(SessionSnapshot snap) { + try { + JsonNode node = MAPPER.valueToTree(snap); + return MAPPER.writeValueAsString(node); + } catch (Exception e) { + return ""; + } + } + + /** Notifies in-process subscribers (used when saveSnapshot changes status). */ + private void notifySubscribers(String finalId, SessionSnapshot result) { + // For disk-based store the polling mechanism handles notifications. + // This method is a no-op: polling detects the file change independently. + // (In-process callers that registered via onSnapshotStateChange will be notified by the poll.) + } + + // ── Builder ─────────────────────────────────────────────────────────────── + + /** + * Builder for {@link FileSessionStore}. + * + * @param the type of custom session state + */ + public static final class Builder { + private final String dir; + private String prefix = DEFAULT_PREFIX; + private int maxPersistedChainLength = 0; + private boolean rejectBranchingSessions = false; + private long snapshotWatchPollIntervalMs = DEFAULT_POLL_INTERVAL_MS; + + private Builder(String dir) { + this.dir = dir; + } + + /** + * Sets the prefix sub-directory (default {@code "global"}). + * + * @param prefix the prefix + * @return this builder + */ + public Builder prefix(String prefix) { + this.prefix = prefix; + return this; + } + + /** + * Sets the maximum number of snapshot files to retain in the parent chain. When set to a + * positive value, older ancestors are deleted after each new save. Zero (default) means + * unlimited. + * + * @param maxPersistedChainLength the maximum chain length, or 0 for unlimited + * @return this builder + */ + public Builder maxPersistedChainLength(int maxPersistedChainLength) { + this.maxPersistedChainLength = maxPersistedChainLength; + return this; + } + + /** + * If {@code true}, {@link LeafSelection#selectLeaf} throws when more than one leaf exists for a + * session. + * + * @param rejectBranchingSessions whether to reject branching + * @return this builder + */ + public Builder rejectBranchingSessions(boolean rejectBranchingSessions) { + this.rejectBranchingSessions = rejectBranchingSessions; + return this; + } + + /** + * Sets the polling interval for snapshot-change subscriptions (default 2000 ms). Use a small + * value (e.g. 100 ms) in tests for faster callback detection. + * + * @param ms the poll interval in milliseconds + * @return this builder + */ + public Builder snapshotWatchPollIntervalMs(long ms) { + this.snapshotWatchPollIntervalMs = ms; + return this; + } + + /** + * Builds a new {@link FileSessionStore}. + * + * @return a new store + */ + public FileSessionStore build() { + return new FileSessionStore<>( + dir, + prefix, + maxPersistedChainLength, + rejectBranchingSessions, + snapshotWatchPollIntervalMs); + } + } +} diff --git a/ai/src/main/java/com/google/genkit/ai/agent/GetSnapshotOptions.java b/ai/src/main/java/com/google/genkit/ai/agent/GetSnapshotOptions.java new file mode 100644 index 000000000..b8482a2a9 --- /dev/null +++ b/ai/src/main/java/com/google/genkit/ai/agent/GetSnapshotOptions.java @@ -0,0 +1,104 @@ +/* + * 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. + * + * SPDX-License-Identifier: Apache-2.0 + */ + +package com.google.genkit.ai.agent; + +/** + * Options for retrieving a session snapshot. + * + *

Exactly one of {@code snapshotId} or {@code sessionId} must be set by callers. When {@code + * snapshotId} is set the store returns that specific snapshot. When {@code sessionId} is set the + * store resolves and returns the latest leaf snapshot for that session. + */ +public final class GetSnapshotOptions { + + private final String snapshotId; + private final String sessionId; + + private GetSnapshotOptions(Builder builder) { + this.snapshotId = builder.snapshotId; + this.sessionId = builder.sessionId; + } + + /** + * Returns the specific snapshot ID to retrieve, or {@code null} if not set. + * + * @return the snapshot ID + */ + public String getSnapshotId() { + return snapshotId; + } + + /** + * Returns the session ID whose latest leaf snapshot should be retrieved, or {@code null} if not + * set. + * + * @return the session ID + */ + public String getSessionId() { + return sessionId; + } + + /** + * Creates a new builder for {@code GetSnapshotOptions}. + * + * @return a new builder + */ + public static Builder builder() { + return new Builder(); + } + + /** Builder for {@link GetSnapshotOptions}. */ + public static final class Builder { + private String snapshotId; + private String sessionId; + + private Builder() {} + + /** + * Sets the specific snapshot ID to retrieve. + * + * @param snapshotId the snapshot ID + * @return this builder + */ + public Builder snapshotId(String snapshotId) { + this.snapshotId = snapshotId; + return this; + } + + /** + * Sets the session ID whose latest leaf snapshot should be retrieved. + * + * @param sessionId the session ID + * @return this builder + */ + public Builder sessionId(String sessionId) { + this.sessionId = sessionId; + return this; + } + + /** + * Builds a new {@code GetSnapshotOptions}. + * + * @return a new {@code GetSnapshotOptions} + */ + public GetSnapshotOptions build() { + return new GetSnapshotOptions(this); + } + } +} diff --git a/ai/src/main/java/com/google/genkit/ai/agent/GetSnapshotRequest.java b/ai/src/main/java/com/google/genkit/ai/agent/GetSnapshotRequest.java new file mode 100644 index 000000000..07aa28a10 --- /dev/null +++ b/ai/src/main/java/com/google/genkit/ai/agent/GetSnapshotRequest.java @@ -0,0 +1,106 @@ +/* + * 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. + * + * SPDX-License-Identifier: Apache-2.0 + */ + +package com.google.genkit.ai.agent; + +import com.fasterxml.jackson.annotation.JsonInclude; +import com.fasterxml.jackson.annotation.JsonProperty; + +/** GetSnapshotRequest is the request body for retrieving a session snapshot. */ +@JsonInclude(JsonInclude.Include.NON_NULL) +public class GetSnapshotRequest { + + @JsonProperty("snapshotId") + private String snapshotId; + + @JsonProperty("sessionId") + private String sessionId; + + /** Default constructor. */ + public GetSnapshotRequest() {} + + private GetSnapshotRequest(Builder builder) { + this.snapshotId = builder.snapshotId; + this.sessionId = builder.sessionId; + } + + /** + * Creates a builder for GetSnapshotRequest. + * + * @return a new builder + */ + public static Builder builder() { + return new Builder(); + } + + /** + * Returns the snapshot ID. + * + * @return the snapshot ID + */ + public String getSnapshotId() { + return snapshotId; + } + + /** + * Sets the snapshot ID. + * + * @param snapshotId the snapshot ID + */ + public void setSnapshotId(String snapshotId) { + this.snapshotId = snapshotId; + } + + /** + * Returns the session ID. + * + * @return the session ID + */ + public String getSessionId() { + return sessionId; + } + + /** + * Sets the session ID. + * + * @param sessionId the session ID + */ + public void setSessionId(String sessionId) { + this.sessionId = sessionId; + } + + /** Builder for GetSnapshotRequest. */ + public static class Builder { + private String snapshotId; + private String sessionId; + + public Builder snapshotId(String snapshotId) { + this.snapshotId = snapshotId; + return this; + } + + public Builder sessionId(String sessionId) { + this.sessionId = sessionId; + return this; + } + + public GetSnapshotRequest build() { + return new GetSnapshotRequest(this); + } + } +} diff --git a/ai/src/main/java/com/google/genkit/ai/agent/InMemorySessionStore.java b/ai/src/main/java/com/google/genkit/ai/agent/InMemorySessionStore.java new file mode 100644 index 000000000..caaac3b02 --- /dev/null +++ b/ai/src/main/java/com/google/genkit/ai/agent/InMemorySessionStore.java @@ -0,0 +1,313 @@ +/* + * 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. + * + * SPDX-License-Identifier: Apache-2.0 + */ + +package com.google.genkit.ai.agent; + +import com.fasterxml.jackson.databind.JsonNode; +import com.fasterxml.jackson.databind.ObjectMapper; +import com.google.genkit.ai.agent.internal.LeafSelection; +import com.google.genkit.core.GenkitException; +import com.google.genkit.core.JsonUtils; +import java.util.ArrayList; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.UUID; +import java.util.function.Consumer; +import java.util.stream.Collectors; + +/** + * Reference in-memory implementation of {@link SessionStore} and {@link SnapshotSubscriber}. + * + *

Used by conformance tests and as the simplest full implementation of the {@code saveSnapshot} + * read-modify-write contract. + * + *

Deep-copy strategy

+ * + *

Snapshots are deep-copied on both write (store) and read (return) using a JSON round-trip via + * the shared {@link ObjectMapper} ({@code valueToTree} → {@code treeToValue(node, + * SessionSnapshot.class)}). Because of Java's type erasure, the round-trip uses the raw type {@code + * SessionSnapshot} (without the {@code } parameter), so the {@code custom} field of {@link + * SessionState} will be deserialized as a {@link java.util.LinkedHashMap} rather than as {@code S}. + * This is acceptable for the in-memory store used dynamically (conformance tests use {@code + * Map} state). + * + *

Subscriber notification

+ * + *

Callbacks are collected under the lock and invoked after releasing the lock to + * prevent deadlock if a callback re-enters the store (e.g. to read the new snapshot). + * + *

Subscriber semantics (Go-compatible)

+ * + *

If the snapshot is already present when {@link #onSnapshotStateChange} is called, the callback + * is not invoked immediately. The callback fires only when a subsequent {@link + * #saveSnapshot} causes a status change (including the first save when existing is {@code null}). + * + * @param the type of custom session state + */ +public final class InMemorySessionStore implements SessionStore, SnapshotSubscriber { + + private static final ObjectMapper MAPPER = JsonUtils.getObjectMapper(); + + /** + * Backing store: snapshotId → raw (erased) SessionSnapshot. We use the raw type internally + * because the JSON round-trip erases {@code }, and the store must hold snapshots of a single + * type anyway. + */ + @SuppressWarnings("rawtypes") + private final Map snapshots = new HashMap<>(); + + /** Subscriber registry: snapshotId → list of callbacks. */ + private final Map>>> subscribers = new HashMap<>(); + + private final boolean rejectBranchingSessions; + + /** + * Creates a new {@code InMemorySessionStore} with {@code rejectBranching} defaulting to false. + */ + public InMemorySessionStore() { + this(false); + } + + /** + * Creates a new {@code InMemorySessionStore}. + * + * @param rejectBranchingSessions if {@code true}, {@link LeafSelection#selectLeaf} will throw + * when more than one leaf is detected for a session + */ + public InMemorySessionStore(boolean rejectBranchingSessions) { + this.rejectBranchingSessions = rejectBranchingSessions; + } + + // ── SnapshotReader ──────────────────────────────────────────────────────────── + + /** + * {@inheritDoc} + * + *

When {@link GetSnapshotOptions#getSnapshotId()} is set, returns a deep copy of the stored + * snapshot (or {@code null} if absent). When {@link GetSnapshotOptions#getSessionId()} is set, + * gathers all snapshots whose {@code sessionId} equals it, applies {@link LeafSelection}, and + * returns a deep copy of the result. If neither id is set, returns {@code null}. + */ + @Override + @SuppressWarnings({"unchecked", "rawtypes"}) + public SessionSnapshot getSnapshot(GetSnapshotOptions opts) { + if (opts == null) { + return null; + } + + synchronized (this) { + if (opts.getSnapshotId() != null) { + SessionSnapshot stored = snapshots.get(opts.getSnapshotId()); + return stored == null ? null : deepCopy(stored); + } + + if (opts.getSessionId() != null) { + String targetSessionId = opts.getSessionId(); + List> matching = + snapshots.values().stream() + .filter( + s -> { + // sessionId can be on the snapshot itself or on its state + String sid = s.getSessionId(); + if (sid == null && s.getState() != null) { + sid = s.getState().getSessionId(); + } + return targetSessionId.equals(sid); + }) + .map(s -> (SessionSnapshot) s) + .collect(Collectors.toList()); + + if (matching.isEmpty()) { + return null; + } + SessionSnapshot leaf = LeafSelection.selectLeaf(matching, rejectBranchingSessions); + return leaf == null ? null : deepCopy(leaf); + } + + return null; + } + } + + // ── SnapshotWriter ──────────────────────────────────────────────────────────── + + /** + * {@inheritDoc} + * + *

Implements the full saveSnapshot contract: + * + *

    + *
  1. Under lock: read existing snapshot (deep copy) identified by {@code snapshotId}. + *
  2. Invoke {@code mutator.apply(existing)}; if result is {@code null}, return {@code null} + * (no write). + *
  3. Determine final id: {@code snapshotId != null ? snapshotId : (result.snapshotId != null ? + * result.snapshotId : UUID.randomUUID())}. Force {@code result.snapshotId = finalId}. + *
  4. Preserve sessionId from existing snapshot when updating. + *
  5. Validate sessionId is non-null and non-empty; throw {@link GenkitException} with {@code + * INVALID_ARGUMENT} otherwise. + *
  6. Default {@code null} status to {@link SnapshotStatus#COMPLETED}. + *
  7. Store deep copy. Collect subscriber callbacks if status changed (existing was null or + * status differs). + *
  8. Invoke collected callbacks after releasing the lock. + *
+ */ + @Override + @SuppressWarnings({"unchecked", "rawtypes"}) + public String saveSnapshot( + String snapshotId, SnapshotMutator mutator, SessionStoreOptions options) { + // Declared outside synchronized block so callbacks can be invoked after lock release. + List>> toNotify; + SessionSnapshot notifyPayload = null; + String finalId; + + synchronized (this) { + // Step 1: read existing + SessionSnapshot existing = + snapshotId != null && snapshots.containsKey(snapshotId) + ? deepCopy((SessionSnapshot) snapshots.get(snapshotId)) + : null; + + // Step 2: apply mutator + SessionSnapshot result = mutator.apply(existing); + if (result == null) { + return null; + } + + // Step 3: determine final id + if (snapshotId != null) { + finalId = snapshotId; + } else if (result.getSnapshotId() != null && !result.getSnapshotId().isBlank()) { + finalId = result.getSnapshotId(); + } else { + finalId = UUID.randomUUID().toString(); + } + result.setSnapshotId(finalId); + + // Step 4: preserve sessionId from existing row + if (existing != null && existing.getSessionId() != null) { + result.setSessionId(existing.getSessionId()); + } + + // Step 4b: fall back to state.sessionId if top-level sessionId is null + if (result.getSessionId() == null && result.getState() != null) { + result.setSessionId(result.getState().getSessionId()); + } + + // Step 4c: validate sessionId + if (result.getSessionId() == null || result.getSessionId().isBlank()) { + throw GenkitException.builder() + .message("snapshot requires sessionId") + .errorCode("INVALID_ARGUMENT") + .build(); + } + + // Step 5: default null status to COMPLETED + if (result.getStatus() == null) { + result.setStatus(SnapshotStatus.COMPLETED); + } + + // Step 6: store deep copy + SessionSnapshot toStore = deepCopy(result); + snapshots.put(finalId, toStore); + + // Step 7: collect subscriber callbacks if status changed (invoke after releasing lock) + SnapshotStatus existingStatus = existing != null ? existing.getStatus() : null; + boolean statusChanged = existingStatus != result.getStatus(); + if (statusChanged) { + List>> cbs = subscribers.get(finalId); + if (cbs != null && !cbs.isEmpty()) { + toNotify = new ArrayList<>(cbs); + notifyPayload = deepCopy(toStore); + } else { + toNotify = new ArrayList<>(); + notifyPayload = null; + } + } else { + toNotify = new ArrayList<>(); + notifyPayload = null; + } + } // lock released here + + // Step 8: invoke callbacks outside the lock to avoid deadlock on re-entry + if (notifyPayload != null) { + for (Consumer> cb : toNotify) { + cb.accept(notifyPayload); + } + } + return finalId; + } + + // ── SnapshotSubscriber ──────────────────────────────────────────────────────── + + /** + * {@inheritDoc} + * + *

Registers {@code cb} to be invoked on subsequent status changes for {@code snapshotId}. If + * the snapshot already exists, the callback is invoked immediately with a deep copy of the + * current snapshot (before returning). The callback will also fire on any subsequent status + * changes. If the snapshot does not currently exist, the callback is registered and will fire on + * the first save (status change from null). Returns an {@link AutoCloseable} that unregisters the + * callback. + */ + @Override + @SuppressWarnings({"unchecked", "rawtypes"}) + public AutoCloseable onSnapshotStateChange( + String snapshotId, Consumer> cb, SessionStoreOptions options) { + SessionSnapshot immediate = null; + synchronized (this) { + subscribers.computeIfAbsent(snapshotId, k -> new ArrayList<>()).add(cb); + SessionSnapshot current = snapshots.get(snapshotId); + if (current != null) { + immediate = deepCopy(current); + } + } + if (immediate != null) { + cb.accept(immediate); + } + return () -> { + synchronized (this) { + List>> cbs = subscribers.get(snapshotId); + if (cbs != null) { + cbs.remove(cb); + if (cbs.isEmpty()) { + subscribers.remove(snapshotId); + } + } + } + }; + } + + // ── deep copy ──────────────────────────────────────────────────────────────── + + /** + * Deep-copies a snapshot via JSON round-trip (valueToTree → treeToValue with raw type). + * + *

Note: due to type erasure, the {@code custom} field of {@link SessionState} is deserialized + * as {@link java.util.LinkedHashMap} rather than {@code S}. This is documented and acceptable for + * the in-memory reference store. + */ + @SuppressWarnings({"unchecked", "rawtypes"}) + private SessionSnapshot deepCopy(SessionSnapshot snapshot) { + try { + JsonNode node = MAPPER.valueToTree(snapshot); + return (SessionSnapshot) MAPPER.treeToValue(node, SessionSnapshot.class); + } catch (Exception e) { + throw new GenkitException("Failed to deep-copy SessionSnapshot: " + e.getMessage(), e); + } + } +} diff --git a/ai/src/main/java/com/google/genkit/ai/agent/RuntimeError.java b/ai/src/main/java/com/google/genkit/ai/agent/RuntimeError.java new file mode 100644 index 000000000..9cb75b22c --- /dev/null +++ b/ai/src/main/java/com/google/genkit/ai/agent/RuntimeError.java @@ -0,0 +1,134 @@ +/* + * 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. + * + * SPDX-License-Identifier: Apache-2.0 + */ + +package com.google.genkit.ai.agent; + +import com.fasterxml.jackson.annotation.JsonInclude; +import com.fasterxml.jackson.annotation.JsonProperty; + +/** RuntimeError represents a runtime error that occurred during agent execution. */ +@JsonInclude(JsonInclude.Include.NON_NULL) +public class RuntimeError { + + @JsonProperty("status") + private String status; + + @JsonProperty("message") + private String message; + + @JsonProperty("details") + private Object details; + + /** Default constructor. */ + public RuntimeError() {} + + private RuntimeError(Builder builder) { + this.status = builder.status; + this.message = builder.message; + this.details = builder.details; + } + + /** + * Creates a builder for RuntimeError. + * + * @return a new builder + */ + public static Builder builder() { + return new Builder(); + } + + /** + * Returns the status code of the error. + * + * @return the status + */ + public String getStatus() { + return status; + } + + /** + * Sets the status code. + * + * @param status the status code + */ + public void setStatus(String status) { + this.status = status; + } + + /** + * Returns the error message. + * + * @return the message + */ + public String getMessage() { + return message; + } + + /** + * Sets the error message. + * + * @param message the message + */ + public void setMessage(String message) { + this.message = message; + } + + /** + * Returns additional error details. + * + * @return the details + */ + public Object getDetails() { + return details; + } + + /** + * Sets additional error details. + * + * @param details the details + */ + public void setDetails(Object details) { + this.details = details; + } + + /** Builder for RuntimeError. */ + public static class Builder { + private String status; + private String message; + private Object details; + + public Builder status(String status) { + this.status = status; + return this; + } + + public Builder message(String message) { + this.message = message; + return this; + } + + public Builder details(Object details) { + this.details = details; + return this; + } + + public RuntimeError build() { + return new RuntimeError(this); + } + } +} diff --git a/ai/src/main/java/com/google/genkit/ai/agent/Session.java b/ai/src/main/java/com/google/genkit/ai/agent/Session.java new file mode 100644 index 000000000..1b6bc6818 --- /dev/null +++ b/ai/src/main/java/com/google/genkit/ai/agent/Session.java @@ -0,0 +1,349 @@ +/* + * 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. + * + * SPDX-License-Identifier: Apache-2.0 + */ + +package com.google.genkit.ai.agent; + +import com.fasterxml.jackson.core.type.TypeReference; +import com.fasterxml.jackson.databind.ObjectMapper; +import com.google.genkit.ai.Message; +import com.google.genkit.core.JsonUtils; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.List; +import java.util.UUID; +import java.util.function.Consumer; +import java.util.function.UnaryOperator; + +/** + * Session is the per-invocation in-memory state holder managed by the agent runtime. It stores + * messages, custom state, and artifacts; provides mutation methods that increment a version counter + * and fire listener callbacks; and implements {@link ArtifactStore} for middleware/tools that do + * not need to know the custom state type. + * + *

Thread-safety: the runtime drives a session from one turn at a time. Mutations are {@code + * synchronized} to guard against concurrent reads in edge cases, but over-engineering with + * lock-free structures is intentionally avoided. + * + * @param the type of the custom state object + */ +public final class Session implements ArtifactStore { + + private static final ObjectMapper MAPPER = JsonUtils.getObjectMapper(); + + private final String sessionId; + + /** Internal messages list — never exposed directly. */ + private final List messages; + + /** Current custom state. */ + private S custom; + + /** Internal artifacts list — never exposed directly. */ + private final List artifacts; + + /** Monotonically increasing mutation counter. */ + private long version; + + /** Callback fired after updateCustom completes. */ + private Runnable onCustomChanged; + + /** Callback fired per artifact added or updated. */ + private Consumer onArtifactChanged; + + /** + * Constructs a new Session from the given initial state. Deep-copies {@code initialState} so that + * subsequent external mutations to the original object cannot affect this session. If {@code + * initialState.sessionId} is null or empty, a new UUID is minted. + * + * @param initialState the initial session state (must not be null) + */ + public Session(SessionState initialState) { + // Mint or preserve sessionId + String sid = initialState.getSessionId(); + this.sessionId = (sid != null && !sid.isEmpty()) ? sid : UUID.randomUUID().toString(); + + // Deep-copy messages + List srcMessages = initialState.getMessages(); + this.messages = srcMessages != null ? deepCopyMessages(srcMessages) : new ArrayList<>(); + + // Custom state: a shallow reference is acceptable for simple values; for Map/POJO we rely on + // the caller providing an immutable or properly-owned object. Deep-copy via ObjectMapper is + // used to avoid aliasing issues when S is a Map or POJO. + this.custom = deepCopyCustom(initialState.getCustom()); + + // Deep-copy artifacts + List srcArtifacts = initialState.getArtifacts(); + this.artifacts = srcArtifacts != null ? deepCopyArtifacts(srcArtifacts) : new ArrayList<>(); + + this.version = 0L; + } + + // ---- Identifiers ---- + + /** + * Returns the session ID. + * + * @return the session ID (never null) + */ + public String sessionId() { + return sessionId; + } + + // ---- Version ---- + + /** + * Returns the current mutation version counter. Increments on every mutating operation. + * + * @return the version + */ + public synchronized long getVersion() { + return version; + } + + // ---- State snapshot ---- + + /** + * Returns a deep copy of the current session state. Callers must not mutate the returned object + * to avoid leaking changes back into the session. + * + * @return a deep copy of the current {@link SessionState} + */ + public synchronized SessionState getState() { + SessionState snapshot = new SessionState<>(); + snapshot.setSessionId(sessionId); + snapshot.setMessages(new ArrayList<>(messages)); + snapshot.setCustom(deepCopyCustom(custom)); + snapshot.setArtifacts(new ArrayList<>(artifacts)); + return snapshot; + } + + // ---- Messages ---- + + /** + * Returns a copy of the current messages list. Mutating the returned list does not affect the + * session. + * + * @return a copy of the messages (never null) + */ + public synchronized List getMessages() { + return new ArrayList<>(messages); + } + + /** + * Appends the given messages to the session. + * + * @param msgs messages to add + */ + public synchronized void addMessages(Message... msgs) { + messages.addAll(Arrays.asList(msgs)); + version++; + } + + /** + * Appends the given messages to the session. + * + * @param msgs messages to add (must not be null) + */ + public synchronized void addMessages(List msgs) { + messages.addAll(msgs); + version++; + } + + /** + * Replaces the session's message list with a copy of the given list. + * + * @param msgs the new messages (must not be null) + */ + public synchronized void setMessages(List msgs) { + messages.clear(); + messages.addAll(msgs); + version++; + } + + /** + * Applies {@code fn} to the current messages list and stores the result. + * + * @param fn the update function + */ + public synchronized void updateMessages(UnaryOperator> fn) { + List updated = fn.apply(new ArrayList<>(messages)); + messages.clear(); + if (updated != null) { + messages.addAll(updated); + } + version++; + } + + // ---- Custom state ---- + + /** + * Returns a deep copy of the current custom state. Mutating the returned value does not affect + * the session — prefer {@link #updateCustom} for modifications. + * + * @return a deep copy of the current custom state, or null if not set + */ + public synchronized S getCustom() { + return deepCopyCustom(custom); + } + + /** + * Applies {@code fn} to the current custom state, stores the result, increments the version, and + * fires the {@code onCustomChanged} listener. + * + * @param fn the update function + */ + public synchronized void updateCustom(UnaryOperator fn) { + custom = fn.apply(custom); + version++; + if (onCustomChanged != null) { + onCustomChanged.run(); + } + } + + // ---- Artifacts (ArtifactStore implementation) ---- + + /** + * Returns a copy of the current artifacts list. Mutating the returned list does not affect the + * session. + * + * @return a copy of the artifacts (never null) + */ + @Override + public synchronized List getArtifacts() { + return new ArrayList<>(artifacts); + } + + /** + * Adds artifacts with deduplication by name. If an artifact with the same non-null name already + * exists, it is replaced in place. Artifacts with null names are always appended. Fires the + * {@code onArtifactChanged} listener for each artifact added or updated. + * + * @param arts artifacts to add + */ + @Override + public synchronized void addArtifacts(Artifact... arts) { + addArtifactList(Arrays.asList(arts)); + } + + /** + * Adds artifacts with deduplication by name (list overload). + * + * @param arts artifacts to add (must not be null) + */ + public synchronized void addArtifacts(List arts) { + addArtifactList(arts); + } + + private void addArtifactList(List arts) { + for (Artifact incoming : arts) { + String name = incoming.getName(); + if (name != null) { + // Replace in place if name matches + boolean replaced = false; + for (int i = 0; i < artifacts.size(); i++) { + if (name.equals(artifacts.get(i).getName())) { + artifacts.set(i, incoming); + replaced = true; + break; + } + } + if (!replaced) { + artifacts.add(incoming); + } + } else { + // Null name: always append + artifacts.add(incoming); + } + version++; + if (onArtifactChanged != null) { + onArtifactChanged.accept(incoming); + } + } + } + + /** + * Applies {@code fn} to the current artifacts list and stores the result. + * + * @param fn the update function + */ + public synchronized void updateArtifacts(UnaryOperator> fn) { + List updated = fn.apply(new ArrayList<>(artifacts)); + artifacts.clear(); + if (updated != null) { + artifacts.addAll(updated); + } + version++; + } + + // ---- Listener hooks ---- + + /** + * Sets the callback to be invoked after {@link #updateCustom} completes. + * + * @param cb the callback (may be null to clear) + */ + public synchronized void setOnCustomChanged(Runnable cb) { + this.onCustomChanged = cb; + } + + /** + * Sets the callback to be invoked for each artifact added or updated via {@link #addArtifacts}. + * + * @param cb the callback (may be null to clear) + */ + public synchronized void setOnArtifactChanged(Consumer cb) { + this.onArtifactChanged = cb; + } + + // ---- Deep-copy helpers ---- + + @SuppressWarnings("unchecked") + private S deepCopyCustom(S value) { + if (value == null) { + return null; + } + try { + // Serialize to JSON and back to obtain a disconnected copy. + String json = MAPPER.writeValueAsString(value); + return (S) MAPPER.readValue(json, value.getClass()); + } catch (Exception e) { + // Inner generic type params are erased here; complex parameterized custom state may not + // be fully deep-copied. + // Fall back to returning the original reference if serialization is not possible. + return value; + } + } + + private List deepCopyMessages(List src) { + try { + String json = MAPPER.writeValueAsString(src); + return MAPPER.readValue(json, new TypeReference>() {}); + } catch (Exception e) { + return new ArrayList<>(src); + } + } + + private List deepCopyArtifacts(List src) { + try { + String json = MAPPER.writeValueAsString(src); + return MAPPER.readValue(json, new TypeReference>() {}); + } catch (Exception e) { + return new ArrayList<>(src); + } + } +} diff --git a/ai/src/main/java/com/google/genkit/ai/agent/SessionRunner.java b/ai/src/main/java/com/google/genkit/ai/agent/SessionRunner.java new file mode 100644 index 000000000..ba85eef5b --- /dev/null +++ b/ai/src/main/java/com/google/genkit/ai/agent/SessionRunner.java @@ -0,0 +1,448 @@ +/* + * 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. + * + * SPDX-License-Identifier: Apache-2.0 + */ + +package com.google.genkit.ai.agent; + +import com.google.genkit.ai.Message; +import com.google.genkit.ai.Part; +import com.google.genkit.ai.Role; +import com.google.genkit.ai.agent.internal.AbortAwareMutator; +import com.google.genkit.core.GenkitException; +import java.time.Instant; +import java.util.List; +import java.util.UUID; +import java.util.function.UnaryOperator; + +/** + * Drives ONE turn of an agent session: validates the input, appends the user message, runs the turn + * body, and persists a snapshot on success ({@link SnapshotStatus#COMPLETED}) or on failure ({@link + * SnapshotStatus#FAILED}). Failure snapshots are swallowed — {@link #runTurn} never rethrows + * exceptions from the turn body. + * + *

Client-managed mode ({@code store == null}): no persistence; the runner tracks the last + * snapshot in-memory only and {@link #lastSnapshotId()} always returns {@code ""}. + * + * @param the type of the custom session state + */ +public final class SessionRunner { + + private final Session session; + + /** Nullable: null means client-managed (no server-side persistence). */ + private final SessionStore store; + + private final SessionStoreOptions opts; + + private int turnIndex; + private SessionSnapshot lastSnapshot; + private String lastSnapshotId; + private SessionState lastGoodState; + private AgentFinishReason lastTurnFinishReason; + private RuntimeError lastTurnError; + + /** + * Constructs a new SessionRunner. + * + * @param session the in-memory session to drive + * @param store the session store, or {@code null} for client-managed mode (no server persistence) + * @param opts options forwarded to store operations + */ + public SessionRunner(Session session, SessionStore store, SessionStoreOptions opts) { + this(session, store, opts, null); + } + + /** + * Constructs a new SessionRunner seeded with the id of the snapshot the session was resumed from. + * + *

Seeding {@link #lastSnapshotId} makes the first post-resume turn chain its {@code parentId} + * to the resumed snapshot (the {@code parentId = lastSnapshotId} logic in {@link #maybeSnapshot} + * then threads correctly across invocations). Pass {@code null} for a fresh session. + * + * @param session the in-memory session to drive + * @param store the session store, or {@code null} for client-managed mode (no server persistence) + * @param opts options forwarded to store operations + * @param seedSnapshotId the id of the snapshot this session was resumed from, or {@code null} + */ + public SessionRunner( + Session session, SessionStore store, SessionStoreOptions opts, String seedSnapshotId) { + this.session = session; + this.store = store; + this.opts = opts != null ? opts : SessionStoreOptions.empty(); + if (seedSnapshotId != null && !seedSnapshotId.isEmpty()) { + this.lastSnapshotId = seedSnapshotId; + } + } + + // ── Session delegates ──────────────────────────────────────────────────────── + + /** + * Returns the session ID. + * + * @return the session ID (never null) + */ + public String sessionId() { + return session.sessionId(); + } + + /** + * Returns a deep copy of the current session state. + * + * @return the current session state + */ + public SessionState getState() { + return session.getState(); + } + + /** + * Returns a copy of the current messages list. + * + * @return the current messages + */ + public List getMessages() { + return session.getMessages(); + } + + /** + * Appends messages to the session. + * + * @param m messages to add + */ + public void addMessages(Message... m) { + session.addMessages(m); + } + + /** + * Returns a deep copy of the current custom state. + * + * @return the custom state + */ + public S getCustom() { + return session.getCustom(); + } + + /** + * Applies {@code fn} to the current custom state and stores the result. + * + * @param fn the update function + */ + public void updateCustom(UnaryOperator fn) { + session.updateCustom(fn); + } + + /** + * Returns a copy of the current artifacts list. + * + * @return the artifacts + */ + public List getArtifacts() { + return session.getArtifacts(); + } + + /** + * Adds artifacts to the session. + * + * @param a artifacts to add + */ + public void addArtifacts(Artifact... a) { + session.addArtifacts(a); + } + + /** + * Returns the underlying session. + * + * @return the session + */ + public Session session() { + return session; + } + + // ── Turn-state accessors ───────────────────────────────────────────────────── + + /** + * Returns the number of turns that have completed (including failed turns). + * + * @return the turn index (0 before any turn has run) + */ + public int turnIndex() { + return turnIndex; + } + + /** + * Returns the last persisted (or in-memory) snapshot, or {@code null} if no turn has completed. + * + * @return the last snapshot + */ + public SessionSnapshot lastSnapshot() { + return lastSnapshot; + } + + /** + * Returns the ID of the last snapshot, or {@code null} if no turn has completed. Returns {@code + * ""} (empty string) in client-managed mode. + * + * @return the last snapshot ID + */ + public String lastSnapshotId() { + return lastSnapshotId; + } + + /** + * Returns the session state as of the last successful turn, or {@code null} if no successful turn + * has completed. + * + * @return the last good state + */ + public SessionState lastGoodState() { + return lastGoodState; + } + + /** + * Returns the finish reason of the last completed turn. + * + * @return the finish reason, or {@code null} if no turn has run + */ + public AgentFinishReason lastTurnFinishReason() { + return lastTurnFinishReason; + } + + /** + * Returns the error from the last failed turn, or {@code null} if the last turn succeeded. + * + * @return the last turn error + */ + public RuntimeError lastTurnError() { + return lastTurnError; + } + + // ── Core turn lifecycle ────────────────────────────────────────────────────── + + /** + * Runs ONE turn: + * + *

    + *
  1. Validates {@code input.message}: role must be null/USER; rejects tool-request and + * tool-response parts → throws {@link GenkitException} with {@code INVALID_ARGUMENT} (API + * misuse; not graceful). + *
  2. Reserves a turn snapshot ID (UUID if store != null, else ""). + *
  3. Appends {@code input.message} to the session if non-null. + *
  4. Runs {@code turnBody.run(input, turnCtx)}. + *
  5. On success: persists a {@code COMPLETED} snapshot; records {@link #lastGoodState}; + * increments turn index. + *
  6. On exception: records {@link #lastTurnError}; persists a {@code FAILED} snapshot (or + * {@code ABORTED} for {@link InterruptedException}); increments turn index; does NOT + * rethrow. + *
+ * + * @param input the agent input for this turn + * @param turnBody the turn body to execute + * @throws GenkitException (INVALID_ARGUMENT) if the input message is invalid — this IS propagated + * (it is API misuse, not a graceful turn failure) + */ + public void runTurn(AgentInput input, TurnBody turnBody) { + // Step 1: validate input + validateInput(input); + + // Step 2: reserve snapshot ID + String turnSnapshotId = (store != null) ? UUID.randomUUID().toString() : ""; + String parentId = lastSnapshotId; // may be null on first turn + TurnContext turnCtx = new TurnContext(turnSnapshotId, parentId, turnIndex); + + // Step 3: append user message (deep-copied to prevent external mutation of history) + if (input != null && input.getMessage() != null) { + session.addMessages(deepCopyMessage(input.getMessage())); + } + + // Steps 4-6: run turnBody + try { + AgentFinishReason finishReason = turnBody.run(input, turnCtx); + if (finishReason == null) { + finishReason = AgentFinishReason.STOP; + } + // Step 5: success path + lastTurnError = null; + lastTurnFinishReason = finishReason; + maybeSnapshot(SnapshotStatus.COMPLETED, finishReason, null, turnSnapshotId, parentId); + lastGoodState = session.getState(); + } catch (InterruptedException ie) { + // ABORTED — record but do NOT write a failed snapshot for aborted turns + Thread.currentThread().interrupt(); // restore interrupted status + lastTurnError = + RuntimeError.builder() + .status("ABORTED") + .message(ie.getMessage() != null ? ie.getMessage() : "interrupted") + .build(); + lastTurnFinishReason = AgentFinishReason.ABORTED; + maybeSnapshot( + SnapshotStatus.ABORTED, + AgentFinishReason.ABORTED, + lastTurnError, + turnSnapshotId, + parentId); + } catch (Exception e) { + // Step 6: failure path — record error, persist FAILED snapshot, swallow. Preserve a + // GenkitException's error code (e.g. INVALID_ARGUMENT from resume-directive validation) so + // the + // graceful FAILED output carries the correct status; default to INTERNAL otherwise. + String status = "INTERNAL"; + Throwable probe = e; + while (probe != null) { + if (probe instanceof GenkitException ge && ge.getErrorCode() != null) { + status = ge.getErrorCode(); + break; + } + probe = probe.getCause(); + } + lastTurnError = + RuntimeError.builder() + .status(status) + .message(e.getMessage() != null ? e.getMessage() : e.getClass().getName()) + .build(); + lastTurnFinishReason = AgentFinishReason.FAILED; + maybeSnapshot( + SnapshotStatus.FAILED, AgentFinishReason.FAILED, lastTurnError, turnSnapshotId, parentId); + } finally { + turnIndex++; + } + } + + // ── Private helpers ────────────────────────────────────────────────────────── + + /** + * Deep-copies a message using the shared ObjectMapper. Returns null if input is null. + * + * @param m the message to copy + * @return a deep copy of the message + * @throws GenkitException if copying fails + */ + private Message deepCopyMessage(Message m) { + if (m == null) return null; + var mapper = com.google.genkit.core.JsonUtils.getObjectMapper(); + try { + return mapper.treeToValue(mapper.valueToTree(m), Message.class); + } catch (Exception e) { + throw GenkitException.builder() + .message("failed to copy message: " + e.getMessage()) + .errorCode("INTERNAL") + .build(); + } + } + + /** + * Validates the input message. Throws {@link GenkitException} with {@code INVALID_ARGUMENT} if: + * + *
    + *
  • The message role is not null, empty, or USER. + *
  • The message contains tool-request or tool-response parts (those go via resume). + *
+ */ + private void validateInput(AgentInput input) { + if (input == null || input.getMessage() == null) { + return; + } + Message msg = input.getMessage(); + + // Validate role + Role role = msg.getRole(); + if (role != null && role != Role.USER) { + throw GenkitException.builder() + .message("input message role must be 'user' (or null); got: " + role) + .errorCode("INVALID_ARGUMENT") + .build(); + } + + // Reject tool-request and tool-response parts on the user message + List parts = msg.getContent(); + if (parts != null) { + for (Part part : parts) { + if (part.getToolRequest() != null) { + throw GenkitException.builder() + .message( + "user message must not contain toolRequest parts; use resume for tool responses") + .errorCode("INVALID_ARGUMENT") + .build(); + } + if (part.getToolResponse() != null) { + throw GenkitException.builder() + .message( + "user message must not contain toolResponse parts; use resume for tool responses") + .errorCode("INVALID_ARGUMENT") + .build(); + } + } + } + } + + /** + * Persists (or tracks in-memory for client-managed mode) a snapshot for this turn. + * + *

If {@code store == null} (client-managed): no persistence; builds a synthetic snapshot and + * stores it in {@link #lastSnapshot}. {@link #lastSnapshotId} stays {@code ""}. + * + *

If {@code store != null}: builds a {@link SessionSnapshot} and calls {@link + * SessionStore#saveSnapshot} with an {@link AbortAwareMutator}-wrapped mutator. Updates {@link + * #lastSnapshot} and {@link #lastSnapshotId} to the persisted values. + */ + private void maybeSnapshot( + SnapshotStatus status, + AgentFinishReason finishReason, + RuntimeError error, + String snapshotId, + String parentId) { + + String now = Instant.now().toString(); + SessionState currentState = session.getState(); + + if (store == null) { + // Client-managed: track in memory only, snapshotId stays "" + SessionSnapshot snap = + SessionSnapshot.builder() + .snapshotId("") + .sessionId(session.sessionId()) + .parentId(parentId) + .createdAt(now) + .updatedAt(now) + .state(currentState) + .status(status) + .finishReason(finishReason) + .error(error) + .build(); + lastSnapshot = snap; + lastSnapshotId = ""; + } else { + // Persist via store + final SessionSnapshot snap = + SessionSnapshot.builder() + .snapshotId(snapshotId) + .sessionId(session.sessionId()) + .parentId(parentId) + .createdAt(now) + .updatedAt(now) + .state(currentState) + .status(status) + .finishReason(finishReason) + .error(error) + .build(); + + SnapshotMutator inner = existing -> snap; + String savedId = store.saveSnapshot(snapshotId, AbortAwareMutator.wrap(inner), opts); + + lastSnapshot = snap; + lastSnapshotId = savedId != null ? savedId : snapshotId; + } + } +} diff --git a/ai/src/main/java/com/google/genkit/ai/agent/SessionSnapshot.java b/ai/src/main/java/com/google/genkit/ai/agent/SessionSnapshot.java new file mode 100644 index 000000000..368acf278 --- /dev/null +++ b/ai/src/main/java/com/google/genkit/ai/agent/SessionSnapshot.java @@ -0,0 +1,345 @@ +/* + * 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. + * + * SPDX-License-Identifier: Apache-2.0 + */ + +package com.google.genkit.ai.agent; + +import com.fasterxml.jackson.annotation.JsonInclude; +import com.fasterxml.jackson.annotation.JsonProperty; + +/** + * SessionSnapshot represents a point-in-time snapshot of an agent session. + * + *

Timestamps ({@code createdAt}, {@code updatedAt}, {@code heartbeatAt}) are stored as RFC-3339 + * strings for exact wire fidelity. + * + * @param the type of custom state + */ +@JsonInclude(JsonInclude.Include.NON_NULL) +public class SessionSnapshot { + + @JsonProperty("snapshotId") + private String snapshotId; + + @JsonProperty("sessionId") + private String sessionId; + + @JsonProperty("parentId") + private String parentId; + + /** RFC-3339 timestamp string. */ + @JsonProperty("createdAt") + private String createdAt; + + /** RFC-3339 timestamp string. */ + @JsonProperty("updatedAt") + private String updatedAt; + + /** RFC-3339 timestamp string. */ + @JsonProperty("heartbeatAt") + private String heartbeatAt; + + @JsonProperty("status") + private SnapshotStatus status; + + @JsonProperty("finishReason") + private AgentFinishReason finishReason; + + @JsonProperty("error") + private RuntimeError error; + + @JsonProperty("state") + private SessionState state; + + /** Default constructor. */ + public SessionSnapshot() {} + + private SessionSnapshot(Builder builder) { + this.snapshotId = builder.snapshotId; + this.sessionId = builder.sessionId; + this.parentId = builder.parentId; + this.createdAt = builder.createdAt; + this.updatedAt = builder.updatedAt; + this.heartbeatAt = builder.heartbeatAt; + this.status = builder.status; + this.finishReason = builder.finishReason; + this.error = builder.error; + this.state = builder.state; + } + + /** + * Creates a builder for SessionSnapshot. + * + * @param the type of custom state + * @return a new builder + */ + public static Builder builder() { + return new Builder<>(); + } + + /** + * Returns the snapshot ID. + * + * @return the snapshot ID + */ + public String getSnapshotId() { + return snapshotId; + } + + /** + * Sets the snapshot ID. + * + * @param snapshotId the snapshot ID + */ + public void setSnapshotId(String snapshotId) { + this.snapshotId = snapshotId; + } + + /** + * Returns the session ID. + * + * @return the session ID + */ + public String getSessionId() { + return sessionId; + } + + /** + * Sets the session ID. + * + * @param sessionId the session ID + */ + public void setSessionId(String sessionId) { + this.sessionId = sessionId; + } + + /** + * Returns the parent snapshot ID. + * + * @return the parent ID + */ + public String getParentId() { + return parentId; + } + + /** + * Sets the parent snapshot ID. + * + * @param parentId the parent ID + */ + public void setParentId(String parentId) { + this.parentId = parentId; + } + + /** + * Returns the creation timestamp (RFC-3339). + * + * @return the created-at timestamp + */ + public String getCreatedAt() { + return createdAt; + } + + /** + * Sets the creation timestamp. + * + * @param createdAt the created-at timestamp (RFC-3339) + */ + public void setCreatedAt(String createdAt) { + this.createdAt = createdAt; + } + + /** + * Returns the last-update timestamp (RFC-3339). + * + * @return the updated-at timestamp + */ + public String getUpdatedAt() { + return updatedAt; + } + + /** + * Sets the last-update timestamp. + * + * @param updatedAt the updated-at timestamp (RFC-3339) + */ + public void setUpdatedAt(String updatedAt) { + this.updatedAt = updatedAt; + } + + /** + * Returns the last-heartbeat timestamp (RFC-3339). + * + * @return the heartbeat-at timestamp + */ + public String getHeartbeatAt() { + return heartbeatAt; + } + + /** + * Sets the last-heartbeat timestamp. + * + * @param heartbeatAt the heartbeat-at timestamp (RFC-3339) + */ + public void setHeartbeatAt(String heartbeatAt) { + this.heartbeatAt = heartbeatAt; + } + + /** + * Returns the snapshot status. + * + * @return the status + */ + public SnapshotStatus getStatus() { + return status; + } + + /** + * Sets the snapshot status. + * + * @param status the status + */ + public void setStatus(SnapshotStatus status) { + this.status = status; + } + + /** + * Returns the finish reason. + * + * @return the finish reason + */ + public AgentFinishReason getFinishReason() { + return finishReason; + } + + /** + * Sets the finish reason. + * + * @param finishReason the finish reason + */ + public void setFinishReason(AgentFinishReason finishReason) { + this.finishReason = finishReason; + } + + /** + * Returns the runtime error, if any. + * + * @return the error, or null if no error + */ + public RuntimeError getError() { + return error; + } + + /** + * Sets the runtime error. + * + * @param error the error + */ + public void setError(RuntimeError error) { + this.error = error; + } + + /** + * Returns the session state at this snapshot. + * + * @return the state + */ + public SessionState getState() { + return state; + } + + /** + * Sets the session state. + * + * @param state the state + */ + public void setState(SessionState state) { + this.state = state; + } + + /** + * Builder for SessionSnapshot. + * + * @param the type of custom state + */ + public static class Builder { + private String snapshotId; + private String sessionId; + private String parentId; + private String createdAt; + private String updatedAt; + private String heartbeatAt; + private SnapshotStatus status; + private AgentFinishReason finishReason; + private RuntimeError error; + private SessionState state; + + public Builder snapshotId(String snapshotId) { + this.snapshotId = snapshotId; + return this; + } + + public Builder sessionId(String sessionId) { + this.sessionId = sessionId; + return this; + } + + public Builder parentId(String parentId) { + this.parentId = parentId; + return this; + } + + public Builder createdAt(String createdAt) { + this.createdAt = createdAt; + return this; + } + + public Builder updatedAt(String updatedAt) { + this.updatedAt = updatedAt; + return this; + } + + public Builder heartbeatAt(String heartbeatAt) { + this.heartbeatAt = heartbeatAt; + return this; + } + + public Builder status(SnapshotStatus status) { + this.status = status; + return this; + } + + public Builder finishReason(AgentFinishReason finishReason) { + this.finishReason = finishReason; + return this; + } + + public Builder error(RuntimeError error) { + this.error = error; + return this; + } + + public Builder state(SessionState state) { + this.state = state; + return this; + } + + public SessionSnapshot build() { + return new SessionSnapshot<>(this); + } + } +} diff --git a/ai/src/main/java/com/google/genkit/ai/agent/SessionState.java b/ai/src/main/java/com/google/genkit/ai/agent/SessionState.java new file mode 100644 index 000000000..9e2996825 --- /dev/null +++ b/ai/src/main/java/com/google/genkit/ai/agent/SessionState.java @@ -0,0 +1,174 @@ +/* + * 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. + * + * SPDX-License-Identifier: Apache-2.0 + */ + +package com.google.genkit.ai.agent; + +import com.fasterxml.jackson.annotation.JsonInclude; +import com.fasterxml.jackson.annotation.JsonProperty; +import com.google.genkit.ai.Message; +import java.util.ArrayList; +import java.util.List; + +/** + * SessionState represents the state of an agent session. + * + * @param the type of custom state + */ +@JsonInclude(JsonInclude.Include.NON_NULL) +public class SessionState { + + @JsonProperty("sessionId") + private String sessionId; + + @JsonProperty("messages") + private List messages; + + @JsonProperty("custom") + private S custom; + + @JsonProperty("artifacts") + private List artifacts; + + /** Default constructor. */ + public SessionState() {} + + private SessionState(Builder builder) { + this.sessionId = builder.sessionId; + this.messages = builder.messages; + this.custom = builder.custom; + this.artifacts = builder.artifacts; + } + + /** + * Creates a builder for SessionState. + * + * @param the type of custom state + * @return a new builder + */ + public static Builder builder() { + return new Builder<>(); + } + + /** + * Returns the session ID. + * + * @return the session ID + */ + public String getSessionId() { + return sessionId; + } + + /** + * Sets the session ID. + * + * @param sessionId the session ID + */ + public void setSessionId(String sessionId) { + this.sessionId = sessionId; + } + + /** + * Returns the session messages. + * + * @return the messages + */ + public List getMessages() { + return messages; + } + + /** + * Sets the session messages. + * + * @param messages the messages + */ + public void setMessages(List messages) { + this.messages = messages != null ? new ArrayList<>(messages) : null; + } + + /** + * Returns the custom state. + * + * @return the custom state + */ + public S getCustom() { + return custom; + } + + /** + * Sets the custom state. + * + * @param custom the custom state + */ + public void setCustom(S custom) { + this.custom = custom; + } + + /** + * Returns the session artifacts. + * + * @return the artifacts + */ + public List getArtifacts() { + return artifacts; + } + + /** + * Sets the session artifacts. + * + * @param artifacts the artifacts + */ + public void setArtifacts(List artifacts) { + this.artifacts = artifacts != null ? new ArrayList<>(artifacts) : null; + } + + /** + * Builder for SessionState. + * + * @param the type of custom state + */ + public static class Builder { + private String sessionId; + private List messages; + private S custom; + private List artifacts; + + public Builder sessionId(String sessionId) { + this.sessionId = sessionId; + return this; + } + + public Builder messages(List messages) { + this.messages = messages; + return this; + } + + public Builder custom(S custom) { + this.custom = custom; + return this; + } + + public Builder artifacts(List artifacts) { + this.artifacts = artifacts; + return this; + } + + public SessionState build() { + return new SessionState<>(this); + } + } +} diff --git a/ai/src/main/java/com/google/genkit/ai/agent/SessionStore.java b/ai/src/main/java/com/google/genkit/ai/agent/SessionStore.java new file mode 100644 index 000000000..da30dc8ef --- /dev/null +++ b/ai/src/main/java/com/google/genkit/ai/agent/SessionStore.java @@ -0,0 +1,32 @@ +/* + * 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. + * + * SPDX-License-Identifier: Apache-2.0 + */ + +package com.google.genkit.ai.agent; + +/** + * Combined read/write interface for a session store. + * + *

Implementations must satisfy both the {@link SnapshotReader} and {@link SnapshotWriter} + * contracts. Optional subscription support is provided by {@link SnapshotSubscriber}. + * + * @param the type of custom session state + * @see SnapshotReader + * @see SnapshotWriter + * @see SnapshotSubscriber + */ +public interface SessionStore extends SnapshotReader, SnapshotWriter {} diff --git a/ai/src/main/java/com/google/genkit/ai/agent/SessionStoreOptions.java b/ai/src/main/java/com/google/genkit/ai/agent/SessionStoreOptions.java new file mode 100644 index 000000000..a1efa602c --- /dev/null +++ b/ai/src/main/java/com/google/genkit/ai/agent/SessionStoreOptions.java @@ -0,0 +1,66 @@ +/* + * 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. + * + * SPDX-License-Identifier: Apache-2.0 + */ + +package com.google.genkit.ai.agent; + +/** + * Options passed to session store operations. + * + *

Currently a thin options holder that provides room for future context (e.g. deadline, + * credentials). Use {@link #empty()} for a default instance. + */ +public final class SessionStoreOptions { + + /** A shared default empty instance. */ + private static final SessionStoreOptions EMPTY = new SessionStoreOptions(new Builder()); + + private SessionStoreOptions(Builder builder) {} + + /** + * Returns a default empty {@code SessionStoreOptions}. + * + * @return the default instance + */ + public static SessionStoreOptions empty() { + return EMPTY; + } + + /** + * Creates a new builder for {@code SessionStoreOptions}. + * + * @return a new builder + */ + public static Builder builder() { + return new Builder(); + } + + /** Builder for {@link SessionStoreOptions}. */ + public static final class Builder { + + private Builder() {} + + /** + * Builds a new {@code SessionStoreOptions}. + * + * @return a new {@code SessionStoreOptions} + */ + public SessionStoreOptions build() { + return new SessionStoreOptions(this); + } + } +} diff --git a/ai/src/main/java/com/google/genkit/ai/agent/SnapshotMutator.java b/ai/src/main/java/com/google/genkit/ai/agent/SnapshotMutator.java new file mode 100644 index 000000000..22bde822a --- /dev/null +++ b/ai/src/main/java/com/google/genkit/ai/agent/SnapshotMutator.java @@ -0,0 +1,42 @@ +/* + * 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. + * + * SPDX-License-Identifier: Apache-2.0 + */ + +package com.google.genkit.ai.agent; + +/** + * The mutator for an atomic read-modify-write on a session snapshot. + * + *

Receives the existing snapshot (or {@code null} if none exists) and returns the snapshot to + * persist, or {@code null} to decline / no-op. + * + *

Contract: implementations MUST be pure functions — the store may retry the + * mutator on transient conflicts, so side effects are not safe. + * + * @param the type of custom session state + */ +@FunctionalInterface +public interface SnapshotMutator { + + /** + * Applies the mutation. + * + * @param existing the existing snapshot, or {@code null} if no snapshot exists yet + * @return the snapshot to persist, or {@code null} to decline (no write performed) + */ + SessionSnapshot apply(SessionSnapshot existing); +} diff --git a/ai/src/main/java/com/google/genkit/ai/agent/SnapshotReader.java b/ai/src/main/java/com/google/genkit/ai/agent/SnapshotReader.java new file mode 100644 index 000000000..3d1f57975 --- /dev/null +++ b/ai/src/main/java/com/google/genkit/ai/agent/SnapshotReader.java @@ -0,0 +1,40 @@ +/* + * 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. + * + * SPDX-License-Identifier: Apache-2.0 + */ + +package com.google.genkit.ai.agent; + +/** + * Read side of a session store. + * + * @param the type of custom session state + */ +public interface SnapshotReader { + + /** + * Retrieves a session snapshot. + * + *

When {@link GetSnapshotOptions#getSnapshotId()} is set the store returns that specific + * snapshot. When {@link GetSnapshotOptions#getSessionId()} is set the store resolves and returns + * the latest leaf snapshot for that session. Returns {@code null} if no matching snapshot exists. + * + * @param opts options specifying which snapshot to retrieve; exactly one of {@code + * snapshotId}/{@code sessionId} must be set + * @return the matching snapshot, or {@code null} if none exists + */ + SessionSnapshot getSnapshot(GetSnapshotOptions opts); +} diff --git a/ai/src/main/java/com/google/genkit/ai/agent/SnapshotStatus.java b/ai/src/main/java/com/google/genkit/ai/agent/SnapshotStatus.java new file mode 100644 index 000000000..7b3a37a4c --- /dev/null +++ b/ai/src/main/java/com/google/genkit/ai/agent/SnapshotStatus.java @@ -0,0 +1,92 @@ +/* + * 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. + * + * SPDX-License-Identifier: Apache-2.0 + */ + +package com.google.genkit.ai.agent; + +import com.fasterxml.jackson.annotation.JsonCreator; +import com.fasterxml.jackson.annotation.JsonValue; + +/** SnapshotStatus represents the status of an agent snapshot. */ +public enum SnapshotStatus { + /** The snapshot is pending. */ + PENDING("pending"), + + /** The snapshot is completed. */ + COMPLETED("completed"), + + /** The snapshot was aborted. */ + ABORTED("aborted"), + + /** The snapshot failed. */ + FAILED("failed"), + + /** The snapshot expired. */ + EXPIRED("expired"); + + private final String value; + + SnapshotStatus(String value) { + this.value = value; + } + + /** + * Returns the string value of the snapshot status. + * + * @return the snapshot status string value + */ + @JsonValue + public String getValue() { + return value; + } + + /** + * Creates a SnapshotStatus from a string value. + * + * @param value the string value + * @return the corresponding SnapshotStatus + * @throws IllegalArgumentException if the value doesn't match any SnapshotStatus + */ + @JsonCreator + public static SnapshotStatus fromValue(String value) { + for (SnapshotStatus status : values()) { + if (status.value.equals(value)) { + return status; + } + } + throw new IllegalArgumentException("Unknown snapshot status: " + value); + } + + /** + * Creates a SnapshotStatus from a string value, treating null or empty strings as COMPLETED. + * + * @param value the string value + * @return the corresponding SnapshotStatus, or COMPLETED if value is null or empty + * @throws IllegalArgumentException if the non-empty value doesn't match any SnapshotStatus + */ + public static SnapshotStatus fromValueOrCompleted(String value) { + if (value == null || value.isEmpty()) { + return COMPLETED; + } + return fromValue(value); + } + + @Override + public String toString() { + return value; + } +} diff --git a/ai/src/main/java/com/google/genkit/ai/agent/SnapshotSubscriber.java b/ai/src/main/java/com/google/genkit/ai/agent/SnapshotSubscriber.java new file mode 100644 index 000000000..bc68cb809 --- /dev/null +++ b/ai/src/main/java/com/google/genkit/ai/agent/SnapshotSubscriber.java @@ -0,0 +1,45 @@ +/* + * 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. + * + * SPDX-License-Identifier: Apache-2.0 + */ + +package com.google.genkit.ai.agent; + +import java.util.function.Consumer; + +/** + * Optional capability: subscribe to a snapshot's status changes. + * + *

Stores that support real-time notifications (e.g. Firestore) implement this interface. + * Required for detach/abort workflows that need to observe when a running snapshot transitions to a + * terminal state. + */ +public interface SnapshotSubscriber { + + /** + * Subscribes to status changes for the snapshot identified by {@code snapshotId}. + * + *

Invokes {@code cb} immediately with the current snapshot and again on every subsequent + * change. The returned {@link AutoCloseable} unsubscribes when closed. + * + * @param snapshotId the snapshot to observe + * @param cb callback invoked with the updated snapshot on each change + * @param options store options + * @return a handle that cancels the subscription when {@link AutoCloseable#close()} is called + */ + AutoCloseable onSnapshotStateChange( + String snapshotId, Consumer> cb, SessionStoreOptions options); +} diff --git a/ai/src/main/java/com/google/genkit/ai/agent/SnapshotWriter.java b/ai/src/main/java/com/google/genkit/ai/agent/SnapshotWriter.java new file mode 100644 index 000000000..0d2737451 --- /dev/null +++ b/ai/src/main/java/com/google/genkit/ai/agent/SnapshotWriter.java @@ -0,0 +1,50 @@ +/* + * 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. + * + * SPDX-License-Identifier: Apache-2.0 + */ + +package com.google.genkit.ai.agent; + +/** + * Write side of a session store. + * + *

The single write primitive is {@link #saveSnapshot}, an atomic read-modify-write. Stores + * enforce the following rules (see individual implementations for details): + * + *

    + *
  • If {@code snapshotId} is {@code null} the store mints a UUID. + *
  • The store preserves {@code sessionId} from the existing snapshot when available. + *
  • An empty {@code sessionId} is rejected with {@code INVALID_ARGUMENT}. + *
  • An empty {@code status} defaults to {@code completed}. + *
  • The {@link SnapshotMutator} may be retried on transient conflicts; it must be pure. + *
+ * + * @param the type of custom session state + */ +public interface SnapshotWriter { + + /** + * Atomically reads the snapshot identified by {@code snapshotId} (or creates a new one), applies + * {@code mutator}, and persists the result. + * + * @param snapshotId the ID of the snapshot to write; if {@code null} the store mints a UUID + * @param mutator the pure read-modify-write function; receives the existing snapshot (or {@code + * null}) and returns the snapshot to persist, or {@code null} to decline + * @param options store options + * @return the snapshot ID that was written, or {@code null} if the mutator declined + */ + String saveSnapshot(String snapshotId, SnapshotMutator mutator, SessionStoreOptions options); +} diff --git a/ai/src/main/java/com/google/genkit/ai/agent/ToolResume.java b/ai/src/main/java/com/google/genkit/ai/agent/ToolResume.java new file mode 100644 index 000000000..890e8cc22 --- /dev/null +++ b/ai/src/main/java/com/google/genkit/ai/agent/ToolResume.java @@ -0,0 +1,109 @@ +/* + * 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. + * + * SPDX-License-Identifier: Apache-2.0 + */ + +package com.google.genkit.ai.agent; + +import com.fasterxml.jackson.annotation.JsonInclude; +import com.fasterxml.jackson.annotation.JsonProperty; +import com.google.genkit.ai.Part; +import java.util.ArrayList; +import java.util.List; + +/** ToolResume provides the parts needed to resume agent execution after a tool interrupt. */ +@JsonInclude(JsonInclude.Include.NON_NULL) +public class ToolResume { + + @JsonProperty("respond") + private List respond; + + @JsonProperty("restart") + private List restart; + + /** Default constructor. */ + public ToolResume() {} + + private ToolResume(Builder builder) { + this.respond = builder.respond; + this.restart = builder.restart; + } + + /** + * Creates a builder for ToolResume. + * + * @return a new builder + */ + public static Builder builder() { + return new Builder(); + } + + /** + * Returns the response parts to resume with. + * + * @return the respond parts + */ + public List getRespond() { + return respond; + } + + /** + * Sets the response parts. + * + * @param respond the respond parts + */ + public void setRespond(List respond) { + this.respond = respond != null ? new ArrayList<>(respond) : null; + } + + /** + * Returns the restart parts. + * + * @return the restart parts + */ + public List getRestart() { + return restart; + } + + /** + * Sets the restart parts. + * + * @param restart the restart parts + */ + public void setRestart(List restart) { + this.restart = restart != null ? new ArrayList<>(restart) : null; + } + + /** Builder for ToolResume. */ + public static class Builder { + private List respond; + private List restart; + + public Builder respond(List respond) { + this.respond = respond; + return this; + } + + public Builder restart(List restart) { + this.restart = restart; + return this; + } + + public ToolResume build() { + return new ToolResume(this); + } + } +} diff --git a/ai/src/main/java/com/google/genkit/ai/agent/TurnBody.java b/ai/src/main/java/com/google/genkit/ai/agent/TurnBody.java new file mode 100644 index 000000000..6aefc9174 --- /dev/null +++ b/ai/src/main/java/com/google/genkit/ai/agent/TurnBody.java @@ -0,0 +1,43 @@ +/* + * 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. + * + * SPDX-License-Identifier: Apache-2.0 + */ + +package com.google.genkit.ai.agent; + +/** + * The body of one agent turn. Implementations run the actual agent logic (e.g. calling the model, + * executing tools) and return the {@link AgentFinishReason} for this turn. + * + *

TurnBody is called by {@link SessionRunner#runTurn} after input validation and message + * appending. On success the runner persists a {@code COMPLETED} snapshot; on exception the runner + * persists a {@code FAILED} snapshot and swallows the exception. + * + * @param the type of the custom session state + */ +@FunctionalInterface +public interface TurnBody { + + /** + * Runs the body of this turn. + * + * @param input the agent input for this turn + * @param turnCtx the per-turn context (snapshot IDs, turn index) + * @return the finish reason for this turn (e.g. {@link AgentFinishReason#STOP}) + * @throws Exception if the turn body fails; the runner will catch and record a FAILED snapshot + */ + AgentFinishReason run(AgentInput input, TurnContext turnCtx) throws Exception; +} diff --git a/ai/src/main/java/com/google/genkit/ai/agent/TurnContext.java b/ai/src/main/java/com/google/genkit/ai/agent/TurnContext.java new file mode 100644 index 000000000..587d32011 --- /dev/null +++ b/ai/src/main/java/com/google/genkit/ai/agent/TurnContext.java @@ -0,0 +1,71 @@ +/* + * 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. + * + * SPDX-License-Identifier: Apache-2.0 + */ + +package com.google.genkit.ai.agent; + +/** + * Per-turn context exposed to {@link TurnBody} implementations. Carries the snapshot IDs reserved + * for this turn and the current turn index. + */ +public final class TurnContext { + + private final String snapshotId; + private final String parentSnapshotId; + private final int turnIndex; + + /** + * Constructs a new TurnContext. + * + * @param snapshotId the snapshot ID reserved for this turn (empty string for client-managed) + * @param parentSnapshotId the snapshot ID from the previous turn, or null if this is the first + * turn + * @param turnIndex the zero-based index of this turn + */ + public TurnContext(String snapshotId, String parentSnapshotId, int turnIndex) { + this.snapshotId = snapshotId; + this.parentSnapshotId = parentSnapshotId; + this.turnIndex = turnIndex; + } + + /** + * Returns the snapshot ID reserved for this turn. + * + * @return the snapshot ID (empty string for client-managed mode) + */ + public String snapshotId() { + return snapshotId; + } + + /** + * Returns the parent snapshot ID (ID of the previous turn's snapshot). + * + * @return the parent snapshot ID, or null if this is the first turn + */ + public String parentSnapshotId() { + return parentSnapshotId; + } + + /** + * Returns the zero-based index of this turn. + * + * @return the turn index + */ + public int turnIndex() { + return turnIndex; + } +} diff --git a/ai/src/main/java/com/google/genkit/ai/agent/TurnEnd.java b/ai/src/main/java/com/google/genkit/ai/agent/TurnEnd.java new file mode 100644 index 000000000..0c772d279 --- /dev/null +++ b/ai/src/main/java/com/google/genkit/ai/agent/TurnEnd.java @@ -0,0 +1,106 @@ +/* + * 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. + * + * SPDX-License-Identifier: Apache-2.0 + */ + +package com.google.genkit.ai.agent; + +import com.fasterxml.jackson.annotation.JsonInclude; +import com.fasterxml.jackson.annotation.JsonProperty; + +/** TurnEnd signals the end of an agent turn in a streaming response. */ +@JsonInclude(JsonInclude.Include.NON_NULL) +public class TurnEnd { + + @JsonProperty("snapshotId") + private String snapshotId; + + @JsonProperty("finishReason") + private AgentFinishReason finishReason; + + /** Default constructor. */ + public TurnEnd() {} + + private TurnEnd(Builder builder) { + this.snapshotId = builder.snapshotId; + this.finishReason = builder.finishReason; + } + + /** + * Creates a builder for TurnEnd. + * + * @return a new builder + */ + public static Builder builder() { + return new Builder(); + } + + /** + * Returns the snapshot ID. + * + * @return the snapshot ID + */ + public String getSnapshotId() { + return snapshotId; + } + + /** + * Sets the snapshot ID. + * + * @param snapshotId the snapshot ID + */ + public void setSnapshotId(String snapshotId) { + this.snapshotId = snapshotId; + } + + /** + * Returns the finish reason. + * + * @return the finish reason + */ + public AgentFinishReason getFinishReason() { + return finishReason; + } + + /** + * Sets the finish reason. + * + * @param finishReason the finish reason + */ + public void setFinishReason(AgentFinishReason finishReason) { + this.finishReason = finishReason; + } + + /** Builder for TurnEnd. */ + public static class Builder { + private String snapshotId; + private AgentFinishReason finishReason; + + public Builder snapshotId(String snapshotId) { + this.snapshotId = snapshotId; + return this; + } + + public Builder finishReason(AgentFinishReason finishReason) { + this.finishReason = finishReason; + return this; + } + + public TurnEnd build() { + return new TurnEnd(this); + } + } +} diff --git a/ai/src/main/java/com/google/genkit/ai/agent/internal/AbortAwareMutator.java b/ai/src/main/java/com/google/genkit/ai/agent/internal/AbortAwareMutator.java new file mode 100644 index 000000000..7da470815 --- /dev/null +++ b/ai/src/main/java/com/google/genkit/ai/agent/internal/AbortAwareMutator.java @@ -0,0 +1,55 @@ +/* + * 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. + * + * SPDX-License-Identifier: Apache-2.0 + */ + +package com.google.genkit.ai.agent.internal; + +import com.google.genkit.ai.agent.SnapshotMutator; +import com.google.genkit.ai.agent.SnapshotStatus; + +/** + * Wraps a {@link SnapshotMutator} so that it returns {@code null} (no-op) when the current row is + * already {@link SnapshotStatus#ABORTED}, preventing a completed or failed write from clobbering a + * concurrent abort. + * + *

Usage: pass {@code AbortAwareMutator.wrap(inner)} to {@code SessionStore.saveSnapshot(...)} + * instead of the bare inner mutator. + */ +public final class AbortAwareMutator { + + private AbortAwareMutator() { + // utility class + } + + /** + * Wraps {@code inner} so that it returns {@code null} when the existing snapshot has status + * {@link SnapshotStatus#ABORTED}; otherwise delegates to {@code inner}. + * + * @param the type of custom session state + * @param inner the underlying mutator to wrap + * @return a new mutator that is abort-aware + */ + public static SnapshotMutator wrap(SnapshotMutator inner) { + return existing -> { + if (existing != null && SnapshotStatus.ABORTED.equals(existing.getStatus())) { + // Existing row is ABORTED — decline the write so we don't clobber it + return null; + } + return inner.apply(existing); + }; + } +} diff --git a/ai/src/main/java/com/google/genkit/ai/agent/internal/AgentActions.java b/ai/src/main/java/com/google/genkit/ai/agent/internal/AgentActions.java new file mode 100644 index 000000000..1fbccf614 --- /dev/null +++ b/ai/src/main/java/com/google/genkit/ai/agent/internal/AgentActions.java @@ -0,0 +1,428 @@ +/* + * 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. + * + * SPDX-License-Identifier: Apache-2.0 + */ + +package com.google.genkit.ai.agent.internal; + +import com.google.genkit.ai.Message; +import com.google.genkit.ai.agent.Agent; +import com.google.genkit.ai.agent.AgentAbortRequest; +import com.google.genkit.ai.agent.AgentAbortResponse; +import com.google.genkit.ai.agent.AgentFinishReason; +import com.google.genkit.ai.agent.AgentFn; +import com.google.genkit.ai.agent.AgentFnContext; +import com.google.genkit.ai.agent.AgentInit; +import com.google.genkit.ai.agent.AgentInput; +import com.google.genkit.ai.agent.AgentOutput; +import com.google.genkit.ai.agent.AgentResult; +import com.google.genkit.ai.agent.AgentSessionContext; +import com.google.genkit.ai.agent.AgentStreamChunk; +import com.google.genkit.ai.agent.Artifact; +import com.google.genkit.ai.agent.ClientTransform; +import com.google.genkit.ai.agent.CustomAgentConfig; +import com.google.genkit.ai.agent.GetSnapshotOptions; +import com.google.genkit.ai.agent.GetSnapshotRequest; +import com.google.genkit.ai.agent.SessionRunner; +import com.google.genkit.ai.agent.SessionSnapshot; +import com.google.genkit.ai.agent.SessionState; +import com.google.genkit.ai.agent.SessionStore; +import com.google.genkit.ai.agent.SessionStoreOptions; +import com.google.genkit.ai.agent.SnapshotStatus; +import com.google.genkit.ai.agent.SnapshotSubscriber; +import com.google.genkit.ai.agent.TurnEnd; +import com.google.genkit.core.Action; +import com.google.genkit.core.ActionDef; +import com.google.genkit.core.ActionType; +import com.google.genkit.core.BidiAction; +import com.google.genkit.core.BidiActionImpl; +import com.google.genkit.core.JsonUtils; +import com.google.genkit.core.Registry; +import com.google.genkit.core.SchemaUtils; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.Optional; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.function.Consumer; + +/** + * Factory for defining custom agents. + * + *

The main entry point is {@link #defineCustomAgent(Registry, CustomAgentConfig, AgentFn)}. + */ +public final class AgentActions { + + private AgentActions() {} + + /** + * Defines a custom agent, registers it (and any companion actions) with the registry, and returns + * the agent instance. + * + * @param the type of custom session state + * @param registry the registry to register the agent and companion actions into + * @param config the agent configuration + * @param fn the agent function implementing per-turn logic + * @return the constructed and registered {@link Agent} + */ + public static Agent defineCustomAgent( + Registry registry, CustomAgentConfig config, AgentFn fn) { + + SessionStore store = config.getStore(); + boolean serverManaged = (store != null); + ClientTransform clientTransform = config.getClientTransform(); + SessionStoreOptions opts = config.getStoreOptions(); + + // Build agent sub-metadata as a plain Map so callers can cast to Map. + // Keys match AgentMetadata's JSON property names so wire-format is identical. + Map agentMetaMap = new HashMap<>(); + agentMetaMap.put("stateManagement", serverManaged ? "server" : "client"); + agentMetaMap.put("abortable", store instanceof SnapshotSubscriber); + + // Generate a JSON schema for the custom state type, reusing the same schema-generation + // mechanism as Tool input/output types (com.google.genkit.core.SchemaUtils, backed by the + // victools SchemaGenerator). Skipped for types with no useful shape to describe: null (not + // specified), Object.class (no fields), and Map (and subtypes) — a dynamic/untyped bag of + // properties that inferSchema can only describe as a bare "type: object" with no properties. + Map stateSchema = inferStateSchema(config.getStateType()); + if (stateSchema != null) { + agentMetaMap.put("stateSchema", stateSchema); + } + + Map metadata = new HashMap<>(); + metadata.put("agent", agentMetaMap); + if (config.getDescription() != null) { + metadata.put("description", config.getDescription()); + } + + // Build the bidi handler + BidiAction.BidiHandler, AgentStreamChunk, AgentInit> handler = + (ctx, init, inputs, sink) -> { + // Step 1: resolve session. API-misuse (wrong init for the state-management mode, + // ownership mismatch, unresolvable snapshot) throws and propagates to the transport. + // Other (non-misuse) pre-turn failures resolve gracefully as a FAILED output (design + // spec §6.4) — guard against a null session before constructing the runner. + SessionResolver.Resolution resolution = + SessionResolver.resolve(store, serverManaged, init, opts); + if (!resolution.isOk()) { + return AgentOutput.builder() + .finishReason(AgentFinishReason.FAILED) + .error(resolution.error()) + .build(); + } + + // Step 2: create session runner, seeded with the id of the snapshot the session was + // resumed from (if any) so the first post-resume turn chains its parentId to it (Gap 4). + SessionRunner runner = + new SessionRunner<>(resolution.session(), store, opts, resolution.sourceSnapshotId()); + + // Step 3: create stream emitter + StreamEmitter emitter = + new StreamEmitter<>(sink != null ? sink : chunk -> {}, JsonUtils.getObjectMapper()); + emitter.attach(runner.session()); + + // Step 4: run the turn loop + Optional next; + while ((next = inputs.next()).isPresent()) { + AgentInput input = next.get(); + + // Detach handling (server-managed only). A client-managed agent has no store to write a + // pending snapshot to, so detach is not applicable there — fall through and process the + // turn normally (graceful no-op for the detach flag; documented in DetachController). + if (input != null && input.getDetach() && serverManaged) { + Consumer detachSink = sink != null ? sink : chunk -> {}; + AtomicBoolean detachAbortSignal = new AtomicBoolean(false); + AgentFnContext detachCtx = + new AgentFnContext(detachSink, detachAbortSignal, ctx, input.getResume()); + final AgentInput detachInput = input; + + // Run the turn in the background with streaming suppressed; finalize the pending + // snapshot on completion. The handler returns DETACHED immediately below. + String pendingId = + DetachController.detach( + runner, + store, + opts, + emitter, + detachAbortSignal, + () -> { + if (detachInput.getMessage() != null) { + runner.addMessages(detachInput.getMessage()); + } + AgentResult res = + AgentSessionContext.call( + runner.session(), () -> callFn(fn, runner, detachCtx)); + if (res != null) { + if (res.getMessage() != null) { + runner.addMessages(res.getMessage()); + } + if (res.getArtifacts() != null) { + runner.addArtifacts(res.getArtifacts().toArray(new Artifact[0])); + } + } + return res != null && res.getFinishReason() != null + ? res.getFinishReason() + : AgentFinishReason.STOP; + }); + + // Return immediately: detached run does not wait for or stream the background work. + return AgentOutput.builder() + .sessionId(runner.sessionId()) + .snapshotId(pendingId) + .finishReason(AgentFinishReason.DETACHED) + .build(); + } + + emitter.beginTurn(); + + Consumer chunkSink = sink != null ? sink : chunk -> {}; + AgentFnContext fnCtx = + new AgentFnContext( + chunkSink, + new AtomicBoolean(false), + ctx, + input != null ? input.getResume() : null); + + runner.runTurn( + input, + (in, turnCtx) -> { + AgentResult res = + AgentSessionContext.call(runner.session(), () -> callFn(fn, runner, fnCtx)); + if (res != null) { + if (res.getMessage() != null) { + runner.addMessages(res.getMessage()); + } + if (res.getArtifacts() != null) { + runner.addArtifacts(res.getArtifacts().toArray(new Artifact[0])); + } + } + return res != null && res.getFinishReason() != null + ? res.getFinishReason() + : AgentFinishReason.STOP; + }); + + // Emit TurnEnd chunk + if (sink != null) { + sink.accept( + AgentStreamChunk.builder() + .turnEnd( + TurnEnd.builder() + .snapshotId(runner.lastSnapshotId()) + .finishReason(runner.lastTurnFinishReason()) + .build()) + .build()); + } + } + + // Step 5: build AgentOutput + AgentOutput.Builder outBuilder = + AgentOutput.builder() + .sessionId(runner.sessionId()) + .finishReason(runner.lastTurnFinishReason()) + .error(runner.lastTurnError()); + + List msgs = runner.getMessages(); + if (!msgs.isEmpty()) { + outBuilder.message(msgs.get(msgs.size() - 1)); + } + + List artifacts = runner.getArtifacts(); + if (!artifacts.isEmpty()) { + outBuilder.artifacts(artifacts); + } + + if (serverManaged) { + String lastSnapshotId = runner.lastSnapshotId(); + if (lastSnapshotId != null && !lastSnapshotId.isEmpty()) { + outBuilder.snapshotId(lastSnapshotId); + } + } else { + SessionState state = runner.getState(); + if (clientTransform != null) { + state = clientTransform.transformState(state); + } + outBuilder.state(state); + } + + return outBuilder.build(); + }; + + // Build the BidiActionImpl + @SuppressWarnings("unchecked") + BidiActionImpl, AgentStreamChunk, AgentInit> impl = + BidiActionImpl., AgentStreamChunk, AgentInit>builder() + .name(config.getName()) + .inputClass(AgentInput.class) + .outputClass((Class>) (Class) AgentOutput.class) + .streamClass(AgentStreamChunk.class) + .initClass((Class>) (Class) AgentInit.class) + .metadata(metadata) + .handler(handler) + .build(); + + // Build companion actions + Action snapshotAction = null; + Action abortAction = null; + + if (store != null) { + snapshotAction = buildSnapshotAction(config.getName(), store, opts); + } + + if (store instanceof SnapshotSubscriber) { + abortAction = buildAbortAction(config.getName(), store, opts); + } + + // Construct the agent + Agent agent = + new Agent<>( + impl, + store, + serverManaged, + snapshotAction, + abortAction, + clientTransform, + opts, + config.getName(), + config.getDescription(), + registry); + + // Register the agent and companion actions + agent.register(registry); + + return agent; + } + + // ── Private helpers ─────────────────────────────────────────────────────────── + + /** + * Generates a JSON schema map for {@code stateType}, reusing {@link SchemaUtils#inferSchema} (the + * same mechanism {@code Tool}/{@code genkit.defineTool} use for input/output schemas). + * + * @param the type of custom session state + * @param stateType the configured state type; may be {@code null} + * @return the generated schema, or {@code null} if {@code stateType} is {@code null}, {@code + * Object.class}, or a {@link Map} (or subtype) — types with no useful shape to describe + */ + private static Map inferStateSchema(Class stateType) { + if (stateType == null || stateType == Object.class || Map.class.isAssignableFrom(stateType)) { + return null; + } + return SchemaUtils.inferSchema(stateType); + } + + /** + * Invokes {@code fn.run(runner, ctx)}, rethrowing any checked exception wrapped in an unchecked + * {@link AgentFnExecutionException} that preserves the original exception's identity (via {@link + * Throwable#getCause()}) and message. + * + *

{@link AgentFn#run} declares {@code throws Exception}, but this is called from inside a + * {@link java.util.function.Supplier} (via {@link AgentSessionContext#call}), which does not + * declare any checked exception. The turn-body lambdas that call this helper already run inside a + * try/catch in {@link SessionRunner#runTurn} (foreground) or the detached-turn background + * runnable (which itself catches {@code Throwable}), so unwrapping is not required there — both + * callers only care that the original exception's message/type is observable, which {@link + * AgentFnExecutionException#getMessage()} and {@link AgentFnExecutionException#getCause()} + * preserve. + * + * @param the type of custom session state + * @param fn the agent function to invoke + * @param runner the session runner for this turn + * @param ctx the per-invocation context for this turn + * @return the agent result returned by {@code fn.run} + */ + private static AgentResult callFn( + AgentFn fn, SessionRunner runner, AgentFnContext ctx) { + try { + return fn.run(runner, ctx); + } catch (RuntimeException e) { + // Unchecked already — propagate as-is so callers see the original type/message unchanged. + throw e; + } catch (Exception e) { + throw new AgentFnExecutionException(e); + } + } + + /** + * Unchecked wrapper used solely to carry a checked exception thrown by {@link AgentFn#run} + * through a {@link java.util.function.Supplier} boundary ({@link AgentSessionContext#call}). + * Callers that catch this should prefer {@link #getCause()} / {@link #getMessage()} (which + * delegates to the cause's message) over the wrapper itself. + */ + static final class AgentFnExecutionException extends RuntimeException { + AgentFnExecutionException(Exception cause) { + super(cause.getMessage(), cause); + } + } + + private static Action buildSnapshotAction( + String agentName, SessionStore store, SessionStoreOptions opts) { + return ActionDef.create( + agentName, + ActionType.AGENT_SNAPSHOT, + null, + null, + GetSnapshotRequest.class, + SessionSnapshot.class, + (ctx, req) -> { + GetSnapshotOptions.Builder optsBuilder = GetSnapshotOptions.builder(); + if (req.getSnapshotId() != null) { + optsBuilder.snapshotId(req.getSnapshotId()); + } + if (req.getSessionId() != null) { + optsBuilder.sessionId(req.getSessionId()); + } + return store.getSnapshot(optsBuilder.build()); + }); + } + + private static Action buildAbortAction( + String agentName, SessionStore store, SessionStoreOptions opts) { + return ActionDef.create( + agentName, + ActionType.AGENT_ABORT, + null, + null, + AgentAbortRequest.class, + AgentAbortResponse.class, + (ctx, req) -> { + String snapshotId = req.getSnapshotId(); + final SnapshotStatus[] resultStatus = {null}; + store.saveSnapshot( + snapshotId, + existing -> { + if (existing == null) { + resultStatus[0] = null; + return null; + } + if (existing.getStatus() != SnapshotStatus.PENDING) { + resultStatus[0] = existing.getStatus(); + return existing; + } + existing.setStatus(SnapshotStatus.ABORTED); + resultStatus[0] = SnapshotStatus.ABORTED; + return existing; + }, + opts); + // Also flip the live in-memory signal for a still-running DETACHED turn registered under + // this snapshot id, mirroring Agent.abort(String) (see PendingAbortRegistry javadoc). + PendingAbortRegistry.signal(snapshotId); + return AgentAbortResponse.builder() + .snapshotId(snapshotId) + .status(resultStatus[0]) + .build(); + }); + } +} diff --git a/ai/src/main/java/com/google/genkit/ai/agent/internal/DetachController.java b/ai/src/main/java/com/google/genkit/ai/agent/internal/DetachController.java new file mode 100644 index 000000000..499ada192 --- /dev/null +++ b/ai/src/main/java/com/google/genkit/ai/agent/internal/DetachController.java @@ -0,0 +1,292 @@ +/* + * 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. + * + * SPDX-License-Identifier: Apache-2.0 + */ + +package com.google.genkit.ai.agent.internal; + +import com.google.genkit.ai.agent.AgentFinishReason; +import com.google.genkit.ai.agent.RuntimeError; +import com.google.genkit.ai.agent.SessionRunner; +import com.google.genkit.ai.agent.SessionSnapshot; +import com.google.genkit.ai.agent.SessionState; +import com.google.genkit.ai.agent.SessionStore; +import com.google.genkit.ai.agent.SessionStoreOptions; +import com.google.genkit.ai.agent.SnapshotMutator; +import com.google.genkit.ai.agent.SnapshotStatus; +import java.time.Instant; +import java.util.UUID; +import java.util.concurrent.Executors; +import java.util.concurrent.ScheduledExecutorService; +import java.util.concurrent.ScheduledFuture; +import java.util.concurrent.ThreadFactory; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicLong; + +/** + * Runs a detached agent turn: writes a {@code PENDING} snapshot, suppresses streaming, starts a + * heartbeat, runs the turn on a background daemon thread, and finalizes the pending snapshot to a + * terminal status when the turn completes. + * + *

Pragmatic approach. The bidi handler returns {@code AgentOutput} + * synchronously after the first detach input. True multi-turn continuation inside the live bidi + * stream would require keeping the stream open while running in the background, which is + * incompatible with the synchronous handler contract. Instead, {@link #detach} reserves the + * snapshot id, kicks off the single turn on a shared daemon executor (streaming suppressed), and + * returns immediately. The client polls {@code getSnapshot} until the snapshot reaches a terminal + * status. This matches the design spec §6.4 and the Go/JS reference behavior for the write side. + * + *

Abort race. Both the heartbeat and the finalize step go through {@link + * AbortAwareMutator}, and the finalize mutator additionally never overwrites an {@code ABORTED} + * row. If the abort companion flips the pending snapshot to {@code ABORTED} while the background + * turn is still running, the heartbeat becomes a no-op (status != PENDING) and the finalize + * declines the write, so the {@code ABORTED} terminal state is preserved. + * + *

Threading. All background work uses shared daemon-thread executors so the JVM + * can exit without blocking and no threads are leaked across invocations. + */ +public final class DetachController { + + private DetachController() {} + + /** + * Heartbeat interval in milliseconds. The heartbeat refreshes {@code heartbeatAt} on the pending + * snapshot so a reader can distinguish a live detached run from a stale (expired) one. + * + *

Package-private and non-final so tests can shrink it for determinism. Production default is + * 30s, matching the Go/JS reference. + */ + static volatile long heartbeatIntervalMillis = 30_000L; + + /** Shared single-thread daemon scheduler for all heartbeats across agents. */ + private static final ScheduledExecutorService HEARTBEAT_SCHEDULER = + Executors.newSingleThreadScheduledExecutor(daemonFactory("genkit-agent-heartbeat")); + + /** Shared daemon executor that runs detached turn bodies. */ + private static final java.util.concurrent.ExecutorService BACKGROUND_EXECUTOR = + Executors.newCachedThreadPool(daemonFactory("genkit-agent-detach")); + + private static ThreadFactory daemonFactory(String prefix) { + AtomicLong counter = new AtomicLong(); + return r -> { + Thread t = new Thread(r, prefix + "-" + counter.incrementAndGet()); + t.setDaemon(true); + return t; + }; + } + + /** + * Overrides the heartbeat interval (test hook). Returns the previous value so tests can restore + * it. Public so tests outside this package can make heartbeat timing deterministic; it is not + * part of the supported API surface. + * + * @param millis the new interval in milliseconds (must be positive) + * @return the previous interval + */ + public static long setHeartbeatIntervalMillisForTest(long millis) { + long prev = heartbeatIntervalMillis; + heartbeatIntervalMillis = millis; + return prev; + } + + /** + * Handles a detached turn for a server-managed agent. + * + *

    + *
  1. Reserves a snapshot id and writes a {@code PENDING} snapshot (cumulative state baseline). + *
  2. Suppresses the stream emitter so the detached run emits no chunks. + *
  3. Schedules a heartbeat that refreshes {@code heartbeatAt} while the snapshot stays + * pending. + *
  4. Runs {@code turnBody} on a shared daemon thread, then finalizes the snapshot + * (abort-aware) to {@code COMPLETED}/{@code FAILED} with the cumulative state and stops the + * heartbeat. + *
+ * + *

Abort signal registration. While the turn is {@code PENDING}, {@code + * abortSignal} (the same {@link AtomicBoolean} handed to the turn's {@code AgentFnContext}) is + * registered in {@link PendingAbortRegistry} under the reserved snapshot id, so an external + * {@code Agent.abort(snapshotId)} call can flip it while the background turn is still running. + * The registration is removed in a {@code finally} block when the turn finishes, regardless of + * outcome. + * + * @param the type of custom session state + * @param runner the session runner (shared with the background thread; {@link + * com.google.genkit.ai.agent.Session} is synchronized) + * @param store the (non-null) server-side store + * @param opts store options + * @param emitter the stream emitter to suppress + * @param abortSignal the abort signal handed to this turn's {@code AgentFnContext}; registered + * under the reserved snapshot id for the lifetime of the background run + * @param turnBody the turn body to execute in the background + * @return the reserved snapshot id for the pending detached run + */ + public static String detach( + SessionRunner runner, + SessionStore store, + SessionStoreOptions opts, + StreamEmitter emitter, + AtomicBoolean abortSignal, + DetachedTurn turnBody) { + + final SessionStoreOptions options = opts != null ? opts : SessionStoreOptions.empty(); + final String snapshotId = UUID.randomUUID().toString(); + final String sessionId = runner.sessionId(); + final String parentId = runner.lastSnapshotId(); + + // Register the abort signal under the reserved snapshot id BEFORE the id becomes externally + // visible (the pending snapshot write below), so there is no window where a caller could + // observe the id via getSnapshot/abort but find no registered signal yet. + PendingAbortRegistry.register(snapshotId, abortSignal); + + // Step 1: write the PENDING snapshot with the current cumulative state as a baseline. + final SessionState baseline = runner.getState(); + final String createdAt = Instant.now().toString(); + SnapshotMutator pendingMutator = + existing -> + SessionSnapshot.builder() + .snapshotId(snapshotId) + .sessionId(sessionId) + .parentId(parentId) + .createdAt(createdAt) + .updatedAt(createdAt) + .heartbeatAt(createdAt) + .status(SnapshotStatus.PENDING) + .state(baseline) + .build(); + store.saveSnapshot(snapshotId, AbortAwareMutator.wrap(pendingMutator), options); + + // Step 2: suppress streaming for the detached run. + emitter.setSuppressed(true); + + // Step 3: start the heartbeat — a no-op once the snapshot leaves PENDING. + final ScheduledFuture[] heartbeatHandle = new ScheduledFuture[1]; + heartbeatHandle[0] = + HEARTBEAT_SCHEDULER.scheduleAtFixedRate( + () -> beat(store, snapshotId, options), + heartbeatIntervalMillis, + heartbeatIntervalMillis, + TimeUnit.MILLISECONDS); + + // Step 4: run the turn body in the background, then finalize. + BACKGROUND_EXECUTOR.execute( + () -> { + SnapshotStatus terminalStatus = SnapshotStatus.COMPLETED; + AgentFinishReason finishReason = AgentFinishReason.STOP; + RuntimeError error = null; + try { + finishReason = turnBody.run(); + if (finishReason == null) { + finishReason = AgentFinishReason.STOP; + } + } catch (Throwable t) { + terminalStatus = SnapshotStatus.FAILED; + finishReason = AgentFinishReason.FAILED; + error = + RuntimeError.builder() + .status("INTERNAL") + .message(t.getMessage() != null ? t.getMessage() : t.getClass().getName()) + .build(); + } finally { + try { + finalizePending( + runner, store, options, snapshotId, terminalStatus, finishReason, error); + } finally { + ScheduledFuture handle = heartbeatHandle[0]; + if (handle != null) { + handle.cancel(false); + } + // Always remove the registration once the turn is no longer running, whatever the + // outcome, so entries never leak and a later abort() call for a reused/unknown id + // cannot spuriously flip a stale signal. + PendingAbortRegistry.unregister(snapshotId); + } + } + }); + + // Step 5: caller returns AgentOutput{finishReason=DETACHED, snapshotId} immediately. + return snapshotId; + } + + /** Refreshes {@code heartbeatAt} iff the snapshot is still pending; otherwise a no-op. */ + private static void beat(SessionStore store, String snapshotId, SessionStoreOptions opts) { + SnapshotMutator mutator = + existing -> { + if (existing == null || existing.getStatus() != SnapshotStatus.PENDING) { + return null; // decline: nothing to refresh + } + existing.setHeartbeatAt(Instant.now().toString()); + return existing; + }; + try { + store.saveSnapshot(snapshotId, AbortAwareMutator.wrap(mutator), opts); + } catch (RuntimeException ignored) { + // A heartbeat failure must never crash the scheduler thread. + } + } + + /** + * Finalizes the pending snapshot to a terminal status with the cumulative session state. Never + * overwrites an {@code ABORTED} row, and declines if the snapshot has gone missing. + */ + private static void finalizePending( + SessionRunner runner, + SessionStore store, + SessionStoreOptions opts, + String snapshotId, + SnapshotStatus terminalStatus, + AgentFinishReason finishReason, + RuntimeError error) { + + final SessionState finalState = runner.getState(); + final String now = Instant.now().toString(); + SnapshotMutator mutator = + existing -> { + if (existing == null) { + return null; // snapshot gone — decline + } + if (existing.getStatus() == SnapshotStatus.ABORTED) { + return existing; // abort won — preserve it + } + existing.setStatus(terminalStatus); + existing.setFinishReason(finishReason); + existing.setError(error); + existing.setState(finalState); + existing.setUpdatedAt(now); + existing.setHeartbeatAt(null); // terminal: no more heartbeats + return existing; + }; + store.saveSnapshot(snapshotId, AbortAwareMutator.wrap(mutator), opts); + } + + /** + * The body of a detached turn. Runs the agent function and applies its result to the session, + * returning the finish reason. Thrown exceptions cause the pending snapshot to finalize as {@code + * FAILED}. + * + * @param the type of custom session state + */ + @FunctionalInterface + public interface DetachedTurn { + /** + * Runs the detached turn body. + * + * @return the finish reason + * @throws Exception if the turn fails + */ + AgentFinishReason run() throws Exception; + } +} diff --git a/ai/src/main/java/com/google/genkit/ai/agent/internal/InProcessTransport.java b/ai/src/main/java/com/google/genkit/ai/agent/internal/InProcessTransport.java new file mode 100644 index 000000000..04b91e963 --- /dev/null +++ b/ai/src/main/java/com/google/genkit/ai/agent/internal/InProcessTransport.java @@ -0,0 +1,122 @@ +/* + * 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. + * + * SPDX-License-Identifier: Apache-2.0 + */ + +package com.google.genkit.ai.agent.internal; + +import com.fasterxml.jackson.databind.JsonNode; +import com.google.genkit.ai.agent.Agent; +import com.google.genkit.ai.agent.AgentInit; +import com.google.genkit.ai.agent.AgentInput; +import com.google.genkit.ai.agent.AgentOutput; +import com.google.genkit.ai.agent.AgentStreamChunk; +import com.google.genkit.ai.agent.AgentTransport; +import com.google.genkit.ai.agent.GetSnapshotRequest; +import com.google.genkit.ai.agent.SessionSnapshot; +import com.google.genkit.ai.agent.SnapshotStatus; +import com.google.genkit.core.ActionContext; +import com.google.genkit.core.BufferedInputSource; +import com.google.genkit.core.GenkitException; +import com.google.genkit.core.JsonUtils; +import java.util.function.Consumer; + +/** + * In-process {@link AgentTransport} that drives a locally-defined {@link Agent} via its bidi + * action. + * + *

Each {@link #runTurn} call feeds exactly one input (offer + end) into a fresh {@link + * BufferedInputSource} and invokes {@code agent.runBidiJson}, deserializing the returned node into + * an {@link AgentOutput}. Snapshot retrieval and abort delegate to the agent's typed facades. + * + * @param the type of custom session state + */ +public final class InProcessTransport implements AgentTransport { + + private final Agent agent; + private final ActionContext ctx; + + /** + * Constructs an in-process transport. + * + * @param agent the agent to drive (must not be null) + * @param ctx the action context used for each turn invocation (must not be null) + */ + public InProcessTransport(Agent agent, ActionContext ctx) { + if (agent == null) { + throw new IllegalArgumentException("agent must not be null"); + } + if (ctx == null) { + throw new IllegalArgumentException("ctx must not be null"); + } + this.agent = agent; + this.ctx = ctx; + } + + @Override + @SuppressWarnings("unchecked") + public AgentOutput runTurn( + AgentInput input, AgentInit init, Consumer onChunk) { + BufferedInputSource inputs = new BufferedInputSource<>(); + inputs.offer(JsonUtils.toJsonNode(input != null ? input : new AgentInput())); + inputs.end(); + + JsonNode initNode = JsonUtils.toJsonNode(init != null ? init : new AgentInit()); + + Consumer sink = + chunkJson -> { + if (onChunk == null) { + return; + } + try { + onChunk.accept( + JsonUtils.getObjectMapper().treeToValue(chunkJson, AgentStreamChunk.class)); + } catch (Exception e) { + throw new GenkitException("Failed to deserialize agent stream chunk", e); + } + }; + + JsonNode outNode; + try { + outNode = agent.runBidiJson(ctx, initNode, inputs, sink); + } catch (GenkitException e) { + throw e; + } catch (Exception e) { + throw new GenkitException("Agent turn failed", e); + } + + try { + return (AgentOutput) JsonUtils.getObjectMapper().treeToValue(outNode, AgentOutput.class); + } catch (Exception e) { + throw new GenkitException("Failed to deserialize agent output", e); + } + } + + @Override + public SessionSnapshot getSnapshot(GetSnapshotRequest req) { + return agent.getSnapshotData(req); + } + + @Override + public SnapshotStatus abort(String snapshotId) { + return agent.abort(snapshotId); + } + + @Override + public boolean serverManaged() { + return agent.serverManaged(); + } +} diff --git a/ai/src/main/java/com/google/genkit/ai/agent/internal/LeafSelection.java b/ai/src/main/java/com/google/genkit/ai/agent/internal/LeafSelection.java new file mode 100644 index 000000000..1d8bc4031 --- /dev/null +++ b/ai/src/main/java/com/google/genkit/ai/agent/internal/LeafSelection.java @@ -0,0 +1,148 @@ +/* + * 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. + * + * SPDX-License-Identifier: Apache-2.0 + */ + +package com.google.genkit.ai.agent.internal; + +import com.google.genkit.ai.agent.SessionSnapshot; +import com.google.genkit.core.GenkitException; +import java.time.Instant; +import java.util.List; +import java.util.Set; +import java.util.stream.Collectors; + +/** + * Utility for selecting the leaf snapshot from a collection of snapshots in a session's + * parent-chain. + * + *

A leaf is a snapshot whose {@code snapshotId} is not referenced as a {@code parentId} + * by any other snapshot in the list. + */ +public final class LeafSelection { + + private LeafSelection() {} + + /** + * Selects the leaf snapshot from the given list. + * + *

Selection rules: + * + *

    + *
  1. Empty list — returns {@code null}. An empty input carries no information + * about a cycle, so returning {@code null} (no snapshot found) is the correct behaviour. + *
  2. Exactly one leaf — returns it directly. + *
  3. Zero leaves in a non-empty list — the chain is corrupt (cycle). Throws + * {@link GenkitException} with error code {@code FAILED_PRECONDITION}. + *
  4. More than one leaf — if {@code rejectBranching} is {@code true}, throws + * {@link GenkitException} with error code {@code FAILED_PRECONDITION}. Otherwise the most + * recent leaf is returned (by {@code createdAt} parsed as {@link Instant}; {@code null} + * {@code createdAt} is treated as the epoch / earliest). Ties are broken by {@code + * snapshotId} in natural (lexicographic) ascending order, taking the greater ID. + *
+ * + * @param the type of custom session state + * @param snapshots the full list of snapshots for a session (may be empty) + * @param rejectBranching if {@code true}, more than one leaf is treated as an error + * @return the selected leaf snapshot, or {@code null} if the input list is empty + * @throws GenkitException with {@code FAILED_PRECONDITION} when the chain contains a cycle or + * (with {@code rejectBranching=true}) multiple leaves + */ + public static SessionSnapshot selectLeaf( + List> snapshots, boolean rejectBranching) { + + if (snapshots == null || snapshots.isEmpty()) { + return null; + } + + // Collect all IDs that appear as a parentId — these are non-leaf nodes. + Set referencedAsParent = + snapshots.stream() + .map(SessionSnapshot::getParentId) + .filter(pid -> pid != null) + .collect(Collectors.toSet()); + + // Leaves are snapshots whose own snapshotId is not in the parent set. + List> leaves = + snapshots.stream() + .filter(s -> !referencedAsParent.contains(s.getSnapshotId())) + .collect(Collectors.toList()); + + if (leaves.isEmpty()) { + throw GenkitException.builder() + .message( + "Session snapshot chain contains a cycle or all snapshots reference a parent:" + + " no leaf found.") + .errorCode("FAILED_PRECONDITION") + .build(); + } + + if (leaves.size() == 1) { + return leaves.get(0); + } + + // Multiple leaves. + if (rejectBranching) { + throw GenkitException.builder() + .message( + "Session snapshot chain has " + + leaves.size() + + " leaves (branching detected). Use rejectBranching=false to select the most" + + " recent leaf automatically.") + .errorCode("FAILED_PRECONDITION") + .build(); + } + + // Pick the most recent leaf; null createdAt treated as epoch (earliest). + return leaves.stream() + .max(LeafSelection::compareByCreatedAtThenId) + .orElseThrow( + () -> + GenkitException.builder() + .message("Unexpected empty leaves stream.") + .errorCode("FAILED_PRECONDITION") + .build()); + } + + /** + * Comparator: earlier createdAt → smaller; null createdAt → epoch (smallest). Tie-break: larger + * snapshotId wins (natural string order ascending, max). + */ + private static int compareByCreatedAtThenId(SessionSnapshot a, SessionSnapshot b) { + Instant ia = parseInstant(a.getCreatedAt()); + Instant ib = parseInstant(b.getCreatedAt()); + int cmp = ia.compareTo(ib); + if (cmp != 0) { + return cmp; + } + // Tie-break: lexicographically larger snapshotId wins. + String idA = a.getSnapshotId() != null ? a.getSnapshotId() : ""; + String idB = b.getSnapshotId() != null ? b.getSnapshotId() : ""; + return idA.compareTo(idB); + } + + /** Parses an RFC-3339 timestamp string; returns {@link Instant#EPOCH} for null or unparseable. */ + private static Instant parseInstant(String rfc3339) { + if (rfc3339 == null || rfc3339.isEmpty()) { + return Instant.EPOCH; + } + try { + return Instant.parse(rfc3339); + } catch (Exception e) { + return Instant.EPOCH; + } + } +} diff --git a/ai/src/main/java/com/google/genkit/ai/agent/internal/PendingAbortRegistry.java b/ai/src/main/java/com/google/genkit/ai/agent/internal/PendingAbortRegistry.java new file mode 100644 index 000000000..6be6fce6f --- /dev/null +++ b/ai/src/main/java/com/google/genkit/ai/agent/internal/PendingAbortRegistry.java @@ -0,0 +1,101 @@ +/* + * 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. + * + * SPDX-License-Identifier: Apache-2.0 + */ + +package com.google.genkit.ai.agent.internal; + +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.atomic.AtomicBoolean; + +/** + * Registry mapping a still-{@code PENDING} detached turn's snapshot id to the live {@link + * AtomicBoolean} abort signal handed to that turn's {@code AgentFnContext}. + * + *

This is the only case where an external caller can know a turn's snapshot id WHILE the turn is + * still running: {@code DetachController.detach(...)} writes the pending snapshot and returns its + * id immediately, before the background work finishes. Foreground turns have no resolvable id until + * after they return, so there is no reachable window in which to abort them externally — the + * registry intentionally only covers detached turns. + * + *

{@link DetachController#detach} registers the signal for a snapshot id when the background + * turn starts, and removes it in a {@code finally} block when the turn finishes (success, failure, + * or abort) so entries never leak. {@code Agent.abort(String)} looks up the signal and flips it (in + * addition to the existing snapshot-store status mutation, which remains the source of truth for + * anything that reads the snapshot after the fact, e.g. a poller that only checks status). + */ +public final class PendingAbortRegistry { + + private static final ConcurrentHashMap SIGNALS = new ConcurrentHashMap<>(); + + private PendingAbortRegistry() {} + + /** + * Registers {@code signal} as the abort flag for the pending detached turn identified by {@code + * snapshotId}. + * + * @param snapshotId the pending snapshot id + * @param signal the abort signal handed to that turn's {@code AgentFnContext} + */ + public static void register(String snapshotId, AtomicBoolean signal) { + if (snapshotId == null || signal == null) { + return; + } + SIGNALS.put(snapshotId, signal); + } + + /** + * Removes the registration for {@code snapshotId}, if present. Safe to call more than once. + * + * @param snapshotId the pending snapshot id + */ + public static void unregister(String snapshotId) { + if (snapshotId == null) { + return; + } + SIGNALS.remove(snapshotId); + } + + /** + * Flips the abort signal registered for {@code snapshotId} to {@code true}, if one is currently + * registered (i.e. the turn is a still-running detached turn). + * + * @param snapshotId the snapshot id to signal + * @return {@code true} if a signal was found and flipped; {@code false} if no turn is currently + * registered under that id (already finished, foreground, or unknown) + */ + public static boolean signal(String snapshotId) { + if (snapshotId == null) { + return false; + } + AtomicBoolean flag = SIGNALS.get(snapshotId); + if (flag == null) { + return false; + } + flag.set(true); + return true; + } + + /** + * Returns the number of currently registered (still-pending, running) detached turns. Test-only + * visibility hook. + * + * @return the number of registered signals + */ + static int size() { + return SIGNALS.size(); + } +} diff --git a/ai/src/main/java/com/google/genkit/ai/agent/internal/PointerDoc.java b/ai/src/main/java/com/google/genkit/ai/agent/internal/PointerDoc.java new file mode 100644 index 000000000..8acfe3d51 --- /dev/null +++ b/ai/src/main/java/com/google/genkit/ai/agent/internal/PointerDoc.java @@ -0,0 +1,115 @@ +/* + * 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. + * + * SPDX-License-Identifier: Apache-2.0 + */ + +package com.google.genkit.ai.agent.internal; + +import com.fasterxml.jackson.annotation.JsonInclude; +import com.fasterxml.jackson.annotation.JsonProperty; + +/** + * Tiny POJO representing the contents of a per-session pointer file. + * + *

The pointer file lives at {@code

//.pointers/.json} and contains a + * reference to the current latest-leaf snapshot for that session, allowing fast lookup without + * scanning all snapshot files. + */ +@JsonInclude(JsonInclude.Include.NON_NULL) +public class PointerDoc { + + /** The snapshot ID of the current latest-leaf snapshot for the session. */ + @JsonProperty("currentSnapshotId") + private String currentSnapshotId; + + /** The {@code createdAt} timestamp of the current latest-leaf snapshot (RFC-3339). */ + @JsonProperty("currentCreatedAt") + private String currentCreatedAt; + + /** The timestamp when this pointer was last updated (RFC-3339). */ + @JsonProperty("updatedAt") + private String updatedAt; + + /** Default constructor (required for Jackson deserialization). */ + public PointerDoc() {} + + /** + * Creates a new {@code PointerDoc}. + * + * @param currentSnapshotId the snapshot ID + * @param currentCreatedAt the snapshot's createdAt timestamp + * @param updatedAt when this pointer was written + */ + public PointerDoc(String currentSnapshotId, String currentCreatedAt, String updatedAt) { + this.currentSnapshotId = currentSnapshotId; + this.currentCreatedAt = currentCreatedAt; + this.updatedAt = updatedAt; + } + + /** + * Returns the current snapshot ID. + * + * @return the snapshot ID + */ + public String getCurrentSnapshotId() { + return currentSnapshotId; + } + + /** + * Sets the current snapshot ID. + * + * @param currentSnapshotId the snapshot ID + */ + public void setCurrentSnapshotId(String currentSnapshotId) { + this.currentSnapshotId = currentSnapshotId; + } + + /** + * Returns the current snapshot's createdAt timestamp. + * + * @return the createdAt timestamp (RFC-3339) + */ + public String getCurrentCreatedAt() { + return currentCreatedAt; + } + + /** + * Sets the current snapshot's createdAt timestamp. + * + * @param currentCreatedAt the createdAt timestamp (RFC-3339) + */ + public void setCurrentCreatedAt(String currentCreatedAt) { + this.currentCreatedAt = currentCreatedAt; + } + + /** + * Returns the timestamp when this pointer was last updated. + * + * @return the updatedAt timestamp (RFC-3339) + */ + public String getUpdatedAt() { + return updatedAt; + } + + /** + * Sets the timestamp when this pointer was last updated. + * + * @param updatedAt the updatedAt timestamp (RFC-3339) + */ + public void setUpdatedAt(String updatedAt) { + this.updatedAt = updatedAt; + } +} diff --git a/ai/src/main/java/com/google/genkit/ai/agent/internal/SessionResolver.java b/ai/src/main/java/com/google/genkit/ai/agent/internal/SessionResolver.java new file mode 100644 index 000000000..ae1c13fe0 --- /dev/null +++ b/ai/src/main/java/com/google/genkit/ai/agent/internal/SessionResolver.java @@ -0,0 +1,331 @@ +/* + * 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. + * + * SPDX-License-Identifier: Apache-2.0 + */ + +package com.google.genkit.ai.agent.internal; + +import com.google.genkit.ai.agent.AgentInit; +import com.google.genkit.ai.agent.GetSnapshotOptions; +import com.google.genkit.ai.agent.RuntimeError; +import com.google.genkit.ai.agent.Session; +import com.google.genkit.ai.agent.SessionSnapshot; +import com.google.genkit.ai.agent.SessionState; +import com.google.genkit.ai.agent.SessionStore; +import com.google.genkit.ai.agent.SessionStoreOptions; +import com.google.genkit.ai.agent.SnapshotStatus; +import com.google.genkit.core.GenkitException; + +/** + * Resolves a {@link Session} from an {@link AgentInit} at the start of an agent invocation. + * + *

Enforces the server-vs-client state-management rules: + * + *

    + *
  • API misuse (wrong init for the state-management mode, ownership mismatch) throws {@link + * GenkitException} with {@code FAILED_PRECONDITION} or {@code INVALID_ARGUMENT}. + *
  • Other pre-turn failures are returned as a {@link Resolution#failure} so the agent can emit + * an {@code AgentOutput} with {@code finishReason=FAILED} rather than surfacing as an + * unhandled exception. + *
+ * + *

Pending-leaf decision

+ * + *

When resolving by sessionId and the latest leaf snapshot has a non-COMPLETED status (e.g. + * PENDING, FAILED, ABORTED), this implementation throws {@link GenkitException} with {@code + * FAILED_PRECONDITION}. This matches Go's {@code resumeSessionFrom} behavior and is simpler than + * the JS approach of walking back to the last completed snapshot. The error is conformance-tested; + * callers should clear or complete any in-progress snapshot before starting a new turn. + * + *

Unknown-snapshot decision

+ * + *

When a specific {@code snapshotId} is provided but not found in the store, this implementation + * throws {@link GenkitException} with {@code INVALID_ARGUMENT}. An unknown snapshotId is treated as + * API misuse (the client provided a bad reference), not as a recoverable pre-turn error. This + * matches the design spec §6.4 note: "unknown snapshot → throw/404". + */ +public final class SessionResolver { + + private SessionResolver() {} + + // ── Resolution ──────────────────────────────────────────────────────────────── + + /** + * Result of session resolution: either a ready {@link Session}, or a graceful pre-turn failure. + * + *

A failure resolution lets the agent generator (Task 4.5) emit an {@code AgentOutput} with + * {@code finishReason=FAILED} instead of throwing. + * + * @param the type of custom session state + */ + public static final class Resolution { + + private final Session session; + private final RuntimeError error; + private final String sourceSnapshotId; + + private Resolution(Session session, RuntimeError error, String sourceSnapshotId) { + this.session = session; + this.error = error; + this.sourceSnapshotId = sourceSnapshotId; + } + + /** + * Creates a successful resolution wrapping the given session. + * + * @param the type of custom session state + * @param session the resolved session (must not be null) + * @return a successful Resolution + */ + public static Resolution ok(Session session) { + return new Resolution<>(session, null, null); + } + + /** + * Creates a successful resolution wrapping the given session and the id of the snapshot it was + * resolved (resumed) from. The source snapshot id seeds the runner's {@code lastSnapshotId} so + * the first post-resume turn chains its {@code parentId} to it (see {@code SessionRunner}). + * + * @param the type of custom session state + * @param session the resolved session (must not be null) + * @param sourceSnapshotId the id of the snapshot this session was resumed from; may be {@code + * null} for a fresh session + * @return a successful Resolution carrying the source snapshot id + */ + public static Resolution ok(Session session, String sourceSnapshotId) { + return new Resolution<>(session, null, sourceSnapshotId); + } + + /** + * Creates a failure resolution carrying a {@link RuntimeError}. + * + *

The agent generator should translate this into an {@code AgentOutput} with {@code + * finishReason=FAILED} rather than throwing. + * + * @param the type of custom session state + * @param error the runtime error describing the failure + * @return a failure Resolution + */ + public static Resolution failure(RuntimeError error) { + return new Resolution<>(null, error, null); + } + + /** + * Returns {@code true} if this is a successful resolution with a ready session. + * + * @return true if ok + */ + public boolean isOk() { + return session != null; + } + + /** + * Returns the resolved session, or {@code null} if this is a failure resolution. + * + * @return the session, or null + */ + public Session session() { + return session; + } + + /** + * Returns the runtime error, or {@code null} if this is a successful resolution. + * + * @return the error, or null + */ + public RuntimeError error() { + return error; + } + + /** + * Returns the id of the snapshot this session was resolved (resumed) from, or {@code null} for + * a fresh session. Used to seed the runner's {@code lastSnapshotId} so the first post-resume + * turn chains its {@code parentId} to the resumed snapshot. + * + * @return the source snapshot id, or {@code null} + */ + public String sourceSnapshotId() { + return sourceSnapshotId; + } + } + + // ── resolve ─────────────────────────────────────────────────────────────────── + + /** + * Resolves a {@link Session} from the given {@link AgentInit}. + * + *

Rules (see class-level Javadoc for details): + * + *

    + *
  1. State-management mismatch → throws {@link GenkitException}. + *
  2. Client-managed: hydrate from {@code init.state} or mint a fresh session. + *
  3. Server-managed: resolve via snapshotId, sessionId, or mint fresh. + *
+ * + * @param the type of custom session state + * @param store the agent's session store, or {@code null} for client-managed agents + * @param serverManaged {@code true} iff the agent has a store (server-managed) + * @param init the agent init (may be {@code null} — treated as all fields null) + * @param opts options forwarded to store operations (may be {@code null}) + * @return a {@link Resolution} wrapping either the resolved session or a graceful failure + * @throws GenkitException with {@code FAILED_PRECONDITION} or {@code INVALID_ARGUMENT} for API + * misuse (wrong init for the state-management mode, or ownership mismatch) + */ + public static Resolution resolve( + SessionStore store, boolean serverManaged, AgentInit init, SessionStoreOptions opts) { + + // Null-safe field extraction — treat null init as all-fields-null. + String snapshotId = init != null ? init.getSnapshotId() : null; + String sessionId = init != null ? init.getSessionId() : null; + SessionState state = init != null ? init.getState() : null; + + // ── Rule 1: state-management mismatch → THROW ──────────────────────────── + + if (serverManaged && state != null) { + throw GenkitException.builder() + .message("Cannot send 'state' to a server-managed agent") + .errorCode("FAILED_PRECONDITION") + .build(); + } + + if (!serverManaged && (sessionId != null || snapshotId != null)) { + String which = sessionId != null ? "sessionId" : "snapshotId"; + throw GenkitException.builder() + .message("Cannot use '" + which + "'/'snapshotId' with a client-managed agent") + .errorCode("FAILED_PRECONDITION") + .build(); + } + + // ── Rule 2: client-managed ──────────────────────────────────────────────── + + if (!serverManaged) { + if (state != null) { + return Resolution.ok(new Session<>(state)); + } + return Resolution.ok(new Session<>(new SessionState<>())); + } + + // ── Rule 3: server-managed ──────────────────────────────────────────────── + + if (snapshotId != null) { + return resolveBySnapshotId(store, snapshotId, sessionId, opts); + } + + if (sessionId != null) { + return resolveBySessionId(store, sessionId, opts); + } + + // No snapshotId, no sessionId → fresh server session (Session mints a new UUID). + return Resolution.ok(new Session<>(new SessionState<>())); + } + + // ── private helpers ─────────────────────────────────────────────────────────── + + /** + * Resolves a session from a specific snapshotId. Throws for unknown, non-COMPLETED, or + * ownership-mismatched snapshots. + */ + private static Resolution resolveBySnapshotId( + SessionStore store, String snapshotId, String callerSessionId, SessionStoreOptions opts) { + + SessionSnapshot snap = + store.getSnapshot(GetSnapshotOptions.builder().snapshotId(snapshotId).build()); + + if (snap == null) { + throw GenkitException.builder() + .message("snapshot not found: " + snapshotId) + .errorCode("INVALID_ARGUMENT") + .build(); + } + + if (snap.getStatus() != SnapshotStatus.COMPLETED) { + throw GenkitException.builder() + .message("snapshot is not resumable: " + snap.getStatus()) + .errorCode("INVALID_ARGUMENT") + .build(); + } + + // Ownership check: if the caller also provided a sessionId it must match. + if (callerSessionId != null) { + String snapSessionId = snap.getSessionId(); + if (snapSessionId == null && snap.getState() != null) { + snapSessionId = snap.getState().getSessionId(); + } + if (!callerSessionId.equals(snapSessionId)) { + throw GenkitException.builder() + .message("snapshot does not belong to session: " + callerSessionId) + .errorCode("INVALID_ARGUMENT") + .build(); + } + } + + return Resolution.ok(new Session<>(hydratedState(snap)), snap.getSnapshotId()); + } + + /** + * Resolves a session from a sessionId. Fresh if unknown; hydrated if COMPLETED leaf; throws for + * non-COMPLETED leaf. + */ + private static Resolution resolveBySessionId( + SessionStore store, String sessionId, SessionStoreOptions opts) { + + SessionSnapshot snap = + store.getSnapshot(GetSnapshotOptions.builder().sessionId(sessionId).build()); + + if (snap == null) { + // Unknown sessionId → fresh session bound to that id. + SessionState fresh = new SessionState<>(); + fresh.setSessionId(sessionId); + return Resolution.ok(new Session<>(fresh)); + } + + if (snap.getStatus() == SnapshotStatus.COMPLETED) { + return Resolution.ok(new Session<>(hydratedState(snap)), snap.getSnapshotId()); + } + + // Latest leaf is pending/failed/aborted/expired → reject resumption. + // Decision: throw FAILED_PRECONDITION (matches Go's resumeSessionFrom). + // The JS runtime walks back to the last completed leaf, but throwing is simpler and + // conformance-correct. The caller must clear or complete in-progress snapshots first. + throw GenkitException.builder() + .message("cannot resume: latest snapshot status=" + snap.getStatus()) + .errorCode("FAILED_PRECONDITION") + .build(); + } + + /** + * Extracts the state from a snapshot, ensuring the returned state carries the snapshot's + * sessionId. Falls back to a fresh {@link SessionState} with the snapshot's sessionId if the + * snapshot's state is null. + */ + private static SessionState hydratedState(SessionSnapshot snap) { + SessionState state = snap.getState(); + if (state != null) { + // Ensure sessionId is set on the state so subsequent snapshots chain correctly. + if (state.getSessionId() == null || state.getSessionId().isEmpty()) { + String sid = snap.getSessionId(); + if (sid != null && !sid.isEmpty()) { + state.setSessionId(sid); + } + } + return state; + } + // Snapshot has no state — create a fresh state carrying the snapshot's sessionId. + SessionState fresh = new SessionState<>(); + fresh.setSessionId(snap.getSessionId()); + return fresh; + } +} diff --git a/ai/src/main/java/com/google/genkit/ai/agent/internal/SnapshotSharding.java b/ai/src/main/java/com/google/genkit/ai/agent/internal/SnapshotSharding.java new file mode 100644 index 000000000..9a65985e3 --- /dev/null +++ b/ai/src/main/java/com/google/genkit/ai/agent/internal/SnapshotSharding.java @@ -0,0 +1,188 @@ +/* + * 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. + * + * SPDX-License-Identifier: Apache-2.0 + */ + +package com.google.genkit.ai.agent.internal; + +import com.fasterxml.jackson.databind.JsonNode; +import com.fasterxml.jackson.databind.ObjectMapper; +import com.fasterxml.jackson.databind.node.NullNode; +import com.google.genkit.core.GenkitException; +import com.google.genkit.core.JsonUtils; +import com.google.genkit.core.jsonpatch.JsonPatch; +import java.nio.charset.StandardCharsets; +import java.util.ArrayList; +import java.util.List; + +/** + * Pure, backend-agnostic sharding / checkpoint / RFC-6902 reconstruction helpers shared by the + * Firestore, DynamoDB, and Cosmos DB session stores. + * + *

Session stores that use the sharded checkpoint + diff + pointer layout persist a session's + * state as a periodic full "checkpoint" (its JSON split into byte-sized shards) plus a chain of + * RFC-6902 "diff" patches from the nearest checkpoint. These helpers implement the pure logic of + * that scheme so every backend behaves identically: + * + *

    + *
  • {@link #shouldCheckpoint} — decides checkpoint vs diff for a new snapshot. + *
  • {@link #shardString} / {@link #reassembleShards} — byte-exact splitting/rejoining of the + * checkpoint JSON. + *
  • {@link #reconstructState} — replays the checkpoint + ordered diffs back into a state node. + *
  • {@link #validateId} — validates snapshot/session identifiers. + *
+ */ +public final class SnapshotSharding { + + private static final ObjectMapper MAPPER = JsonUtils.getObjectMapper(); + + private SnapshotSharding() {} + + /** + * Decides whether to write a full checkpoint instead of a diff. + * + *

A checkpoint is written when: + * + *

    + *
  • there is no usable parent (session root, or parent orphaned/missing); or + *
  • the depth from the nearest checkpoint reaches {@code checkpointInterval}; or + *
  • the diff against the parent would exceed {@code shardSize} bytes. + *
+ * + * @param parentExists whether a usable parent snapshot exists + * @param depthFromCheckpoint the number of diffs (inclusive) that would separate this snapshot + * from its nearest checkpoint + * @param checkpointInterval the configured checkpoint interval + * @param diffSizeBytes the serialized size of the candidate diff in bytes + * @param shardSize the configured shard size in bytes + * @return {@code true} if a full checkpoint should be written + */ + public static boolean shouldCheckpoint( + boolean parentExists, + int depthFromCheckpoint, + int checkpointInterval, + int diffSizeBytes, + int shardSize) { + if (!parentExists) { + return true; + } + if (depthFromCheckpoint >= checkpointInterval) { + return true; + } + return diffSizeBytes > shardSize; + } + + /** + * Splits a UTF-8 string into shards of at most {@code shardSize} bytes each. + * + *

Sharding is byte-exact: shards are concatenated by raw bytes (not chars) so multi-byte UTF-8 + * sequences split across a boundary reassemble correctly. + * + * @param value the string to shard + * @param shardSize the maximum shard size in bytes (must be {@code >= 1}) + * @return the ordered shards; each shard holds at most {@code shardSize} raw UTF-8 bytes encoded + * as an ISO-8859-1 string (1 char per byte) for byte-exact reassembly by {@link + * #reassembleShards} + */ + public static List shardString(String value, int shardSize) { + if (shardSize < 1) { + throw new IllegalArgumentException("shardSize must be >= 1"); + } + byte[] bytes = value.getBytes(StandardCharsets.UTF_8); + List shards = new ArrayList<>(); + if (bytes.length == 0) { + shards.add(""); + return shards; + } + for (int offset = 0; offset < bytes.length; offset += shardSize) { + int len = Math.min(shardSize, bytes.length - offset); + // Encode raw bytes as ISO-8859-1 so each byte maps 1:1 to a char, allowing byte-exact + // reassembly even when a multi-byte UTF-8 sequence is split across a shard boundary. + shards.add(new String(bytes, offset, len, StandardCharsets.ISO_8859_1)); + } + return shards; + } + + /** + * Reassembles shards produced by {@link #shardString} back into the original string. + * + * @param shards the ordered shard list + * @return the reassembled original string + */ + public static String reassembleShards(List shards) { + int total = 0; + byte[][] parts = new byte[shards.size()][]; + for (int i = 0; i < shards.size(); i++) { + parts[i] = shards.get(i).getBytes(StandardCharsets.ISO_8859_1); + total += parts[i].length; + } + byte[] all = new byte[total]; + int pos = 0; + for (byte[] part : parts) { + System.arraycopy(part, 0, all, pos, part.length); + pos += part.length; + } + return new String(all, StandardCharsets.UTF_8); + } + + /** + * Reconstructs a state node from a checkpoint JSON string and a list of RFC-6902 patch JSON + * strings, applied in order via {@link JsonPatch#apply}. + * + * @param checkpointJson the full checkpoint state JSON + * @param diffPatchesJson the ordered list of opaque JSON-string patches (each an RFC-6902 op + * array) + * @return the reconstructed state node + * @throws Exception if any JSON cannot be parsed + */ + public static JsonNode reconstructState(String checkpointJson, List diffPatchesJson) + throws Exception { + JsonNode state = + (checkpointJson == null || checkpointJson.isEmpty()) + ? NullNode.getInstance() + : MAPPER.readTree(checkpointJson); + for (String patchJson : diffPatchesJson) { + if (patchJson == null || patchJson.isEmpty()) { + continue; + } + JsonNode patch = MAPPER.readTree(patchJson); + state = JsonPatch.apply(state, patch); + } + return state; + } + + /** + * Validates a snapshot/session document id. + * + * @param id the id to validate + * @throws GenkitException with {@code INVALID_ARGUMENT} if the id is null, empty, or contains a + * forward slash + */ + public static void validateId(String id) { + if (id == null || id.isEmpty()) { + throw GenkitException.builder() + .message("id must be non-empty") + .errorCode("INVALID_ARGUMENT") + .build(); + } + if (id.contains("/")) { + throw GenkitException.builder() + .message("id must not contain '/': " + id) + .errorCode("INVALID_ARGUMENT") + .build(); + } + } +} diff --git a/ai/src/main/java/com/google/genkit/ai/agent/internal/StreamEmitter.java b/ai/src/main/java/com/google/genkit/ai/agent/internal/StreamEmitter.java new file mode 100644 index 000000000..1b4f4bdb8 --- /dev/null +++ b/ai/src/main/java/com/google/genkit/ai/agent/internal/StreamEmitter.java @@ -0,0 +1,130 @@ +/* + * 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. + * + * SPDX-License-Identifier: Apache-2.0 + */ + +package com.google.genkit.ai.agent.internal; + +import com.fasterxml.jackson.databind.JsonNode; +import com.fasterxml.jackson.databind.ObjectMapper; +import com.google.genkit.ai.agent.AgentStreamChunk; +import com.google.genkit.ai.agent.Artifact; +import com.google.genkit.ai.agent.Session; +import com.google.genkit.core.jsonpatch.JsonPatch; +import java.util.function.Consumer; + +/** + * Bridges Session state-change events to AgentStreamChunk emission for one invocation. + * + *

Custom-state changes become customPatch chunks (first-of-turn = whole-doc replace, then + * incremental diffs); artifact add/update become artifact chunks. Suppressed entirely while + * detached. + * + * @param the type of the custom session state object + */ +public final class StreamEmitter { + + private final Consumer sink; + private final ObjectMapper mapper; + + /** Whether this emitter is suppressed (e.g. for detached runs). */ + private volatile boolean suppressed = false; + + /** + * True at the start of each turn. The first customPatch of a turn is always a whole-document + * replace; subsequent ones are incremental diffs. + */ + private boolean firstCustomPatchInTurn = true; + + /** + * The JsonNode of custom state after the last emitted customPatch. Used to compute incremental + * diffs. Null before the first patch of a turn. + */ + private JsonNode lastCustomJson = null; + + /** + * Constructs a new StreamEmitter. + * + * @param sink receives AgentStreamChunk instances; must not be null + * @param mapper ObjectMapper for converting custom state S to JsonNode; must not be null + */ + public StreamEmitter(Consumer sink, ObjectMapper mapper) { + this.sink = sink; + this.mapper = mapper; + } + + /** + * Attaches this emitter to a Session. Registers onCustomChanged and onArtifactChanged listeners + * that emit customPatch and artifact chunks respectively. + * + *

Since {@link Session#setOnCustomChanged(Runnable)} takes a {@code Runnable}, the current + * custom state is retrieved via {@code session.getCustom()} inside the callback. + * + * @param session the session to attach to + */ + public void attach(Session session) { + session.setOnCustomChanged( + () -> { + if (suppressed) { + return; + } + S custom = session.getCustom(); + JsonNode cur = mapper.valueToTree(custom); + + if (firstCustomPatchInTurn) { + // First patch of this turn: whole-document replace + JsonNode patch = JsonPatch.wholeDocumentReplace(cur); + sink.accept(AgentStreamChunk.builder().customPatch(patch).build()); + firstCustomPatchInTurn = false; + } else { + // Subsequent patches: incremental diff + JsonNode patch = JsonPatch.diff(lastCustomJson, cur); + if (patch.size() > 0) { + sink.accept(AgentStreamChunk.builder().customPatch(patch).build()); + } + // Skip emitting if the diff is empty (no-op update) + } + lastCustomJson = cur; + }); + + session.setOnArtifactChanged( + (Artifact artifact) -> { + if (suppressed) { + return; + } + sink.accept(AgentStreamChunk.builder().artifact(artifact).build()); + }); + } + + /** + * Call at the start of each turn. The next customPatch emitted will be a whole-document replace + * rather than an incremental diff. + */ + public void beginTurn() { + firstCustomPatchInTurn = true; + lastCustomJson = null; + } + + /** + * Controls whether chunks are emitted. When {@code true}, all custom-state and artifact changes + * are silently dropped. Use for detached (non-streaming) runs. + * + * @param suppressed {@code true} to suppress emission; {@code false} to re-enable + */ + public void setSuppressed(boolean suppressed) { + this.suppressed = suppressed; + } +} diff --git a/ai/src/main/java/com/google/genkit/ai/agent/package-info.java b/ai/src/main/java/com/google/genkit/ai/agent/package-info.java new file mode 100644 index 000000000..72b925123 --- /dev/null +++ b/ai/src/main/java/com/google/genkit/ai/agent/package-info.java @@ -0,0 +1,23 @@ +/* + * 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. + * + * SPDX-License-Identifier: Apache-2.0 + */ + +/** + * This package provides types for the Genkit agents feature, including enums and models for agent + * state snapshots, agent execution lifecycle, and agent-related operations. + */ +package com.google.genkit.ai.agent; diff --git a/ai/src/main/java/com/google/genkit/ai/session/Chat.java b/ai/src/main/java/com/google/genkit/ai/session/Chat.java deleted file mode 100644 index dc75299bd..000000000 --- a/ai/src/main/java/com/google/genkit/ai/session/Chat.java +++ /dev/null @@ -1,986 +0,0 @@ -/* - * 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. - * - * SPDX-License-Identifier: Apache-2.0 - */ - -package com.google.genkit.ai.session; - -import com.google.genkit.ai.*; -import com.google.genkit.ai.telemetry.ModelTelemetryHelper; -import com.google.genkit.core.Action; -import com.google.genkit.core.ActionContext; -import com.google.genkit.core.ActionType; -import com.google.genkit.core.GenkitException; -import com.google.genkit.core.JsonUtils; -import com.google.genkit.core.Registry; -import java.util.ArrayList; -import java.util.Collections; -import java.util.HashMap; -import java.util.List; -import java.util.Map; -import java.util.function.Consumer; - -/** - * Chat represents a conversation within a session thread. - * - *

Chat provides a simple interface for multi-turn conversations with automatic history - * management. Messages are persisted to the session store after each interaction. - * - *

Example usage: - * - * 

{@code
- * // Simple chat
- * Chat chat = session.chat();
- * ModelResponse response = chat.send("Hello!");
- *
- * // Chat with system prompt
- * Chat chat = session.chat(
- *     ChatOptions.builder()
- *         .model("openai/gpt-4o")
- *         .system("You are a helpful assistant.")
- *         .build());
- *
- * // Multi-turn conversation
- * chat.send("What is the capital of France?");
- * chat.send("And what about Germany?"); // Context is preserved
- * }
- * - * @param the type of the session state - */ -public class Chat { - - private static final String PREAMBLE_KEY = "preamble"; - - private final Session session; - private final String threadName; - private final ChatOptions originalOptions; - private final Registry registry; - private final Map effectiveAgentRegistry; - private List history; - - /** Pending interrupt requests from the last send. */ - private List pendingInterrupts; - - /** Current agent context (mutable for handoffs). */ - private String currentAgentName; - - private String currentSystem; - private String currentModel; - private List> currentTools; - - /** - * Creates a new Chat instance. - * - * @param session the parent session - * @param threadName the thread name - * @param options the chat options - * @param registry the Genkit registry - * @param sessionAgentRegistry the agent registry from the session (may be null) - */ - Chat( - Session session, - String threadName, - ChatOptions options, - Registry registry, - Map sessionAgentRegistry) { - this.session = session; - this.threadName = threadName; - this.originalOptions = options; - this.registry = registry; - // Use agent registry from options if provided, otherwise fall back to session's - // registry - this.effectiveAgentRegistry = - options.getAgentRegistry() != null ? options.getAgentRegistry() : sessionAgentRegistry; - this.history = new ArrayList<>(session.getMessages(threadName)); - this.pendingInterrupts = new ArrayList<>(); - - // Initialize current context from options - this.currentAgentName = null; - this.currentSystem = options.getSystem(); - this.currentModel = options.getModel(); - this.currentTools = options.getTools(); - } - - /** - * Sends a message and gets a response. - * - *

This method: - * - *

    - *
  1. Adds the user message to history - *
  2. Builds a request with all conversation history - *
  3. Sends to the model and gets a response - *
  4. Adds the model response to history - *
  5. Persists the updated history - *
- * - * @param text the user message - * @return the model response - * @throws GenkitException if generation fails - */ - public ModelResponse send(String text) throws GenkitException { - return send(Message.user(text)); - } - - /** - * Sends a message and gets a response. - * - * @param message the message to send - * @return the model response - * @throws GenkitException if generation fails - */ - public ModelResponse send(Message message) throws GenkitException { - return send(message, null); - } - - /** - * Sends a message with send options and gets a response. - * - * @param text the user message - * @param sendOptions additional options for this send - * @return the model response - * @throws GenkitException if generation fails - */ - public ModelResponse send(String text, SendOptions sendOptions) throws GenkitException { - return send(Message.user(text), sendOptions); - } - - /** - * Sends a message with send options and gets a response. - * - * @param message the message to send - * @param sendOptions additional options for this send - * @return the model response - * @throws GenkitException if generation fails - */ - public ModelResponse send(Message message, SendOptions sendOptions) throws GenkitException { - // Clear any pending interrupts from previous send - pendingInterrupts.clear(); - - // Check if we're resuming from an interrupt - ResumeOptions resumeOptions = (sendOptions != null) ? sendOptions.getResumeOptions() : null; - if (resumeOptions != null) { - return resumeFromInterrupt(resumeOptions, sendOptions); - } - - // Add user message to history - history.add(message); - - return executeGenerationLoop(sendOptions); - } - - /** - * Resumes generation after an interrupt. - * - * @param resumeOptions the resume options containing tool responses - * @param sendOptions additional send options - * @return the model response - * @throws GenkitException if generation fails - */ - private ModelResponse resumeFromInterrupt(ResumeOptions resumeOptions, SendOptions sendOptions) - throws GenkitException { - // Add tool responses to history - if (resumeOptions.getRespond() != null && !resumeOptions.getRespond().isEmpty()) { - List responseParts = new ArrayList<>(); - for (ToolResponse response : resumeOptions.getRespond()) { - Part part = new Part(); - part.setToolResponse(response); - responseParts.add(part); - } - Message toolResponseMessage = new Message(); - toolResponseMessage.setRole(Role.TOOL); - toolResponseMessage.setContent(responseParts); - history.add(toolResponseMessage); - } - - // Handle restart requests by re-executing those tools - if (resumeOptions.getRestart() != null && !resumeOptions.getRestart().isEmpty()) { - ActionContext ctx = new ActionContext(registry); - List restartParts = - executeToolsWithInterruptHandling(ctx, resumeOptions.getRestart(), sendOptions); - - // Check if any restarts also triggered interrupts - if (!pendingInterrupts.isEmpty()) { - persistHistory(); - return createInterruptResponse(); - } - - Message toolResponseMessage = new Message(); - toolResponseMessage.setRole(Role.TOOL); - toolResponseMessage.setContent(restartParts); - history.add(toolResponseMessage); - } - - return executeGenerationLoop(sendOptions); - } - - /** Executes the main generation loop with tool handling. */ - private ModelResponse executeGenerationLoop(SendOptions sendOptions) throws GenkitException { - // Build the request - ModelRequest request = buildRequest(sendOptions); - - // Get the model - String modelName = resolveModelName(sendOptions); - if (modelName == null) { - throw new GenkitException("No model specified. Set model in ChatOptions or SendOptions."); - } - - Model model = getModel(modelName); - ActionContext ctx = new ActionContext(registry); - - // Handle tool execution loop with session context - int maxTurns = resolveMaxTurns(sendOptions); - int turn = 0; - - while (turn < maxTurns) { - // Make request effectively final for lambda - final ModelRequest finalRequest = request; - final String flowName = ctx.getFlowName(); - ModelResponse response; - try { - response = - SessionContext.runWithSession( - session, - () -> - ModelTelemetryHelper.runWithTelemetry( - modelName, - flowName != null ? flowName : "chat", - "/chat/" + (threadName != null ? threadName : "default"), - finalRequest, - r -> model.run(ctx, r))); - } catch (GenkitException e) { - throw e; - } catch (Exception e) { - throw new GenkitException("Error during model execution", e); - } - - // Check for tool requests - List toolRequests = extractToolRequests(response); - if (toolRequests.isEmpty()) { - // No tool calls, add response to history and persist - Message responseMessage = response.getMessage(); - if (responseMessage != null) { - history.add(responseMessage); - persistHistory(); - } - return response; - } - - // Execute tools with interrupt handling - List toolResponseParts = - executeToolsWithInterruptHandling(ctx, toolRequests, sendOptions); - - // Add assistant message with tool requests to history - Message assistantMessage = response.getMessage(); - if (assistantMessage != null) { - history.add(assistantMessage); - } - - // Check if any tools triggered interrupts - if (!pendingInterrupts.isEmpty()) { - persistHistory(); - return createInterruptResponse(); - } - - // Add tool response message to history - Message toolResponseMessage = new Message(); - toolResponseMessage.setRole(Role.TOOL); - toolResponseMessage.setContent(toolResponseParts); - history.add(toolResponseMessage); - - // Rebuild request with updated history - request = buildRequest(sendOptions); - turn++; - } - - throw new GenkitException("Max tool execution turns (" + maxTurns + ") exceeded"); - } - - /** - * Sends a message with streaming response. - * - * @param text the user message - * @param streamCallback callback for each response chunk - * @return the final aggregated response - * @throws GenkitException if generation fails - */ - public ModelResponse sendStream(String text, Consumer streamCallback) - throws GenkitException { - return sendStream(Message.user(text), null, streamCallback); - } - - /** - * Sends a message with streaming response. - * - * @param message the message to send - * @param sendOptions additional options for this send - * @param streamCallback callback for each response chunk - * @return the final aggregated response - * @throws GenkitException if generation fails - */ - public ModelResponse sendStream( - Message message, SendOptions sendOptions, Consumer streamCallback) - throws GenkitException { - // Add user message to history - history.add(message); - - // Build the request - ModelRequest request = buildRequest(sendOptions); - - // Get the model - String modelName = resolveModelName(sendOptions); - if (modelName == null) { - throw new GenkitException("No model specified. Set model in ChatOptions or SendOptions."); - } - - Model model = getModel(modelName); - if (!model.supportsStreaming()) { - throw new GenkitException("Model " + modelName + " does not support streaming"); - } - - ActionContext ctx = new ActionContext(registry); - final String flowName = ctx.getFlowName(); - ModelResponse response = - ModelTelemetryHelper.runWithTelemetryStreaming( - modelName, - flowName != null ? flowName : "chat", - "/chat/" + (threadName != null ? threadName : "default"), - request, - r -> model.run(ctx, r, streamCallback)); - - // Add response to history and persist - Message responseMessage = response.getMessage(); - if (responseMessage != null) { - history.add(responseMessage); - persistHistory(); - } - - return response; - } - - /** - * Gets the current conversation history. - * - * @return a copy of the message history - */ - public List getHistory() { - return new ArrayList<>(history); - } - - /** - * Gets the session. - * - * @return the parent session - */ - public Session getSession() { - return session; - } - - /** - * Gets the thread name. - * - * @return the thread name - */ - public String getThreadName() { - return threadName; - } - - /** - * Gets the pending interrupt requests from the last send. - * - *

If the last {@link #send} call returned with interrupts, this method returns the list of - * pending interrupts that need to be resolved before continuing. - * - * @return the list of pending interrupt requests (empty if none) - */ - public List getPendingInterrupts() { - return new ArrayList<>(pendingInterrupts); - } - - /** - * Checks if there are pending interrupts. - * - * @return true if there are pending interrupts - */ - public boolean hasPendingInterrupts() { - return !pendingInterrupts.isEmpty(); - } - - /** - * Gets the current agent name. - * - *

Returns null if no agent handoff has occurred, otherwise returns the name of the agent that - * the conversation was most recently handed off to. - * - * @return the current agent name, or null if no handoff has occurred - */ - public String getCurrentAgentName() { - return currentAgentName; - } - - /** Builds a ModelRequest from current history and options. */ - private ModelRequest buildRequest(SendOptions sendOptions) { - ModelRequest.Builder builder = ModelRequest.builder(); - - // Build messages list with system prompt (preamble) - List messages = new ArrayList<>(); - - // Add system prompt if specified (use current context for handoffs) - String systemPrompt = currentSystem; - if (systemPrompt != null && !systemPrompt.isEmpty()) { - Message systemMessage = Message.system(systemPrompt); - // Mark as preamble in metadata - Map metadata = new HashMap<>(); - metadata.put(PREAMBLE_KEY, true); - systemMessage.setMetadata(metadata); - messages.add(systemMessage); - } - - // Add conversation history (excluding any existing preamble) - for (Message msg : history) { - if (!isPreamble(msg)) { - messages.add(msg); - } - } - - builder.messages(messages); - - // Add tools (use current context for handoffs) - List> tools = resolveTools(sendOptions); - if (tools != null && !tools.isEmpty()) { - List toolDefs = new ArrayList<>(); - for (Tool tool : tools) { - toolDefs.add(tool.getDefinition()); - } - builder.tools(toolDefs); - } - - // Add config - if (originalOptions.getConfig() != null) { - builder.config(convertConfigToMap(originalOptions.getConfig())); - } - - // Add output config - if (originalOptions.getOutput() != null) { - builder.output(originalOptions.getOutput()); - } - - return builder.build(); - } - - /** Checks if a message is a preamble (system prompt). */ - private boolean isPreamble(Message message) { - if (message.getMetadata() == null) { - return false; - } - Object preamble = message.getMetadata().get(PREAMBLE_KEY); - return Boolean.TRUE.equals(preamble); - } - - /** Resolves the model name from options. */ - private String resolveModelName(SendOptions sendOptions) { - if (sendOptions != null && sendOptions.getModel() != null) { - return sendOptions.getModel(); - } - // Use current context (which may have been updated by handoff) - return currentModel; - } - - /** Resolves the max turns from options. */ - private int resolveMaxTurns(SendOptions sendOptions) { - if (sendOptions != null && sendOptions.getMaxTurns() != null) { - return sendOptions.getMaxTurns(); - } - if (originalOptions.getMaxTurns() != null) { - return originalOptions.getMaxTurns(); - } - return 5; // Default - } - - /** Resolves the tools from options. */ - private List> resolveTools(SendOptions sendOptions) { - if (sendOptions != null && sendOptions.getTools() != null) { - return sendOptions.getTools(); - } - // Use current context (which may have been updated by handoff) - return currentTools; - } - - /** Gets a model by name from the registry. */ - private Model getModel(String name) { - Action action = registry.lookupAction(ActionType.MODEL, name); - if (action == null) { - throw new GenkitException("Model not found: " + name); - } - return (Model) action; - } - - /** Extracts tool requests from a model response. */ - private List extractToolRequests(ModelResponse response) { - List requests = new ArrayList<>(); - if (response.getCandidates() != null) { - for (Candidate candidate : response.getCandidates()) { - if (candidate.getMessage() != null && candidate.getMessage().getContent() != null) { - for (Part part : candidate.getMessage().getContent()) { - if (part.getToolRequest() != null) { - requests.add(part.getToolRequest()); - } - } - } - } - } - return requests; - } - - /** Executes tools and returns response parts. */ - private List executeTools( - ActionContext ctx, List toolRequests, SendOptions sendOptions) { - List responseParts = new ArrayList<>(); - List> tools = resolveTools(sendOptions); - - for (ToolRequest toolRequest : toolRequests) { - String toolName = toolRequest.getName(); - Object toolInput = toolRequest.getInput(); - - Tool tool = findTool(toolName, tools); - if (tool == null) { - Part errorPart = new Part(); - ToolResponse errorResponse = - new ToolResponse( - toolRequest.getRef(), - toolName, - Collections.singletonMap("error", "Tool not found: " + toolName)); - errorPart.setToolResponse(errorResponse); - responseParts.add(errorPart); - continue; - } - - try { - @SuppressWarnings("unchecked") - Tool typedTool = (Tool) tool; - - // Convert the input to the expected type if necessary - final Object convertedInput; - Class inputClass = typedTool.getInputClass(); - if (inputClass != null && toolInput != null && !inputClass.isInstance(toolInput)) { - convertedInput = JsonUtils.convert(toolInput, inputClass); - } else { - convertedInput = toolInput; - } - - Object result = - SessionContext.runWithSession(session, () -> typedTool.run(ctx, convertedInput)); - - Part responsePart = new Part(); - ToolResponse toolResponse = new ToolResponse(toolRequest.getRef(), toolName, result); - responsePart.setToolResponse(toolResponse); - responseParts.add(responsePart); - } catch (Exception e) { - Part errorPart = new Part(); - ToolResponse errorResponse = - new ToolResponse( - toolRequest.getRef(), - toolName, - Collections.singletonMap("error", "Tool execution failed: " + e.getMessage())); - errorPart.setToolResponse(errorResponse); - responseParts.add(errorPart); - } - } - - return responseParts; - } - - /** - * Executes tools with interrupt handling and returns response parts. - * - *

When a tool throws {@link ToolInterruptException}, the interrupt is captured and added to - * the pending interrupts list. The tool execution continues for other tools, and an interrupt - * response is returned after all tools have been processed. - * - *

When a tool throws {@link AgentHandoffException}, the chat context is switched to the target - * agent (system prompt, tools, model), enabling multi-agent conversations. - */ - private List executeToolsWithInterruptHandling( - ActionContext ctx, List toolRequests, SendOptions sendOptions) { - List responseParts = new ArrayList<>(); - List> tools = resolveTools(sendOptions); - - for (ToolRequest toolRequest : toolRequests) { - String toolName = toolRequest.getName(); - Object toolInput = toolRequest.getInput(); - - Tool tool = findTool(toolName, tools); - if (tool == null) { - Part errorPart = new Part(); - ToolResponse errorResponse = - new ToolResponse( - toolRequest.getRef(), - toolName, - Collections.singletonMap("error", "Tool not found: " + toolName)); - errorPart.setToolResponse(errorResponse); - responseParts.add(errorPart); - continue; - } - - try { - @SuppressWarnings("unchecked") - Tool typedTool = (Tool) tool; - - // Convert the input to the expected type if necessary - final Object convertedInput; - Class inputClass = typedTool.getInputClass(); - if (inputClass != null && toolInput != null && !inputClass.isInstance(toolInput)) { - convertedInput = JsonUtils.convert(toolInput, inputClass); - } else { - convertedInput = toolInput; - } - - Object result = - SessionContext.runWithSession(session, () -> typedTool.run(ctx, convertedInput)); - - Part responsePart = new Part(); - ToolResponse toolResponse = new ToolResponse(toolRequest.getRef(), toolName, result); - responsePart.setToolResponse(toolResponse); - responseParts.add(responsePart); - } catch (AgentHandoffException e) { - // Handle agent handoff - switch context to the target agent - handleAgentHandoff(e); - - // Add a response indicating the handoff - Part handoffPart = new Part(); - Map handoffOutput = new HashMap<>(); - handoffOutput.put("transferred", true); - handoffOutput.put("transferredTo", e.getTargetAgentName()); - handoffOutput.put("message", "Conversation transferred to " + e.getTargetAgentName()); - ToolResponse handoffResponse = - new ToolResponse(toolRequest.getRef(), toolName, handoffOutput); - handoffPart.setToolResponse(handoffResponse); - responseParts.add(handoffPart); - } catch (ToolInterruptException e) { - // Capture the interrupt - InterruptRequest interruptRequest = new InterruptRequest(toolRequest, e.getMetadata()); - pendingInterrupts.add(interruptRequest); - - // Add a placeholder response indicating interruption - Part interruptPart = new Part(); - Map interruptOutput = new HashMap<>(); - interruptOutput.put("__interrupt", true); - interruptOutput.put("metadata", e.getMetadata()); - ToolResponse interruptResponse = - new ToolResponse(toolRequest.getRef(), toolName, interruptOutput); - interruptPart.setToolResponse(interruptResponse); - responseParts.add(interruptPart); - } catch (Exception e) { - Part errorPart = new Part(); - ToolResponse errorResponse = - new ToolResponse( - toolRequest.getRef(), - toolName, - Collections.singletonMap("error", "Tool execution failed: " + e.getMessage())); - errorPart.setToolResponse(errorResponse); - responseParts.add(errorPart); - } - } - - return responseParts; - } - - /** Handles an agent handoff by switching the chat context. */ - private void handleAgentHandoff(AgentHandoffException handoff) { - AgentConfig targetConfig = handoff.getTargetAgentConfig(); - currentAgentName = handoff.getTargetAgentName(); - - // Update system prompt - if (targetConfig.getSystem() != null) { - currentSystem = targetConfig.getSystem(); - } - - // Update model if specified - if (targetConfig.getModel() != null) { - currentModel = targetConfig.getModel(); - } - - // Update tools - include the agent's tools plus sub-agent tools - List> newTools = new ArrayList<>(); - if (targetConfig.getTools() != null) { - newTools.addAll(targetConfig.getTools()); - } - - // Add sub-agents as tools if agent registry is available - if (targetConfig.getAgents() != null && effectiveAgentRegistry != null) { - for (AgentConfig subAgentConfig : targetConfig.getAgents()) { - Agent subAgent = effectiveAgentRegistry.get(subAgentConfig.getName()); - if (subAgent != null) { - newTools.add(subAgent.asTool()); - } - } - } - - currentTools = newTools; - } - - /** Creates a response indicating the generation was interrupted. */ - private ModelResponse createInterruptResponse() { - ModelResponse response = new ModelResponse(); - - // Create a message indicating interruption - Message interruptMessage = new Message(); - interruptMessage.setRole(Role.MODEL); - - Part textPart = new Part(); - textPart.setText("[Generation interrupted - awaiting user input]"); - interruptMessage.setContent(List.of(textPart)); - - // Add interrupt metadata - Map metadata = new HashMap<>(); - metadata.put("interrupted", true); - metadata.put("interruptCount", pendingInterrupts.size()); - List> interruptData = new ArrayList<>(); - for (InterruptRequest interrupt : pendingInterrupts) { - Map data = new HashMap<>(); - data.put("toolName", interrupt.getToolRequest().getName()); - data.put("toolRef", interrupt.getToolRequest().getRef()); - data.put("metadata", interrupt.getMetadata()); - interruptData.add(data); - } - metadata.put("interrupts", interruptData); - interruptMessage.setMetadata(metadata); - - // Create candidate - Candidate candidate = new Candidate(); - candidate.setMessage(interruptMessage); - candidate.setFinishReason(FinishReason.OTHER); - - response.setCandidates(List.of(candidate)); - - return response; - } - - /** Finds a tool by name. */ - private Tool findTool(String toolName, List> tools) { - if (tools != null) { - for (Tool tool : tools) { - if (tool.getName().equals(toolName)) { - return tool; - } - } - } - - // Try registry - Action action = registry.lookupAction(ActionType.TOOL, toolName); - if (action instanceof Tool) { - return (Tool) action; - } - - return null; - } - - /** Persists the current history to the session store. */ - private void persistHistory() { - session.updateMessages(threadName, history).join(); - } - - /** Converts GenerationConfig to a Map for the ModelRequest. */ - private Map convertConfigToMap(GenerationConfig config) { - Map configMap = new HashMap<>(); - if (config.getTemperature() != null) { - configMap.put("temperature", config.getTemperature()); - } - if (config.getMaxOutputTokens() != null) { - configMap.put("maxOutputTokens", config.getMaxOutputTokens()); - } - if (config.getTopP() != null) { - configMap.put("topP", config.getTopP()); - } - if (config.getTopK() != null) { - configMap.put("topK", config.getTopK()); - } - if (config.getStopSequences() != null) { - configMap.put("stopSequences", config.getStopSequences()); - } - if (config.getPresencePenalty() != null) { - configMap.put("presencePenalty", config.getPresencePenalty()); - } - if (config.getFrequencyPenalty() != null) { - configMap.put("frequencyPenalty", config.getFrequencyPenalty()); - } - if (config.getSeed() != null) { - configMap.put("seed", config.getSeed()); - } - if (config.getCustom() != null) { - configMap.putAll(config.getCustom()); - } - return configMap; - } - - /** Options for individual send operations. */ - public static class SendOptions { - private String model; - private List> tools; - private Integer maxTurns; - private ResumeOptions resumeOptions; - - /** Default constructor. */ - public SendOptions() {} - - /** - * Gets the model name. - * - * @return the model name - */ - public String getModel() { - return model; - } - - /** - * Sets the model name. - * - * @param model the model name - */ - public void setModel(String model) { - this.model = model; - } - - /** - * Gets the tools. - * - * @return the tools - */ - public List> getTools() { - return tools; - } - - /** - * Sets the tools. - * - * @param tools the tools - */ - public void setTools(List> tools) { - this.tools = tools; - } - - /** - * Gets the max turns. - * - * @return the max turns - */ - public Integer getMaxTurns() { - return maxTurns; - } - - /** - * Sets the max turns. - * - * @param maxTurns the max turns - */ - public void setMaxTurns(Integer maxTurns) { - this.maxTurns = maxTurns; - } - - /** - * Gets the resume options. - * - * @return the resume options - */ - public ResumeOptions getResumeOptions() { - return resumeOptions; - } - - /** - * Sets the resume options. - * - * @param resumeOptions the resume options - */ - public void setResumeOptions(ResumeOptions resumeOptions) { - this.resumeOptions = resumeOptions; - } - - /** - * Creates a builder for SendOptions. - * - * @return a new builder - */ - public static Builder builder() { - return new Builder(); - } - - /** Builder for SendOptions. */ - public static class Builder { - private String model; - private List> tools; - private Integer maxTurns; - private ResumeOptions resumeOptions; - - /** - * Sets the model name. - * - * @param model the model name - * @return this builder - */ - public Builder model(String model) { - this.model = model; - return this; - } - - /** - * Sets the tools. - * - * @param tools the tools - * @return this builder - */ - public Builder tools(List> tools) { - this.tools = tools; - return this; - } - - /** - * Sets the max turns. - * - * @param maxTurns the max turns - * @return this builder - */ - public Builder maxTurns(Integer maxTurns) { - this.maxTurns = maxTurns; - return this; - } - - /** - * Sets the resume options for resuming after an interrupt. - * - * @param resumeOptions the resume options - * @return this builder - */ - public Builder resumeOptions(ResumeOptions resumeOptions) { - this.resumeOptions = resumeOptions; - return this; - } - - /** - * Builds the SendOptions. - * - * @return the built SendOptions - */ - public SendOptions build() { - SendOptions options = new SendOptions(); - options.setModel(model); - options.setTools(tools); - options.setMaxTurns(maxTurns); - options.setResumeOptions(resumeOptions); - return options; - } - } - } -} diff --git a/ai/src/main/java/com/google/genkit/ai/session/ChatOptions.java b/ai/src/main/java/com/google/genkit/ai/session/ChatOptions.java deleted file mode 100644 index c16d8f8cf..000000000 --- a/ai/src/main/java/com/google/genkit/ai/session/ChatOptions.java +++ /dev/null @@ -1,322 +0,0 @@ -/* - * 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. - * - * SPDX-License-Identifier: Apache-2.0 - */ - -package com.google.genkit.ai.session; - -import com.google.genkit.ai.Agent; -import com.google.genkit.ai.GenerationConfig; -import com.google.genkit.ai.OutputConfig; -import com.google.genkit.ai.Tool; -import java.util.List; -import java.util.Map; - -/** - * ChatOptions provides configuration options for creating a Chat instance. - * - * @param the type of the session state - */ -public class ChatOptions { - - private String model; - private String system; - private List> tools; - private OutputConfig output; - private GenerationConfig config; - private Map context; - private Integer maxTurns; - private Map agentRegistry; - - /** Default constructor. */ - public ChatOptions() {} - - /** - * Gets the model name. - * - * @return the model name - */ - public String getModel() { - return model; - } - - /** - * Sets the model name. - * - * @param model the model name - */ - public void setModel(String model) { - this.model = model; - } - - /** - * Gets the system prompt. - * - * @return the system prompt - */ - public String getSystem() { - return system; - } - - /** - * Sets the system prompt. - * - * @param system the system prompt - */ - public void setSystem(String system) { - this.system = system; - } - - /** - * Gets the available tools. - * - * @return the tools - */ - public List> getTools() { - return tools; - } - - /** - * Sets the available tools. - * - * @param tools the tools - */ - public void setTools(List> tools) { - this.tools = tools; - } - - /** - * Gets the output configuration. - * - * @return the output configuration - */ - public OutputConfig getOutput() { - return output; - } - - /** - * Sets the output configuration. - * - * @param output the output configuration - */ - public void setOutput(OutputConfig output) { - this.output = output; - } - - /** - * Gets the generation configuration. - * - * @return the generation configuration - */ - public GenerationConfig getConfig() { - return config; - } - - /** - * Sets the generation configuration. - * - * @param config the generation configuration - */ - public void setConfig(GenerationConfig config) { - this.config = config; - } - - /** - * Gets the additional context. - * - * @return the context - */ - public Map getContext() { - return context; - } - - /** - * Sets the additional context. - * - * @param context the context - */ - public void setContext(Map context) { - this.context = context; - } - - /** - * Gets the maximum conversation turns. - * - * @return the max turns - */ - public Integer getMaxTurns() { - return maxTurns; - } - - /** - * Sets the maximum conversation turns. - * - * @param maxTurns the max turns - */ - public void setMaxTurns(Integer maxTurns) { - this.maxTurns = maxTurns; - } - - /** - * Gets the agent registry for multi-agent handoffs. - * - * @return the agent registry - */ - public Map getAgentRegistry() { - return agentRegistry; - } - - /** - * Sets the agent registry for multi-agent handoffs. - * - * @param agentRegistry the agent registry - */ - public void setAgentRegistry(Map agentRegistry) { - this.agentRegistry = agentRegistry; - } - - /** - * Creates a builder for ChatOptions. - * - * @param the state type - * @return a new builder - */ - public static Builder builder() { - return new Builder<>(); - } - - /** - * Builder for ChatOptions. - * - * @param the state type - */ - public static class Builder { - private String model; - private String system; - private List> tools; - private OutputConfig output; - private GenerationConfig config; - private Map context; - private Integer maxTurns; - private Map agentRegistry; - - /** - * Sets the model name. - * - * @param model the model name - * @return this builder - */ - public Builder model(String model) { - this.model = model; - return this; - } - - /** - * Sets the system prompt. - * - * @param system the system prompt - * @return this builder - */ - public Builder system(String system) { - this.system = system; - return this; - } - - /** - * Sets the available tools. - * - * @param tools the tools - * @return this builder - */ - public Builder tools(List> tools) { - this.tools = tools; - return this; - } - - /** - * Sets the output configuration. - * - * @param output the output configuration - * @return this builder - */ - public Builder output(OutputConfig output) { - this.output = output; - return this; - } - - /** - * Sets the generation configuration. - * - * @param config the generation configuration - * @return this builder - */ - public Builder config(GenerationConfig config) { - this.config = config; - return this; - } - - /** - * Sets the additional context. - * - * @param context the context - * @return this builder - */ - public Builder context(Map context) { - this.context = context; - return this; - } - - /** - * Sets the maximum conversation turns. - * - * @param maxTurns the max turns - * @return this builder - */ - public Builder maxTurns(Integer maxTurns) { - this.maxTurns = maxTurns; - return this; - } - - /** - * Sets the agent registry for multi-agent handoffs. - * - * @param agentRegistry the agent registry - * @return this builder - */ - public Builder agentRegistry(Map agentRegistry) { - this.agentRegistry = agentRegistry; - return this; - } - - /** - * Builds the ChatOptions. - * - * @return the built ChatOptions - */ - public ChatOptions build() { - ChatOptions options = new ChatOptions<>(); - options.setModel(model); - options.setSystem(system); - options.setTools(tools); - options.setOutput(output); - options.setConfig(config); - options.setContext(context); - options.setMaxTurns(maxTurns); - options.setAgentRegistry(agentRegistry); - return options; - } - } -} diff --git a/ai/src/main/java/com/google/genkit/ai/session/InMemorySessionStore.java b/ai/src/main/java/com/google/genkit/ai/session/InMemorySessionStore.java deleted file mode 100644 index 0484db629..000000000 --- a/ai/src/main/java/com/google/genkit/ai/session/InMemorySessionStore.java +++ /dev/null @@ -1,83 +0,0 @@ -/* - * 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. - * - * SPDX-License-Identifier: Apache-2.0 - */ - -package com.google.genkit.ai.session; - -import java.util.Map; -import java.util.concurrent.CompletableFuture; -import java.util.concurrent.ConcurrentHashMap; - -/** - * InMemorySessionStore is a simple in-memory implementation of SessionStore. - * - *

This implementation is suitable for: - * - *

    - *
  • Development and testing - *
  • Single-instance deployments - *
  • Prototyping - *
- * - *

Note: Sessions are lost when the application restarts. For production use cases - * requiring persistence, implement a database-backed SessionStore. - * - * @param the type of the custom session state - */ -public class InMemorySessionStore implements SessionStore { - - private final Map> data = new ConcurrentHashMap<>(); - - /** Creates a new InMemorySessionStore. */ - public InMemorySessionStore() {} - - @Override - public CompletableFuture> get(String sessionId) { - return CompletableFuture.completedFuture(data.get(sessionId)); - } - - @Override - public CompletableFuture save(String sessionId, SessionData sessionData) { - data.put(sessionId, sessionData); - return CompletableFuture.completedFuture(null); - } - - @Override - public CompletableFuture delete(String sessionId) { - data.remove(sessionId); - return CompletableFuture.completedFuture(null); - } - - @Override - public CompletableFuture exists(String sessionId) { - return CompletableFuture.completedFuture(data.containsKey(sessionId)); - } - - /** - * Returns the number of sessions currently stored. - * - * @return the session count - */ - public int size() { - return data.size(); - } - - /** Clears all sessions from the store. */ - public void clear() { - data.clear(); - } -} diff --git a/ai/src/main/java/com/google/genkit/ai/session/Session.java b/ai/src/main/java/com/google/genkit/ai/session/Session.java deleted file mode 100644 index 7a1b97ae2..000000000 --- a/ai/src/main/java/com/google/genkit/ai/session/Session.java +++ /dev/null @@ -1,322 +0,0 @@ -/* - * 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. - * - * SPDX-License-Identifier: Apache-2.0 - */ - -package com.google.genkit.ai.session; - -import com.google.genkit.ai.Agent; -import com.google.genkit.ai.Message; -import com.google.genkit.core.Registry; -import java.util.ArrayList; -import java.util.List; -import java.util.Map; -import java.util.UUID; -import java.util.concurrent.CompletableFuture; -import java.util.function.Supplier; - -/** - * Session represents a stateful chat session that persists conversation history and custom state - * across multiple interactions. - * - *

Sessions provide: - * - *

    - *
  • Persistent conversation threads - *
  • Custom state management - *
  • Multiple named chat threads within a session - *
  • Automatic history management - *
- * - *

Example usage: - * - * 

{@code
- * // Create a session with initial state
- * Session session = genkit
- *     .createSession(SessionOptions.builder().initialState(new MyState("John")).build());
- *
- * // Create a chat and interact
- * Chat chat = session.chat();
- * ModelResponse response = chat.send("Hello!");
- *
- * // Access session state
- * MyState state = session.getState();
- * }
- * - * @param the type of the custom session state - */ -public class Session { - - /** Default thread name for chat conversations. */ - public static final String DEFAULT_THREAD = "main"; - - private final String id; - private SessionData sessionData; - private final SessionStore store; - private final Registry registry; - private final Supplier> chatFactory; - private final Map agentRegistry; - - /** - * Creates a new Session. - * - * @param registry the Genkit registry - * @param store the session store - * @param sessionData the initial session data - * @param chatFactory factory for creating Chat instances - * @param agentRegistry the agent registry for multi-agent handoffs (may be null) - */ - Session( - Registry registry, - SessionStore store, - SessionData sessionData, - Supplier> chatFactory, - Map agentRegistry) { - this.registry = registry; - this.store = store; - this.sessionData = sessionData; - this.id = sessionData.getId(); - this.chatFactory = chatFactory; - this.agentRegistry = agentRegistry; - } - - /** - * Gets the session ID. - * - * @return the unique session identifier - */ - public String getId() { - return id; - } - - /** - * Gets the current session state. - * - * @return the session state, or null if not set - */ - public S getState() { - return sessionData.getState(); - } - - /** - * Updates the session state and persists it. - * - * @param state the new state - * @return a CompletableFuture that completes when the state is saved - */ - public CompletableFuture updateState(S state) { - sessionData.setState(state); - return store.save(id, sessionData); - } - - /** - * Gets the message history for a thread. - * - * @param threadName the thread name - * @return the list of messages in the thread - */ - public List getMessages(String threadName) { - List messages = sessionData.getThread(threadName); - return messages != null ? new ArrayList<>(messages) : new ArrayList<>(); - } - - /** - * Gets the message history for the default thread. - * - * @return the list of messages - */ - public List getMessages() { - return getMessages(DEFAULT_THREAD); - } - - /** - * Updates the messages for a thread and persists them. - * - * @param threadName the thread name - * @param messages the messages to save - * @return a CompletableFuture that completes when saved - */ - public CompletableFuture updateMessages(String threadName, List messages) { - sessionData.setThread(threadName, messages); - return store.save(id, sessionData); - } - - /** - * Creates a new Chat instance for the default thread. - * - * @return a new Chat instance - */ - public Chat chat() { - return chat(DEFAULT_THREAD, ChatOptions.builder().build()); - } - - /** - * Creates a new Chat instance with options. - * - * @param options the chat options - * @return a new Chat instance - */ - public Chat chat(ChatOptions options) { - return chat(DEFAULT_THREAD, options); - } - - /** - * Creates a new Chat instance for a specific thread. - * - * @param threadName the thread name - * @return a new Chat instance - */ - public Chat chat(String threadName) { - return chat(threadName, ChatOptions.builder().build()); - } - - /** - * Creates a new Chat instance for a specific thread with options. - * - * @param threadName the thread name - * @param options the chat options - * @return a new Chat instance - */ - public Chat chat(String threadName, ChatOptions options) { - return new Chat<>(this, threadName, options, registry, agentRegistry); - } - - /** - * Gets the session store. - * - * @return the session store - */ - public SessionStore getStore() { - return store; - } - - /** - * Gets the registry. - * - * @return the registry - */ - public Registry getRegistry() { - return registry; - } - - /** - * Gets the agent registry for multi-agent handoffs. - * - * @return the agent registry, or null if not set - */ - public Map getAgentRegistry() { - return agentRegistry; - } - - /** - * Gets the session data. - * - * @return the session data - */ - public SessionData getSessionData() { - return sessionData; - } - - /** - * Serializes the session to JSON-compatible data. - * - * @return the session data - */ - public SessionData toJSON() { - return sessionData; - } - - /** - * Creates a new Session with a generated ID. - * - * @param the state type - * @param registry the Genkit registry - * @param options the session options - * @return a new Session - */ - public static Session create(Registry registry, SessionOptions options) { - return create(registry, options, null); - } - - /** - * Creates a new Session with a generated ID and agent registry. - * - * @param the state type - * @param registry the Genkit registry - * @param options the session options - * @param agentRegistry the agent registry for multi-agent handoffs (may be null) - * @return a new Session - */ - public static Session create( - Registry registry, SessionOptions options, Map agentRegistry) { - String sessionId = - options.getSessionId() != null ? options.getSessionId() : UUID.randomUUID().toString(); - - SessionStore store = - options.getStore() != null ? options.getStore() : new InMemorySessionStore<>(); - - SessionData data = - SessionData.builder().id(sessionId).state(options.getInitialState()).build(); - - // Save initial session data - store.save(sessionId, data).join(); - - return new Session<>(registry, store, data, null, agentRegistry); - } - - /** - * Loads an existing session from a store. - * - * @param the state type - * @param registry the Genkit registry - * @param sessionId the session ID to load - * @param options the session options (must include store) - * @return a CompletableFuture containing the loaded session, or null if not found - */ - public static CompletableFuture> load( - Registry registry, String sessionId, SessionOptions options) { - return load(registry, sessionId, options, null); - } - - /** - * Loads an existing session from a store with agent registry. - * - * @param the state type - * @param registry the Genkit registry - * @param sessionId the session ID to load - * @param options the session options (must include store) - * @param agentRegistry the agent registry for multi-agent handoffs (may be null) - * @return a CompletableFuture containing the loaded session, or null if not found - */ - public static CompletableFuture> load( - Registry registry, - String sessionId, - SessionOptions options, - Map agentRegistry) { - SessionStore store = - options.getStore() != null ? options.getStore() : new InMemorySessionStore<>(); - - return store - .get(sessionId) - .thenApply( - data -> { - if (data == null) { - return null; - } - return new Session<>(registry, store, data, null, agentRegistry); - }); - } -} diff --git a/ai/src/main/java/com/google/genkit/ai/session/SessionContext.java b/ai/src/main/java/com/google/genkit/ai/session/SessionContext.java deleted file mode 100644 index 234c216b9..000000000 --- a/ai/src/main/java/com/google/genkit/ai/session/SessionContext.java +++ /dev/null @@ -1,167 +0,0 @@ -/* - * 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. - * - * SPDX-License-Identifier: Apache-2.0 - */ - -package com.google.genkit.ai.session; - -import com.google.genkit.core.GenkitException; -import java.util.concurrent.Callable; - -/** - * Provides access to the current session context. - * - *

This class uses ThreadLocal to store the current session, making it accessible from within - * tool execution. This enables tools to access and modify session state during their execution. - * - *

Example usage in a tool: - * - * 

{@code
- * Tool myTool = genkit.defineTool(
- *     ToolConfig.builder()
- *         .name("myTool")
- *         .description("A tool that accesses session state")
- *         .inputSchema(Input.class)
- *         .outputSchema(Output.class)
- *         .build(),
- *     (input, ctx) -> {
- *       // Access current session from within tool
- *       Session session = SessionContext.currentSession();
- *       MyState state = session.getState();
- *
- *       // Update session state
- *       session.updateState(new MyState(state.getName(), state.getCount() + 1));
- *
- *       return new Output("Updated");
- *     });
- * }
- */ -public final class SessionContext { - - private static final ThreadLocal> CURRENT_SESSION = new ThreadLocal<>(); - - private SessionContext() {} - - /** - * Gets the current session. - * - * @param the session state type - * @return the current session - * @throws SessionException if not running within a session - */ - @SuppressWarnings("unchecked") - public static Session currentSession() { - Session session = CURRENT_SESSION.get(); - if (session == null) { - throw new SessionException("Not running within a session context"); - } - return (Session) session; - } - - /** - * Gets the current session if available. - * - * @param the session state type - * @return the current session, or null if not in a session context - */ - @SuppressWarnings("unchecked") - public static Session getCurrentSession() { - return (Session) CURRENT_SESSION.get(); - } - - /** - * Checks if currently running within a session context. - * - * @return true if in a session context - */ - public static boolean hasSession() { - return CURRENT_SESSION.get() != null; - } - - /** - * Runs a function within a session context. - * - * @param the session state type - * @param the return type - * @param session the session to use - * @param callable the function to run - * @return the result of the function - * @throws Exception if the function throws an exception - */ - public static T runWithSession(Session session, Callable callable) throws Exception { - Session previous = CURRENT_SESSION.get(); - try { - CURRENT_SESSION.set(session); - return callable.call(); - } finally { - if (previous != null) { - CURRENT_SESSION.set(previous); - } else { - CURRENT_SESSION.remove(); - } - } - } - - /** - * Runs a runnable within a session context. - * - * @param the session state type - * @param session the session to use - * @param runnable the runnable to execute - */ - public static void runWithSession(Session session, Runnable runnable) { - Session previous = CURRENT_SESSION.get(); - try { - CURRENT_SESSION.set(session); - runnable.run(); - } finally { - if (previous != null) { - CURRENT_SESSION.set(previous); - } else { - CURRENT_SESSION.remove(); - } - } - } - - /** - * Sets the current session. This is typically called internally by Chat. - * - * @param session the session to set - */ - public static void setSession(Session session) { - if (session != null) { - CURRENT_SESSION.set(session); - } else { - CURRENT_SESSION.remove(); - } - } - - /** Clears the current session. */ - public static void clearSession() { - CURRENT_SESSION.remove(); - } - - /** Exception thrown when session operations fail. */ - public static class SessionException extends GenkitException { - public SessionException(String message) { - super(message); - } - - public SessionException(String message, Throwable cause) { - super(message, cause); - } - } -} diff --git a/ai/src/main/java/com/google/genkit/ai/session/SessionData.java b/ai/src/main/java/com/google/genkit/ai/session/SessionData.java deleted file mode 100644 index a7292b61a..000000000 --- a/ai/src/main/java/com/google/genkit/ai/session/SessionData.java +++ /dev/null @@ -1,242 +0,0 @@ -/* - * 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. - * - * SPDX-License-Identifier: Apache-2.0 - */ - -package com.google.genkit.ai.session; - -import com.fasterxml.jackson.annotation.JsonInclude; -import com.fasterxml.jackson.annotation.JsonProperty; -import com.google.genkit.ai.Message; -import java.util.ArrayList; -import java.util.HashMap; -import java.util.List; -import java.util.Map; - -/** - * SessionData represents the persistent data structure for a session, including state and - * conversation threads. - * - * @param the type of the custom session state - */ -@JsonInclude(JsonInclude.Include.NON_NULL) -public class SessionData { - - /** The unique identifier for this session. */ - @JsonProperty("id") - private String id; - - /** Custom user-defined state associated with the session. */ - @JsonProperty("state") - private S state; - - /** - * Named conversation threads. Each thread is identified by a string key (default is "main") and - * contains a list of messages. - */ - @JsonProperty("threads") - private Map> threads; - - /** Default constructor. */ - public SessionData() { - this.threads = new HashMap<>(); - } - - /** - * Creates a new SessionData with the given ID. - * - * @param id the session ID - */ - public SessionData(String id) { - this.id = id; - this.threads = new HashMap<>(); - } - - /** - * Creates a new SessionData with the given ID and initial state. - * - * @param id the session ID - * @param state the initial state - */ - public SessionData(String id, S state) { - this.id = id; - this.state = state; - this.threads = new HashMap<>(); - } - - /** - * Gets the session ID. - * - * @return the session ID - */ - public String getId() { - return id; - } - - /** - * Sets the session ID. - * - * @param id the session ID - */ - public void setId(String id) { - this.id = id; - } - - /** - * Gets the session state. - * - * @return the session state - */ - public S getState() { - return state; - } - - /** - * Sets the session state. - * - * @param state the session state - */ - public void setState(S state) { - this.state = state; - } - - /** - * Gets all conversation threads. - * - * @return the threads map - */ - public Map> getThreads() { - return threads; - } - - /** - * Sets all conversation threads. - * - * @param threads the threads map - */ - public void setThreads(Map> threads) { - this.threads = threads; - } - - /** - * Gets a specific thread by name. - * - * @param threadName the thread name - * @return the list of messages in the thread, or null if not found - */ - public List getThread(String threadName) { - return threads.get(threadName); - } - - /** - * Gets or creates a thread by name. - * - * @param threadName the thread name - * @return the list of messages in the thread - */ - public List getOrCreateThread(String threadName) { - return threads.computeIfAbsent(threadName, k -> new ArrayList<>()); - } - - /** - * Sets messages for a specific thread. - * - * @param threadName the thread name - * @param messages the messages to set - */ - public void setThread(String threadName, List messages) { - threads.put(threadName, new ArrayList<>(messages)); - } - - /** - * Creates a builder for SessionData. - * - * @param the state type - * @return a new builder - */ - public static Builder builder() { - return new Builder<>(); - } - - /** - * Builder for SessionData. - * - * @param the state type - */ - public static class Builder { - private String id; - private S state; - private Map> threads = new HashMap<>(); - - /** - * Sets the session ID. - * - * @param id the session ID - * @return this builder - */ - public Builder id(String id) { - this.id = id; - return this; - } - - /** - * Sets the session state. - * - * @param state the session state - * @return this builder - */ - public Builder state(S state) { - this.state = state; - return this; - } - - /** - * Sets the conversation threads. - * - * @param threads the threads map - * @return this builder - */ - public Builder threads(Map> threads) { - this.threads = new HashMap<>(threads); - return this; - } - - /** - * Adds a thread. - * - * @param threadName the thread name - * @param messages the messages - * @return this builder - */ - public Builder thread(String threadName, List messages) { - this.threads.put(threadName, new ArrayList<>(messages)); - return this; - } - - /** - * Builds the SessionData. - * - * @return the built SessionData - */ - public SessionData build() { - SessionData data = new SessionData<>(); - data.setId(id); - data.setState(state); - data.setThreads(threads); - return data; - } - } -} diff --git a/ai/src/main/java/com/google/genkit/ai/session/SessionOptions.java b/ai/src/main/java/com/google/genkit/ai/session/SessionOptions.java deleted file mode 100644 index 109d4370a..000000000 --- a/ai/src/main/java/com/google/genkit/ai/session/SessionOptions.java +++ /dev/null @@ -1,155 +0,0 @@ -/* - * 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. - * - * SPDX-License-Identifier: Apache-2.0 - */ - -package com.google.genkit.ai.session; - -/** - * SessionOptions provides configuration options for creating or loading sessions. - * - * @param the type of the custom session state - */ -public class SessionOptions { - - private SessionStore store; - private S initialState; - private String sessionId; - - /** Default constructor. */ - public SessionOptions() {} - - /** - * Gets the session store. - * - * @return the session store - */ - public SessionStore getStore() { - return store; - } - - /** - * Sets the session store. - * - * @param store the session store - */ - public void setStore(SessionStore store) { - this.store = store; - } - - /** - * Gets the initial state. - * - * @return the initial state - */ - public S getInitialState() { - return initialState; - } - - /** - * Sets the initial state. - * - * @param initialState the initial state - */ - public void setInitialState(S initialState) { - this.initialState = initialState; - } - - /** - * Gets the session ID. - * - * @return the session ID - */ - public String getSessionId() { - return sessionId; - } - - /** - * Sets the session ID. - * - * @param sessionId the session ID - */ - public void setSessionId(String sessionId) { - this.sessionId = sessionId; - } - - /** - * Creates a builder for SessionOptions. - * - * @param the state type - * @return a new builder - */ - public static Builder builder() { - return new Builder<>(); - } - - /** - * Builder for SessionOptions. - * - * @param the state type - */ - public static class Builder { - private SessionStore store; - private S initialState; - private String sessionId; - - /** - * Sets the session store. - * - * @param store the session store - * @return this builder - */ - public Builder store(SessionStore store) { - this.store = store; - return this; - } - - /** - * Sets the initial state. - * - * @param initialState the initial state - * @return this builder - */ - public Builder initialState(S initialState) { - this.initialState = initialState; - return this; - } - - /** - * Sets the session ID. - * - * @param sessionId the session ID - * @return this builder - */ - public Builder sessionId(String sessionId) { - this.sessionId = sessionId; - return this; - } - - /** - * Builds the SessionOptions. - * - * @return the built SessionOptions - */ - public SessionOptions build() { - SessionOptions options = new SessionOptions<>(); - options.setStore(store); - options.setInitialState(initialState); - options.setSessionId(sessionId); - return options; - } - } -} diff --git a/ai/src/main/java/com/google/genkit/ai/session/SessionStore.java b/ai/src/main/java/com/google/genkit/ai/session/SessionStore.java deleted file mode 100644 index 3fb752953..000000000 --- a/ai/src/main/java/com/google/genkit/ai/session/SessionStore.java +++ /dev/null @@ -1,75 +0,0 @@ -/* - * 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. - * - * SPDX-License-Identifier: Apache-2.0 - */ - -package com.google.genkit.ai.session; - -import java.util.concurrent.CompletableFuture; - -/** - * SessionStore is an interface for persisting session data. - * - *

Implementations can provide different storage backends such as: - * - *

    - *
  • In-memory storage (for development/testing) - *
  • Database storage (for production) - *
  • Redis or other distributed cache - *
  • File-based storage - *
- * - * @param the type of the custom session state - */ -public interface SessionStore { - - /** - * Retrieves a session by its ID. - * - * @param sessionId the session ID - * @return a CompletableFuture containing the session data, or null if not found - */ - CompletableFuture> get(String sessionId); - - /** - * Saves session data. - * - * @param sessionId the session ID - * @param data the session data to save - * @return a CompletableFuture that completes when the save is done - */ - CompletableFuture save(String sessionId, SessionData data); - - /** - * Deletes a session by its ID. - * - * @param sessionId the session ID - * @return a CompletableFuture that completes when the deletion is done - */ - default CompletableFuture delete(String sessionId) { - return CompletableFuture.completedFuture(null); - } - - /** - * Checks if a session exists. - * - * @param sessionId the session ID - * @return a CompletableFuture containing true if the session exists - */ - default CompletableFuture exists(String sessionId) { - return get(sessionId).thenApply(data -> data != null); - } -} diff --git a/ai/src/main/java/com/google/genkit/ai/session/package-info.java b/ai/src/main/java/com/google/genkit/ai/session/package-info.java deleted file mode 100644 index faaa08ad2..000000000 --- a/ai/src/main/java/com/google/genkit/ai/session/package-info.java +++ /dev/null @@ -1,73 +0,0 @@ -/* - * 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. - * - * SPDX-License-Identifier: Apache-2.0 - */ - -/** - * Provides session management for multi-turn agent conversations with persistence. - * - *

The session package provides a stateful layer on top of Genkit's generation capabilities, - * enabling: - * - *

    - *
  • Persistent conversation history across multiple interactions - *
  • Custom session state management - *
  • Multiple named conversation threads within a session - *
  • Pluggable storage backends via {@link com.google.genkit.ai.session.SessionStore} - *
- * - *

Key Components

- * - *
    - *
  • {@link com.google.genkit.ai.session.Session} - The main entry point for session management - *
  • {@link com.google.genkit.ai.session.Chat} - Manages conversations within a session thread - *
  • {@link com.google.genkit.ai.session.SessionStore} - Interface for session persistence - *
  • {@link com.google.genkit.ai.session.InMemorySessionStore} - Default in-memory - * implementation - *
- * - *

Example Usage

- * - *

Create a session with custom state: - * - * 

- * Session session = genkit
- *     .createSession(SessionOptions.builder().initialState(new MyState("John")).build());
- *
- * Chat chat = session.chat(
- *     ChatOptions.builder().model("openai/gpt-4o").system("You are a helpful assistant.").build());
- *
- * // Multi-turn conversation (history is preserved automatically)
- * chat.send("What is the capital of France?");
- * chat.send("And what about Germany?");
- *
- * // Access session state
- * MyState state = session.getState();
- *
- * // Load an existing session
- * Session loadedSession = genkit.loadSession(sessionId, options).get();
- *
- * - *

Custom Session Stores

- * - *

Implement {@link com.google.genkit.ai.session.SessionStore} to provide custom persistence - * backends (e.g., database, Redis, file system). - * - * @see com.google.genkit.ai.session.Session - * @see com.google.genkit.ai.session.Chat - * @see com.google.genkit.ai.session.SessionStore - */ -package com.google.genkit.ai.session; diff --git a/ai/src/test/java/com/google/genkit/ai/AgentConfigTest.java b/ai/src/test/java/com/google/genkit/ai/AgentConfigTest.java deleted file mode 100644 index 75f05f9de..000000000 --- a/ai/src/test/java/com/google/genkit/ai/AgentConfigTest.java +++ /dev/null @@ -1,174 +0,0 @@ -/* - * 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. - * - * SPDX-License-Identifier: Apache-2.0 - */ - -package com.google.genkit.ai; - -import static org.junit.jupiter.api.Assertions.*; - -import java.util.HashMap; -import java.util.List; -import java.util.Map; -import org.junit.jupiter.api.Test; - -/** Unit tests for AgentConfig. */ -class AgentConfigTest { - - /** Helper to create a simple test tool. */ - private Tool createTestTool(String name) { - Map schema = new HashMap<>(); - schema.put("type", "string"); - return new Tool<>( - name, "Test tool " + name, schema, schema, String.class, (ctx, input) -> "result"); - } - - @Test - void testBuilderWithAllFields() { - Tool tool1 = createTestTool("tool1"); - Tool tool2 = createTestTool("tool2"); - - AgentConfig subAgent = AgentConfig.builder().name("subAgent").description("Sub agent").build(); - - GenerationConfig genConfig = GenerationConfig.builder().temperature(0.7).build(); - - OutputConfig outputConfig = new OutputConfig(); - outputConfig.setFormat(OutputFormat.JSON); - - AgentConfig config = - AgentConfig.builder() - .name("mainAgent") - .description("Main agent description") - .system("You are a helpful assistant.") - .model("openai/gpt-4o") - .tools(List.of(tool1, tool2)) - .agents(List.of(subAgent)) - .config(genConfig) - .output(outputConfig) - .build(); - - assertEquals("mainAgent", config.getName()); - assertEquals("Main agent description", config.getDescription()); - assertEquals("You are a helpful assistant.", config.getSystem()); - assertEquals("openai/gpt-4o", config.getModel()); - assertEquals(2, config.getTools().size()); - assertEquals(1, config.getAgents().size()); - assertEquals("subAgent", config.getAgents().get(0).getName()); - assertEquals(genConfig, config.getConfig()); - assertEquals(outputConfig, config.getOutput()); - } - - @Test - void testBuilderWithMinimalFields() { - AgentConfig config = - AgentConfig.builder().name("simpleAgent").description("A simple agent").build(); - - assertEquals("simpleAgent", config.getName()); - assertEquals("A simple agent", config.getDescription()); - assertNull(config.getSystem()); - assertNull(config.getModel()); - assertNull(config.getTools()); - assertNull(config.getAgents()); - assertNull(config.getConfig()); - assertNull(config.getOutput()); - } - - @Test - void testDefaultConstructor() { - AgentConfig config = new AgentConfig(); - - assertNull(config.getName()); - assertNull(config.getDescription()); - assertNull(config.getSystem()); - assertNull(config.getModel()); - assertNull(config.getTools()); - assertNull(config.getAgents()); - assertNull(config.getConfig()); - assertNull(config.getOutput()); - } - - @Test - void testSetters() { - AgentConfig config = new AgentConfig(); - Tool tool = createTestTool("testTool"); - AgentConfig subAgent = AgentConfig.builder().name("sub").build(); - GenerationConfig genConfig = GenerationConfig.builder().build(); - OutputConfig outputConfig = new OutputConfig(); - - config.setName("agent"); - config.setDescription("desc"); - config.setSystem("system"); - config.setModel("model"); - config.setTools(List.of(tool)); - config.setAgents(List.of(subAgent)); - config.setConfig(genConfig); - config.setOutput(outputConfig); - - assertEquals("agent", config.getName()); - assertEquals("desc", config.getDescription()); - assertEquals("system", config.getSystem()); - assertEquals("model", config.getModel()); - assertEquals(1, config.getTools().size()); - assertEquals(1, config.getAgents().size()); - assertEquals(genConfig, config.getConfig()); - assertEquals(outputConfig, config.getOutput()); - } - - @Test - void testNestedAgents() { - AgentConfig level3 = AgentConfig.builder().name("level3").description("Level 3 agent").build(); - - AgentConfig level2 = - AgentConfig.builder() - .name("level2") - .description("Level 2 agent") - .agents(List.of(level3)) - .build(); - - AgentConfig level1 = - AgentConfig.builder() - .name("level1") - .description("Level 1 agent") - .agents(List.of(level2)) - .build(); - - assertEquals("level1", level1.getName()); - assertEquals(1, level1.getAgents().size()); - assertEquals("level2", level1.getAgents().get(0).getName()); - assertEquals(1, level1.getAgents().get(0).getAgents().size()); - assertEquals("level3", level1.getAgents().get(0).getAgents().get(0).getName()); - } - - @Test - void testMultipleToolsAndAgents() { - Tool tool1 = createTestTool("tool1"); - Tool tool2 = createTestTool("tool2"); - Tool tool3 = createTestTool("tool3"); - - AgentConfig sub1 = AgentConfig.builder().name("sub1").build(); - AgentConfig sub2 = AgentConfig.builder().name("sub2").build(); - - AgentConfig config = - AgentConfig.builder() - .name("main") - .tools(List.of(tool1, tool2, tool3)) - .agents(List.of(sub1, sub2)) - .build(); - - assertEquals(3, config.getTools().size()); - assertEquals(2, config.getAgents().size()); - } -} diff --git a/ai/src/test/java/com/google/genkit/ai/AgentTest.java b/ai/src/test/java/com/google/genkit/ai/AgentTest.java deleted file mode 100644 index f0e31618e..000000000 --- a/ai/src/test/java/com/google/genkit/ai/AgentTest.java +++ /dev/null @@ -1,185 +0,0 @@ -/* - * 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. - * - * SPDX-License-Identifier: Apache-2.0 - */ - -package com.google.genkit.ai; - -import static org.junit.jupiter.api.Assertions.*; - -import java.util.HashMap; -import java.util.List; -import java.util.Map; -import org.junit.jupiter.api.Test; - -/** Unit tests for Agent. */ -class AgentTest { - - /** Helper to create a simple test tool. */ - private Tool createTestTool(String name) { - Map schema = new HashMap<>(); - schema.put("type", "string"); - return new Tool<>( - name, "Test tool " + name, schema, schema, String.class, (ctx, input) -> "result"); - } - - @Test - void testAgentCreation() { - AgentConfig config = - AgentConfig.builder() - .name("testAgent") - .description("Test agent description") - .system("You are a test agent.") - .model("test-model") - .build(); - - Agent agent = new Agent(config); - - assertEquals("testAgent", agent.getName()); - assertEquals("Test agent description", agent.getDescription()); - assertEquals("You are a test agent.", agent.getSystem()); - assertEquals("test-model", agent.getModel()); - assertEquals(config, agent.getConfig()); - } - - @Test - void testAsTool() { - AgentConfig config = - AgentConfig.builder().name("delegateAgent").description("Handles delegated tasks").build(); - - Agent agent = new Agent(config); - Tool, Agent.AgentTransferResult> tool = agent.asTool(); - - assertNotNull(tool); - assertEquals("delegateAgent", tool.getDefinition().getName()); - assertEquals("Handles delegated tasks", tool.getDefinition().getDescription()); - assertNotNull(tool.getInputSchema()); - assertNotNull(tool.getOutputSchema()); - } - - @Test - void testGetToolDefinition() { - AgentConfig config = AgentConfig.builder().name("myAgent").description("My agent").build(); - - Agent agent = new Agent(config); - ToolDefinition def = agent.getToolDefinition(); - - assertEquals("myAgent", def.getName()); - assertEquals("My agent", def.getDescription()); - assertNotNull(def.getInputSchema()); - } - - @Test - void testGetTools() { - Tool tool1 = createTestTool("tool1"); - Tool tool2 = createTestTool("tool2"); - - AgentConfig config = AgentConfig.builder().name("agent").tools(List.of(tool1, tool2)).build(); - - Agent agent = new Agent(config); - - assertEquals(2, agent.getTools().size()); - assertTrue(agent.getTools().contains(tool1)); - assertTrue(agent.getTools().contains(tool2)); - } - - @Test - void testGetAgents() { - AgentConfig sub1 = AgentConfig.builder().name("sub1").build(); - AgentConfig sub2 = AgentConfig.builder().name("sub2").build(); - - AgentConfig config = AgentConfig.builder().name("parent").agents(List.of(sub1, sub2)).build(); - - Agent agent = new Agent(config); - - assertEquals(2, agent.getAgents().size()); - } - - @Test - void testGetAllToolsWithNoSubAgents() { - Tool tool1 = createTestTool("tool1"); - Tool tool2 = createTestTool("tool2"); - - AgentConfig config = AgentConfig.builder().name("agent").tools(List.of(tool1, tool2)).build(); - - Agent agent = new Agent(config); - Map registry = new HashMap<>(); - - List> allTools = agent.getAllTools(registry); - - assertEquals(2, allTools.size()); - } - - @Test - void testGetAllToolsWithSubAgents() { - Tool parentTool = createTestTool("parentTool"); - - AgentConfig subConfig = AgentConfig.builder().name("subAgent").description("Sub agent").build(); - - AgentConfig config = - AgentConfig.builder() - .name("parent") - .tools(List.of(parentTool)) - .agents(List.of(subConfig)) - .build(); - - Agent parent = new Agent(config); - Agent subAgent = new Agent(subConfig); - - Map registry = new HashMap<>(); - registry.put("subAgent", subAgent); - - List> allTools = parent.getAllTools(registry); - - // Should have parent tool + sub-agent as tool - assertEquals(2, allTools.size()); - } - - @Test - void testAgentTransferResult() { - Agent.AgentTransferResult result = new Agent.AgentTransferResult("targetAgent"); - - assertEquals("targetAgent", result.getTransferredTo()); - assertTrue(result.isTransferred()); - assertTrue(result.toString().contains("targetAgent")); - } - - @Test - void testToString() { - AgentConfig config = AgentConfig.builder().name("myAgent").build(); - - Agent agent = new Agent(config); - String str = agent.toString(); - - assertTrue(str.contains("myAgent")); - assertTrue(str.contains("Agent")); - } - - @Test - void testNullToolsAndAgents() { - AgentConfig config = AgentConfig.builder().name("minimal").build(); - - Agent agent = new Agent(config); - - assertNull(agent.getTools()); - assertNull(agent.getAgents()); - - Map registry = new HashMap<>(); - List> allTools = agent.getAllTools(registry); - - assertTrue(allTools.isEmpty()); - } -} diff --git a/ai/src/test/java/com/google/genkit/ai/agent/AgentChatTest.java b/ai/src/test/java/com/google/genkit/ai/agent/AgentChatTest.java new file mode 100644 index 000000000..7d8beb543 --- /dev/null +++ b/ai/src/test/java/com/google/genkit/ai/agent/AgentChatTest.java @@ -0,0 +1,596 @@ +/* + * 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. + * + * SPDX-License-Identifier: Apache-2.0 + */ + +package com.google.genkit.ai.agent; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertNotEquals; +import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertNull; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import com.fasterxml.jackson.databind.JsonNode; +import com.google.genkit.ai.Message; +import com.google.genkit.ai.ModelResponseChunk; +import com.google.genkit.ai.Part; +import com.google.genkit.ai.Role; +import com.google.genkit.ai.ToolRequest; +import com.google.genkit.ai.agent.internal.AgentActions; +import com.google.genkit.core.ActionContext; +import com.google.genkit.core.BufferedInputSource; +import com.google.genkit.core.DefaultRegistry; +import com.google.genkit.core.JsonUtils; +import com.google.genkit.core.Registry; +import java.util.ArrayList; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.TimeUnit; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; + +/** TDD tests for the {@link AgentChat} programmatic client (Tasks 5.1 + 5.2). */ +class AgentChatTest { + + private Registry registry; + private ActionContext ctx; + + @BeforeEach + void setUp() { + registry = new DefaultRegistry(); + ctx = new ActionContext(registry); + } + + // ── helpers ────────────────────────────────────────────────────────────────── + + /** An AgentFn that echoes the latest user text back as an assistant message. */ + private static AgentFn> echoFn() { + return (sess, fnCtx) -> { + String userText = latestUserText(sess.getMessages()); + return AgentResult.builder() + .message(Message.model("echo: " + userText)) + .finishReason(AgentFinishReason.STOP) + .build(); + }; + } + + /** + * An AgentFn that maintains a custom turn counter, streams a model chunk, and increments the + * counter (which produces a {@code customPatch} stream chunk). + */ + private static AgentFn> countingFn() { + return (sess, fnCtx) -> { + String userText = latestUserText(sess.getMessages()); + // Emit a model chunk so sendStream callbacks receive text. + fnCtx + .sendChunk() + .accept( + AgentStreamChunk.builder() + .modelChunk(ModelResponseChunk.text("chunk: " + userText)) + .build()); + // Mutate custom state -> produces a customPatch chunk. + sess.updateCustom( + cur -> { + Map next = cur != null ? new HashMap<>(cur) : new HashMap<>(); + int count = + next.get("count") instanceof Number ? ((Number) next.get("count")).intValue() : 0; + next.put("count", count + 1); + return next; + }); + return AgentResult.builder() + .message(Message.model("echo: " + userText)) + .finishReason(AgentFinishReason.STOP) + .build(); + }; + } + + private static String latestUserText(List msgs) { + for (int i = msgs.size() - 1; i >= 0; i--) { + if (msgs.get(i).getRole() == Role.USER) { + return msgs.get(i).getText(); + } + } + return ""; + } + + private static Agent> serverAgent( + Registry registry, String name, AgentFn> fn) { + CustomAgentConfig> config = + CustomAgentConfig.>builder() + .name(name) + .store(new InMemorySessionStore<>()) + .build(); + return AgentActions.defineCustomAgent(registry, config, fn); + } + + private static Agent> clientAgent( + Registry registry, String name, AgentFn> fn) { + CustomAgentConfig> config = + CustomAgentConfig.>builder().name(name).build(); + return AgentActions.defineCustomAgent(registry, config, fn); + } + + // ── server-managed multi-turn ───────────────────────────────────────────────── + + @Test + void testServerManagedMultiTurn() throws Exception { + Agent> agent = serverAgent(registry, "svr", echoFn()); + + AgentChat> chat = agent.chat(ctx); + + AgentResponse> r1 = chat.send("hi"); + assertEquals("echo: hi", r1.text()); + assertNotNull(r1.snapshotId()); + assertFalse(r1.snapshotId().isEmpty()); + String snap1 = chat.snapshotId(); + assertNotNull(snap1); + + AgentResponse> r2 = chat.send("again"); + assertEquals("echo: again", r2.text()); + // The second turn must have resumed the first: snapshot advanced. + assertNotNull(chat.snapshotId()); + assertFalse(chat.snapshotId().equals(snap1), "snapshot should advance across turns"); + + // History accumulated on the server: 2 user + 2 model messages. + SessionSnapshot> snap = + agent.getSnapshotData(GetSnapshotRequest.builder().snapshotId(chat.snapshotId()).build()); + assertNotNull(snap); + assertEquals(4, snap.getState().getMessages().size()); + } + + // ── client-managed multi-turn (state round-trip) ────────────────────────────── + + @Test + void testClientManagedMultiTurn() throws Exception { + Agent> agent = clientAgent(registry, "cli", echoFn()); + + AgentChat> chat = agent.chat(ctx); + + AgentResponse> r1 = chat.send("one"); + assertEquals("echo: one", r1.text()); + // After turn 1: 1 user + 1 model message. + assertEquals(2, chat.messages().size()); + + AgentResponse> r2 = chat.send("two"); + assertEquals("echo: two", r2.text()); + // Turn 2 carried the prior state forward: 2 user + 2 model messages. + assertEquals(4, chat.messages().size()); + assertNotNull(r2.state()); + assertEquals(4, r2.state().getMessages().size()); + } + + // ── sendStream delivers chunks ──────────────────────────────────────────────── + + @Test + void testSendStreamDeliversChunks() throws Exception { + Agent> agent = serverAgent(registry, "stream", countingFn()); + + AgentChat> chat = agent.chat(ctx); + + List>> chunks = new ArrayList<>(); + AgentResponse> resp = chat.sendStream("hello", chunks::add); + + assertEquals("echo: hello", resp.text()); + // At least the model chunk should have been delivered. + boolean sawText = chunks.stream().anyMatch(c -> "chunk: hello".equals(c.text())); + assertTrue(sawText, "expected a model chunk with text"); + } + + // ── custom state via chat.state() and chunk.custom() ────────────────────────── + + @Test + void testCustomStateAndChunkCustom() throws Exception { + Agent> agent = serverAgent(registry, "custom", countingFn()); + + AgentChat> chat = agent.chat(ctx); + + List>> chunks = new ArrayList<>(); + chat.sendStream("first", chunks::add); + + // chat.state() reflects custom state after the turn. + assertNotNull(chat.state()); + assertEquals(1, ((Number) chat.state().get("count")).intValue()); + + // A customPatch chunk should have made chunk.custom() reflect the post-patch state. + AgentChunk> patchChunk = + chunks.stream().filter(c -> c.custom() != null).reduce((a, b) -> b).orElse(null); + assertNotNull(patchChunk, "expected a chunk carrying custom state"); + assertEquals(1, ((Number) patchChunk.custom().get("count")).intValue()); + + // Second turn increments again. + chat.send("second"); + assertEquals(2, ((Number) chat.state().get("count")).intValue()); + } + + // ── client-managed custom state round-trip ──────────────────────────────────── + + @Test + void testClientManagedCustomStateRoundTrip() throws Exception { + Agent> agent = clientAgent(registry, "cliCustom", countingFn()); + + AgentChat> chat = agent.chat(ctx); + chat.send("a"); + assertEquals(1, ((Number) chat.state().get("count")).intValue()); + chat.send("b"); + assertEquals(2, ((Number) chat.state().get("count")).intValue()); + } + + // ── loadChat resumes a snapshot ─────────────────────────────────────────────── + + @Test + void testLoadChatResumesSnapshot() throws Exception { + Agent> agent = serverAgent(registry, "load", echoFn()); + + AgentChat> chat = agent.chat(ctx); + chat.send("hi"); + String snapshotId = chat.snapshotId(); + String sessionId = chat.sessionId(); + assertNotNull(snapshotId); + + // Hydrate a fresh chat from the snapshot. + AgentChat> loaded = + agent.loadChat(ctx, GetSnapshotRequest.builder().snapshotId(snapshotId).build()); + assertEquals(snapshotId, loaded.snapshotId()); + assertEquals(sessionId, loaded.sessionId()); + + // Its next send resumes that snapshot: history grows to 4 messages. + loaded.send("again"); + SessionSnapshot> snap = + agent.getSnapshotData(GetSnapshotRequest.builder().snapshotId(loaded.snapshotId()).build()); + assertEquals(4, snap.getState().getMessages().size()); + } + + // ── behavior 1: server-managed response.state() is null, snapshotId present ─── + + @Test + void testServerManagedResponseStateIsNull() throws Exception { + Agent> agent = serverAgent(registry, "svrState", echoFn()); + AgentChat> chat = agent.chat(ctx); + + AgentResponse> resp = chat.send("hi"); + + assertNull(resp.state(), "server-managed agents must not echo full state inline"); + assertNotNull(resp.snapshotId()); + assertFalse(resp.snapshotId().isEmpty()); + } + + // ── behavior 2: client-managed response.state() is non-null, no snapshotId ─── + + @Test + void testClientManagedResponseStateIsNonNull() throws Exception { + Agent> agent = clientAgent(registry, "cliState", echoFn()); + AgentChat> chat = agent.chat(ctx); + + AgentResponse> resp = chat.send("hi"); + + assertNotNull(resp.state(), "client-managed agents must echo full inline state"); + assertTrue(resp.snapshotId() == null || resp.snapshotId().isEmpty()); + } + + // ── behavior 3: interrupted turn surfaces a non-empty interrupts() list ────── + + @Test + void testInterruptedResponseHasInterruptsList() throws Exception { + AgentFn> interruptingFn = + (sess, fnCtx) -> { + Part toolRequestPart = new Part(); + ToolRequest tr = new ToolRequest(); + tr.setName("confirmAction"); + tr.setRef("ref-1"); + toolRequestPart.setToolRequest(tr); + Message msg = new Message(Role.MODEL, List.of(toolRequestPart)); + return AgentResult.builder() + .message(msg) + .finishReason(AgentFinishReason.INTERRUPTED) + .build(); + }; + Agent> agent = serverAgent(registry, "interrupt", interruptingFn); + AgentChat> chat = agent.chat(ctx); + + AgentResponse> resp = chat.send("please confirm"); + + assertEquals(AgentFinishReason.INTERRUPTED, resp.finishReason()); + assertFalse(resp.interrupts().isEmpty(), "expected a non-empty interrupts() list"); + assertEquals("confirmAction", resp.interrupts().get(0).name()); + } + + // ── behavior 4: resume() continues and yields a non-interrupted response ──── + + @Test + void testResumeAfterInterrupt() throws Exception { + // First invocation (session has only the 1 user message so far) interrupts; every + // subsequent invocation (history already carries that message plus the interrupted + // assistant reply) stops normally. turnIndex() cannot be used here because each top-level + // send()/resume() call resolves a fresh SessionRunner (turnIndex always starts at 0). + AgentFn> fn = + (sess, fnCtx) -> { + if (sess.getMessages().size() <= 1) { + Part toolRequestPart = new Part(); + ToolRequest tr = new ToolRequest(); + tr.setName("confirmAction"); + toolRequestPart.setToolRequest(tr); + Message msg = new Message(Role.MODEL, List.of(toolRequestPart)); + return AgentResult.builder() + .message(msg) + .finishReason(AgentFinishReason.INTERRUPTED) + .build(); + } + return AgentResult.builder() + .message(Message.model("resumed-ok")) + .finishReason(AgentFinishReason.STOP) + .build(); + }; + Agent> agent = serverAgent(registry, "resumeFlow", fn); + AgentChat> chat = agent.chat(ctx); + + AgentResponse> interrupted = chat.send("please confirm"); + assertEquals(AgentFinishReason.INTERRUPTED, interrupted.finishReason()); + assertFalse(interrupted.interrupts().isEmpty()); + + Part respondPart = new Part(); + respondPart.setText("confirmed"); + AgentResponse> resumed = chat.resume(List.of(respondPart)); + + assertNotEquals(AgentFinishReason.INTERRUPTED, resumed.finishReason()); + assertEquals(AgentFinishReason.STOP, resumed.finishReason()); + assertEquals("resumed-ok", resumed.text()); + } + + // ── behavior 5: sendStream delivers a chunk carrying turnEnd ────────────────── + + @Test + void testSendStreamDeliversTurnEndChunk() throws Exception { + Agent> agent = serverAgent(registry, "turnEnd", echoFn()); + AgentChat> chat = agent.chat(ctx); + + List>> chunks = new ArrayList<>(); + chat.sendStream("hi", chunks::add); + + boolean sawTurnEnd = + chunks.stream().anyMatch(c -> c.raw() != null && c.raw().getTurnEnd() != null); + assertTrue(sawTurnEnd, "expected at least one chunk carrying a non-null turnEnd"); + } + + // ── behavior 6: sendStream delivers a customPatch-derived chunk.custom() ───── + + @Test + void testSendStreamDeliverCustomChunk() throws Exception { + Agent> agent = serverAgent(registry, "customChunk", countingFn()); + AgentChat> chat = agent.chat(ctx); + + List>> chunks = new ArrayList<>(); + chat.sendStream("hi", chunks::add); + + AgentChunk> withCustom = + chunks.stream().filter(c -> c.custom() != null).findFirst().orElse(null); + assertNotNull(withCustom, "expected a chunk with non-null custom() derived from a customPatch"); + assertEquals(1, ((Number) withCustom.custom().get("count")).intValue()); + } + + // ── behavior 7: ctx.sendChunk(modelChunk) is delivered to sendStream callback ─ + + @Test + void testCustomAgentSendChunkIsDeliveredToSendStream() throws Exception { + AgentFn> fn = + (sess, fnCtx) -> { + fnCtx + .sendChunk() + .accept( + AgentStreamChunk.builder() + .modelChunk(ModelResponseChunk.text("hello-from-agentfn")) + .build()); + return AgentResult.builder() + .message(Message.model("done")) + .finishReason(AgentFinishReason.STOP) + .build(); + }; + Agent> agent = serverAgent(registry, "sendChunk", fn); + AgentChat> chat = agent.chat(ctx); + + List>> chunks = new ArrayList<>(); + chat.sendStream("hi", chunks::add); + + boolean sawChunk = chunks.stream().anyMatch(c -> "hello-from-agentfn".equals(c.text())); + assertTrue(sawChunk, "expected the model chunk sent via ctx.sendChunk() to reach the callback"); + } + + // ── behavior 8: artifacts added by the AgentFn appear in the response ──────── + + @Test + void testArtifactsPresentInResponse() throws Exception { + AgentFn> fn = + (sess, fnCtx) -> { + sess.addArtifacts(Artifact.builder().name("report").parts(List.of()).build()); + return AgentResult.builder() + .message(Message.model("done")) + .finishReason(AgentFinishReason.STOP) + .build(); + }; + Agent> agent = serverAgent(registry, "artifacts", fn); + AgentChat> chat = agent.chat(ctx); + + AgentResponse> resp = chat.send("hi"); + + assertFalse( + resp.artifacts().isEmpty(), "expected the added artifact to be present on the response"); + assertEquals("report", resp.artifacts().get(0).getName()); + } + + // ── behavior 9: clientTransform redacts custom state returned to the caller ── + + @Test + void testClientTransformRedactsState() throws Exception { + ClientTransform> redact = + state -> { + if (state == null) { + return state; + } + Map custom = state.getCustom(); + if (custom != null) { + Map redacted = new HashMap<>(custom); + redacted.put("count", 0); + state.setCustom(redacted); + } + return state; + }; + CustomAgentConfig> config = + CustomAgentConfig.>builder() + .name("redact") + .clientTransform(redact) + .build(); + Agent> agent = + AgentActions.defineCustomAgent(registry, config, countingFn()); + + AgentChat> chat = agent.chat(ctx); + + AgentResponse> r1 = chat.send("a"); + // Without the clientTransform, countingFn's first turn would set count=1; the transform must + // rewrite the outgoing state so the caller observes the redacted value (0) instead. + assertEquals(0, ((Number) r1.state().getCustom().get("count")).intValue()); + assertEquals(0, ((Number) chat.state().get("count")).intValue()); + + // Client-managed mode round-trips the transformed state as the next turn's input, so the + // redaction is applied on every turn the caller observes. + AgentResponse> r2 = chat.send("b"); + assertEquals(0, ((Number) r2.state().getCustom().get("count")).intValue()); + } + + // ── behavior 11: a thrown AgentFn does not throw from chat.send(); FAILED result ─ + + @Test + void testFailedTurnPopulatesFinishReasonFailed() throws Exception { + AgentFn> throwingFn = + (sess, fnCtx) -> { + throw new RuntimeException("boom"); + }; + Agent> agent = serverAgent(registry, "failing", throwingFn); + AgentChat> chat = agent.chat(ctx); + + // chat.send() must not throw even though the AgentFn throws. + AgentResponse> resp = chat.send("hi"); + + assertEquals(AgentFinishReason.FAILED, resp.finishReason()); + assertNotNull(resp.raw().getError()); + assertEquals("boom", resp.raw().getError().getMessage()); + } + + // ── behavior 12: sendStream also surfaces a FAILED turnEnd for a thrown AgentFn ─ + + @Test + void testSendStreamDeliversFailedTurnEnd() throws Exception { + AgentFn> throwingFn = + (sess, fnCtx) -> { + throw new RuntimeException("stream-boom"); + }; + Agent> agent = serverAgent(registry, "failingStream", throwingFn); + AgentChat> chat = agent.chat(ctx); + + List>> chunks = new ArrayList<>(); + AgentResponse> resp = chat.sendStream("hi", chunks::add); + + assertEquals(AgentFinishReason.FAILED, resp.finishReason()); + boolean sawFailedTurnEnd = + chunks.stream() + .anyMatch( + c -> + c.raw() != null + && c.raw().getTurnEnd() != null + && c.raw().getTurnEnd().getFinishReason() == AgentFinishReason.FAILED); + assertTrue(sawFailedTurnEnd, "expected a turnEnd chunk with finishReason FAILED"); + } + + // ── behavior 15: loadChat resumes a detached turn after it completes ───────── + + private static JsonNode initJson(AgentInit> init) { + return JsonUtils.toJsonNode(init != null ? init : new AgentInit>()); + } + + private static BufferedInputSource inputSourceWith(AgentInput... inputs) { + BufferedInputSource src = new BufferedInputSource<>(); + for (AgentInput in : inputs) { + src.offer(JsonUtils.toJsonNode(in)); + } + src.end(); + return src; + } + + private SessionSnapshot> pollForStatus( + Agent> agent, String snapshotId, SnapshotStatus until) throws Exception { + long deadline = System.currentTimeMillis() + 5000; + SessionSnapshot> snap = null; + while (System.currentTimeMillis() < deadline) { + snap = agent.getSnapshotData(GetSnapshotRequest.builder().snapshotId(snapshotId).build()); + if (snap != null && snap.getStatus() == until) { + return snap; + } + Thread.sleep(20); + } + return snap; + } + + @Test + void testLoadChatAfterDetachedTurnCompletes() throws Exception { + InMemorySessionStore> store = new InMemorySessionStore<>(); + CustomAgentConfig> config = + CustomAgentConfig.>builder().name("detachedLoad").store(store).build(); + + CountDownLatch started = new CountDownLatch(1); + CountDownLatch gate = new CountDownLatch(1); + AgentFn> gatedFn = + (sess, fnCtx) -> { + started.countDown(); + if (!gate.await(5, TimeUnit.SECONDS)) { + throw new IllegalStateException("gate not released in time"); + } + String userText = latestUserText(sess.getMessages()); + return AgentResult.builder() + .message(Message.model("echo: " + userText)) + .finishReason(AgentFinishReason.STOP) + .build(); + }; + Agent> agent = AgentActions.defineCustomAgent(registry, config, gatedFn); + + AgentInput detachInput = + AgentInput.builder().message(Message.user("first")).detach(true).build(); + JsonNode out = agent.runBidiJson(ctx, initJson(null), inputSourceWith(detachInput), c -> {}); + String snapshotId = out.get("snapshotId").asText(); + assertTrue(started.await(5, TimeUnit.SECONDS), "background turn should have started"); + + gate.countDown(); + SessionSnapshot> completed = + pollForStatus(agent, snapshotId, SnapshotStatus.COMPLETED); + assertNotNull(completed); + assertEquals(SnapshotStatus.COMPLETED, completed.getStatus()); + + // loadChat resumes the detached-then-completed snapshot; a follow-up send extends history. + AgentChat> loaded = + agent.loadChat(ctx, GetSnapshotRequest.builder().snapshotId(snapshotId).build()); + assertEquals(snapshotId, loaded.snapshotId()); + assertEquals(2, loaded.messages().size(), "prior detached turn's messages must be present"); + + AgentResponse> resp = loaded.send("second"); + assertEquals("echo: second", resp.text()); + + SessionSnapshot> finalSnap = + agent.getSnapshotData(GetSnapshotRequest.builder().snapshotId(loaded.snapshotId()).build()); + assertNotNull(finalSnap); + assertEquals(4, finalSnap.getState().getMessages().size()); + } +} diff --git a/ai/src/test/java/com/google/genkit/ai/agent/CoreFixesTest.java b/ai/src/test/java/com/google/genkit/ai/agent/CoreFixesTest.java new file mode 100644 index 000000000..1b06e320d --- /dev/null +++ b/ai/src/test/java/com/google/genkit/ai/agent/CoreFixesTest.java @@ -0,0 +1,412 @@ +/* + * 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. + * + * SPDX-License-Identifier: Apache-2.0 + */ + +package com.google.genkit.ai.agent; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertSame; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import com.fasterxml.jackson.databind.JsonNode; +import com.google.genkit.ai.Message; +import com.google.genkit.ai.Part; +import com.google.genkit.ai.Role; +import com.google.genkit.ai.ToolRequest; +import com.google.genkit.ai.ToolResponse; +import com.google.genkit.ai.agent.internal.AgentActions; +import com.google.genkit.core.ActionContext; +import com.google.genkit.core.DefaultRegistry; +import com.google.genkit.core.Registry; +import java.util.List; +import java.util.Map; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicReference; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; + +/** + * Tests proving the 4 core-runtime fixes: + * + *

    + *
  1. Fix 1 (custom-AgentFn half): {@code ctx.resume()} carries the real {@code respond}/{@code + * restart} parts on a resume turn (see {@code GenkitBetaTest} / {@code + * AgentConformanceTest}'s "interrupt resume state accumulation" for the generate-backed + * half). + *
  2. Fix 2: {@code agent.abort(snapshotId)} flips {@code AgentFnContext.isAborted()} for a + * still-running DETACHED turn. + *
  3. Fix 3: {@code AgentSessionContext.current()} is bound to the real {@link Session} during a + * running turn. + *
  4. Fix 4: a custom agent's registered metadata includes a generated {@code stateSchema} for a + * non-trivial POJO state type. + *
+ */ +class CoreFixesTest { + + private Registry registry; + private ActionContext ctx; + + @BeforeEach + void setUp() { + registry = new DefaultRegistry(); + ctx = new ActionContext(registry); + } + + // ── Fix 1: ctx.resume() reaches a custom AgentFn ───────────────────────────── + + @Test + void customAgentFnObservesResumeRespondParts() throws Exception { + InMemorySessionStore> store = new InMemorySessionStore<>(); + CustomAgentConfig> config = + CustomAgentConfig.>builder().name("resumeFn").store(store).build(); + + AtomicReference observedResume = new AtomicReference<>(); + AgentFn> fn = + (sess, fnCtx) -> { + if (fnCtx.resume() == null) { + // Turn 1: model "interrupts". + Part toolRequestPart = new Part(); + ToolRequest tr = new ToolRequest(); + tr.setName("confirmAction"); + tr.setRef("ref-1"); + toolRequestPart.setToolRequest(tr); + return AgentResult.builder() + .message(new Message(Role.MODEL, List.of(toolRequestPart))) + .finishReason(AgentFinishReason.INTERRUPTED) + .build(); + } + // Turn 2 (resume): record what ctx.resume() actually carried. + observedResume.set(fnCtx.resume()); + return AgentResult.builder() + .message(Message.model("resumed")) + .finishReason(AgentFinishReason.STOP) + .build(); + }; + + Agent> agent = AgentActions.defineCustomAgent(registry, config, fn); + AgentChat> chat = agent.chat(ctx); + + chat.send("please confirm"); + + Part respondPart = new Part(); + respondPart.setToolResponse(new ToolResponse("ref-1", "confirmAction", Map.of("ok", true))); + chat.resume(List.of(respondPart)); + + ToolResume resume = observedResume.get(); + assertNotNull(resume, "expected the custom AgentFn to observe a non-null ctx.resume()"); + assertNotNull(resume.getRespond(), "expected resume.getRespond() to be populated"); + assertEquals(1, resume.getRespond().size()); + ToolResponse tresp = resume.getRespond().get(0).getToolResponse(); + assertNotNull(tresp, "expected the respond part to carry a real ToolResponse"); + assertEquals("confirmAction", tresp.getName()); + assertEquals("ref-1", tresp.getRef()); + assertEquals(Map.of("ok", true), tresp.getOutput()); + } + + @Test + void nonResumeTurnHasNullCtxResume() throws Exception { + CustomAgentConfig> config = + CustomAgentConfig.>builder().name("noResume").build(); + AtomicReference observed = new AtomicReference<>(); + AgentFn> fn = + (sess, fnCtx) -> { + observed.set(fnCtx.resume()); + return AgentResult.builder() + .message(Message.model("ok")) + .finishReason(AgentFinishReason.STOP) + .build(); + }; + Agent> agent = AgentActions.defineCustomAgent(registry, config, fn); + agent.chat(ctx).send("hi"); + + // observed.get() was actually set (fn ran); a plain (non-resume) turn must see ctx.resume() + // == null, not accidentally reuse a stale ToolResume. + assertEquals(null, observed.get()); + } + + // ── Fix 2: abort() signals a running DETACHED turn ─────────────────────────── + + @Test + void abortSignalsRunningDetachedTurn() throws Exception { + InMemorySessionStore> store = new InMemorySessionStore<>(); + CustomAgentConfig> config = + CustomAgentConfig.>builder().name("abortable").store(store).build(); + + CountDownLatch started = new CountDownLatch(1); + CountDownLatch fnFinished = new CountDownLatch(1); + AtomicBoolean observedAborted = new AtomicBoolean(false); + AgentFn> gatedFn = + (sess, fnCtx) -> { + started.countDown(); + try { + long deadline = System.currentTimeMillis() + 5000; + while (!fnCtx.isAborted() && System.currentTimeMillis() < deadline) { + Thread.sleep(10); + } + observedAborted.set(fnCtx.isAborted()); + return AgentResult.builder() + .message(Message.model(fnCtx.isAborted() ? "stopped" : "timed-out")) + .finishReason( + fnCtx.isAborted() ? AgentFinishReason.ABORTED : AgentFinishReason.STOP) + .build(); + } finally { + // Signal that THIS function has genuinely observed/reacted to the abort (or timed + // out), independent of when the store-level snapshot status flips: Agent.abort() + // flips the stored status synchronously in the calling thread, which can race ahead + // of this background thread noticing the AtomicBoolean flip on its next poll tick. + fnFinished.countDown(); + } + }; + + Agent> agent = AgentActions.defineCustomAgent(registry, config, gatedFn); + + com.fasterxml.jackson.databind.JsonNode initJson = + com.google.genkit.core.JsonUtils.toJsonNode(new AgentInit>()); + com.google.genkit.core.BufferedInputSource inputs = + new com.google.genkit.core.BufferedInputSource<>(); + AgentInput detachInput = AgentInput.builder().message(Message.user("go")).detach(true).build(); + inputs.offer(com.google.genkit.core.JsonUtils.toJsonNode(detachInput)); + inputs.end(); + + JsonNode out = agent.runBidiJson(ctx, initJson, inputs, c -> {}); + String snapshotId = out.get("snapshotId").asText(); + assertNotNull(snapshotId); + assertFalse(snapshotId.isEmpty()); + + // The background turn is definitely running now (it counted down `started`), and is + // definitely still blocked in its poll loop (it only exits on abort or a 5s timeout). + assertTrue(started.await(5, TimeUnit.SECONDS), "background turn should have started"); + + SnapshotStatus statusAfterAbort = agent.abort(snapshotId); + assertEquals(SnapshotStatus.ABORTED, statusAfterAbort); + + // Wait for the AgentFn itself to finish reacting (NOT for the store's snapshot status, which + // Agent.abort() flips synchronously in the calling thread — that would race ahead of the + // background thread's next poll tick and prove nothing about whether the fn actually saw the + // signal). This is the crux of Fix 2: the running function's own isAborted() check must have + // observed true. + assertTrue(fnFinished.await(5, TimeUnit.SECONDS), "the gated AgentFn should have finished"); + assertTrue( + observedAborted.get(), + "the running AgentFn must have observed ctx.isAborted() == true and reacted to it"); + + // The store-level status must also reflect the abort once the background turn finalizes. + SessionSnapshot> finalSnap = pollForTerminal(agent, snapshotId, 5000); + assertNotNull(finalSnap); + assertEquals(SnapshotStatus.ABORTED, finalSnap.getStatus()); + } + + @Test + void abortOfUnknownSnapshotIdIsANoOpForRegistry() { + InMemorySessionStore> store = new InMemorySessionStore<>(); + CustomAgentConfig> config = + CustomAgentConfig.>builder().name("abortUnknown").store(store).build(); + AgentFn> fn = + (sess, fnCtx) -> + AgentResult.builder() + .message(Message.model("x")) + .finishReason(AgentFinishReason.STOP) + .build(); + Agent> agent = AgentActions.defineCustomAgent(registry, config, fn); + + // No such snapshot exists at all (never registered, never persisted): abort() must not throw. + SnapshotStatus result = agent.abort("does-not-exist"); + assertEquals(null, result); + } + + private SessionSnapshot> pollForTerminal( + Agent> agent, String snapshotId, long timeoutMs) + throws InterruptedException { + long deadline = System.currentTimeMillis() + timeoutMs; + SessionSnapshot> snap = null; + while (System.currentTimeMillis() < deadline) { + snap = agent.getSnapshotData(GetSnapshotRequest.builder().snapshotId(snapshotId).build()); + if (snap != null + && (snap.getStatus() == SnapshotStatus.COMPLETED + || snap.getStatus() == SnapshotStatus.FAILED + || snap.getStatus() == SnapshotStatus.ABORTED)) { + return snap; + } + Thread.sleep(20); + } + return snap; + } + + // ── Fix 3: AgentSessionContext is bound during a running turn ─────────────── + + @Test + void agentSessionContextIsBoundDuringForegroundTurn() throws Exception { + CustomAgentConfig> config = + CustomAgentConfig.>builder().name("ctxBound").build(); + + AtomicReference> observed = new AtomicReference<>(); + AgentFn> fn = + (sess, fnCtx) -> { + observed.set(AgentSessionContext.current()); + return AgentResult.builder() + .message(Message.model("ok")) + .finishReason(AgentFinishReason.STOP) + .build(); + }; + Agent> agent = AgentActions.defineCustomAgent(registry, config, fn); + + // Not bound before/after any turn runs. + assertEquals(null, AgentSessionContext.current()); + + agent.chat(ctx).send("hi"); + + assertNotNull(observed.get(), "expected AgentSessionContext.current() to be non-null mid-turn"); + + // Unbound again after the turn completes. + assertEquals(null, AgentSessionContext.current()); + } + + @Test + void agentSessionContextIsBoundDuringDetachedTurn() throws Exception { + InMemorySessionStore> store = new InMemorySessionStore<>(); + CustomAgentConfig> config = + CustomAgentConfig.>builder() + .name("ctxBoundDetached") + .store(store) + .build(); + + CountDownLatch done = new CountDownLatch(1); + AtomicReference> observed = new AtomicReference<>(); + AgentFn> fn = + (sess, fnCtx) -> { + observed.set(AgentSessionContext.current()); + done.countDown(); + return AgentResult.builder() + .message(Message.model("ok")) + .finishReason(AgentFinishReason.STOP) + .build(); + }; + Agent> agent = AgentActions.defineCustomAgent(registry, config, fn); + + JsonNode initJson = + com.google.genkit.core.JsonUtils.toJsonNode(new AgentInit>()); + com.google.genkit.core.BufferedInputSource inputs = + new com.google.genkit.core.BufferedInputSource<>(); + AgentInput detachInput = AgentInput.builder().message(Message.user("go")).detach(true).build(); + inputs.offer(com.google.genkit.core.JsonUtils.toJsonNode(detachInput)); + inputs.end(); + + agent.runBidiJson(ctx, initJson, inputs, c -> {}); + assertTrue(done.await(5, TimeUnit.SECONDS), "detached turn should have completed"); + + assertNotNull( + observed.get(), + "expected AgentSessionContext.current() to be non-null during detached turn"); + } + + // ── Fix 4: stateSchema populated for a custom agent with a POJO state type ── + + /** Non-trivial POJO state type with named fields, used to prove schema generation. */ + public static class ReviewState { + private String status; + private List reviewers; + + public String getStatus() { + return status; + } + + public void setStatus(String status) { + this.status = status; + } + + public List getReviewers() { + return reviewers; + } + + public void setReviewers(List reviewers) { + this.reviewers = reviewers; + } + } + + @Test + @SuppressWarnings("unchecked") + void customAgentMetadataIncludesStateSchemaForPojoStateType() { + CustomAgentConfig config = + CustomAgentConfig.builder() + .name("typedCustom") + .stateType(ReviewState.class) + .build(); + AgentFn fn = + (sess, fnCtx) -> + AgentResult.builder() + .message(Message.model("x")) + .finishReason(AgentFinishReason.STOP) + .build(); + + Agent agent = AgentActions.defineCustomAgent(registry, config, fn); + + Map metadata = agent.getMetadata(); + assertNotNull(metadata); + Map agentMeta = (Map) metadata.get("agent"); + assertNotNull(agentMeta, "expected an \"agent\" sub-map in metadata"); + + Object stateSchemaObj = agentMeta.get("stateSchema"); + assertNotNull(stateSchemaObj, "expected a generated stateSchema for the ReviewState POJO"); + Map stateSchema = (Map) stateSchemaObj; + Map properties = (Map) stateSchema.get("properties"); + assertNotNull(properties); + assertTrue(properties.containsKey("status")); + assertTrue(properties.containsKey("reviewers")); + } + + @Test + @SuppressWarnings("unchecked") + void customAgentMetadataOmitsStateSchemaForMapStateType() { + CustomAgentConfig> config = + CustomAgentConfig.>builder().name("mapStateAgent").build(); + AgentFn> fn = + (sess, fnCtx) -> + AgentResult.builder() + .message(Message.model("x")) + .finishReason(AgentFinishReason.STOP) + .build(); + + Agent> agent = AgentActions.defineCustomAgent(registry, config, fn); + + Map agentMeta = (Map) agent.getMetadata().get("agent"); + assertNotNull(agentMeta); + assertFalse( + agentMeta.containsKey("stateSchema"), + "no stateType was configured (defaults to a dynamic Map), so no stateSchema is expected"); + } + + @Test + void agentFnContextNewConstructorDoesNotBreakOldOnes() { + // Fix 1's new constructor overload must not have broken the pre-existing ones (used by other + // callers). Exercise all four directly. + AgentFnContext c1 = new AgentFnContext(chunk -> {}, new AtomicBoolean(false)); + assertFalse(c1.isAborted()); + assertEquals(null, c1.resume()); + + AgentFnContext c2 = new AgentFnContext(chunk -> {}, new AtomicBoolean(false), ctx); + assertSame(ctx, c2.context()); + assertEquals(null, c2.resume()); + + ToolResume tr = ToolResume.builder().build(); + AgentFnContext c3 = new AgentFnContext(chunk -> {}, new AtomicBoolean(false), ctx, tr); + assertSame(tr, c3.resume()); + } +} diff --git a/ai/src/test/java/com/google/genkit/ai/agent/DefineCustomAgentTest.java b/ai/src/test/java/com/google/genkit/ai/agent/DefineCustomAgentTest.java new file mode 100644 index 000000000..a008731e1 --- /dev/null +++ b/ai/src/test/java/com/google/genkit/ai/agent/DefineCustomAgentTest.java @@ -0,0 +1,416 @@ +/* + * 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. + * + * SPDX-License-Identifier: Apache-2.0 + */ + +package com.google.genkit.ai.agent; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertNotEquals; +import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertNull; +import static org.junit.jupiter.api.Assertions.assertSame; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import com.fasterxml.jackson.databind.JsonNode; +import com.google.genkit.ai.Message; +import com.google.genkit.ai.Role; +import com.google.genkit.ai.agent.internal.AgentActions; +import com.google.genkit.core.ActionContext; +import com.google.genkit.core.BufferedInputSource; +import com.google.genkit.core.DefaultRegistry; +import com.google.genkit.core.JsonUtils; +import com.google.genkit.core.Registry; +import java.time.Instant; +import java.util.ArrayList; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.atomic.AtomicInteger; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; + +/** TDD tests for {@code defineCustomAgent} (Task 4.5). */ +class DefineCustomAgentTest { + + private Registry registry; + private ActionContext ctx; + + @BeforeEach + void setUp() { + registry = new DefaultRegistry(); + ctx = new ActionContext(registry); + } + + // ── helpers ────────────────────────────────────────────────────────────────── + + /** An AgentFn that echoes the user's text back as an assistant message. */ + private static AgentFn> echoFn() { + return (sess, fnCtx) -> { + List msgs = sess.getMessages(); + String userText = ""; + for (int i = msgs.size() - 1; i >= 0; i--) { + if (msgs.get(i).getRole() == Role.USER) { + userText = msgs.get(i).getText(); + break; + } + } + return AgentResult.builder() + .message(Message.model("echo: " + userText)) + .finishReason(AgentFinishReason.STOP) + .build(); + }; + } + + private static JsonNode initJson(AgentInit> init) { + return JsonUtils.toJsonNode(init != null ? init : new AgentInit>()); + } + + private static BufferedInputSource inputSourceWith(AgentInput... inputs) { + BufferedInputSource src = new BufferedInputSource<>(); + for (AgentInput in : inputs) { + src.offer(JsonUtils.toJsonNode(in)); + } + src.end(); + return src; + } + + private static AgentInput userInput(String text) { + return AgentInput.builder().message(Message.user(text)).build(); + } + + // ── Server-managed single turn ─────────────────────────────────────────────── + + @Test + void testServerManagedSingleTurn() throws Exception { + InMemorySessionStore> store = new InMemorySessionStore<>(); + CustomAgentConfig> config = + CustomAgentConfig.>builder() + .name("svr") + .description("server agent") + .store(store) + .build(); + + Agent> agent = AgentActions.defineCustomAgent(registry, config, echoFn()); + + // Registered actions present + assertNotNull(registry.lookupAction("/agent/svr")); + assertNotNull(registry.lookupAction("/agent-snapshot/svr")); + assertNotNull(registry.lookupAction("/agent-abort/svr")); + + // metadata.agent reflects server state-management + abortable + JsonNode agentMeta = JsonUtils.toJsonNode(agent.getMetadata().get("agent")); + assertNotNull(agentMeta); + assertEquals("server", agentMeta.get("stateManagement").asText()); + assertTrue(agentMeta.get("abortable").asBoolean()); + + // Drive one turn + List chunks = new ArrayList<>(); + JsonNode out = + agent.runBidiJson(ctx, initJson(null), inputSourceWith(userInput("hi")), chunks::add); + + assertNotNull(out); + // server-managed: snapshotId present, no inline state + String snapshotId = out.get("snapshotId").asText(); + assertNotNull(snapshotId); + assertFalse(snapshotId.isEmpty()); + assertTrue(out.get("state") == null || out.get("state").isNull()); + // message echoed + assertEquals("echo: hi", out.get("message").get("content").get(0).get("text").asText()); + + // a TurnEnd chunk emitted + boolean sawTurnEnd = chunks.stream().anyMatch(c -> c.has("turnEnd")); + assertTrue(sawTurnEnd, "expected a turnEnd chunk"); + } + + // ── Client-managed ─────────────────────────────────────────────────────────── + + @Test + void testClientManaged() throws Exception { + CustomAgentConfig> config = + CustomAgentConfig.>builder().name("cli").build(); + + Agent> agent = AgentActions.defineCustomAgent(registry, config, echoFn()); + + // Only the agent action is registered + assertNotNull(registry.lookupAction("/agent/cli")); + assertNull(registry.lookupAction("/agent-snapshot/cli")); + assertNull(registry.lookupAction("/agent-abort/cli")); + + JsonNode agentMeta = JsonUtils.toJsonNode(agent.getMetadata().get("agent")); + assertEquals("client", agentMeta.get("stateManagement").asText()); + assertFalse(agentMeta.get("abortable").asBoolean()); + + JsonNode out = + agent.runBidiJson(ctx, initJson(null), inputSourceWith(userInput("yo")), c -> {}); + + // client-managed: inline state present, no snapshotId + assertNotNull(out.get("state")); + assertFalse(out.get("state").isNull()); + assertTrue(out.get("snapshotId") == null || out.get("snapshotId").isNull()); + // state contains messages (user + model) + JsonNode msgs = out.get("state").get("messages"); + assertEquals(2, msgs.size()); + } + + // ── Resume by sessionId (server) ───────────────────────────────────────────── + + @Test + void testResumeBySessionId() throws Exception { + InMemorySessionStore> store = new InMemorySessionStore<>(); + CustomAgentConfig> config = + CustomAgentConfig.>builder().name("res").store(store).build(); + Agent> agent = AgentActions.defineCustomAgent(registry, config, echoFn()); + + // Turn 1 + JsonNode out1 = + agent.runBidiJson(ctx, initJson(null), inputSourceWith(userInput("one")), c -> {}); + String sessionId = out1.get("sessionId").asText(); + assertNotNull(sessionId); + + // Turn 2: resume by sessionId + AgentInit> init2 = + AgentInit.>builder().sessionId(sessionId).build(); + JsonNode out2 = + agent.runBidiJson(ctx, initJson(init2), inputSourceWith(userInput("two")), c -> {}); + + String snapshotId = out2.get("snapshotId").asText(); + GetSnapshotRequest req = GetSnapshotRequest.builder().snapshotId(snapshotId).build(); + SessionSnapshot> snap = agent.getSnapshotData(req); + assertNotNull(snap); + // 2 user + 2 model messages accumulated + assertEquals(4, snap.getState().getMessages().size()); + } + + // ── Multi-turn in one invocation ───────────────────────────────────────────── + + @Test + void testMultiTurnSingleInvocation() throws Exception { + InMemorySessionStore> store = new InMemorySessionStore<>(); + CustomAgentConfig> config = + CustomAgentConfig.>builder().name("multi").store(store).build(); + Agent> agent = AgentActions.defineCustomAgent(registry, config, echoFn()); + + List chunks = new ArrayList<>(); + JsonNode out = + agent.runBidiJson( + ctx, + initJson(null), + inputSourceWith(userInput("first"), userInput("second")), + chunks::add); + + long turnEnds = chunks.stream().filter(c -> c.has("turnEnd")).count(); + assertEquals(2, turnEnds, "expected two turnEnd chunks"); + // Final output reflects the 2nd turn + assertEquals("echo: second", out.get("message").get("content").get(0).get("text").asText()); + + // 2 user + 2 model messages persisted + GetSnapshotRequest req = + GetSnapshotRequest.builder().sessionId(out.get("sessionId").asText()).build(); + SessionSnapshot> snap = agent.getSnapshotData(req); + assertEquals(4, snap.getState().getMessages().size()); + } + + // ── getSnapshot companion ──────────────────────────────────────────────────── + + @Test + void testGetSnapshotCompanion() throws Exception { + InMemorySessionStore> store = new InMemorySessionStore<>(); + CustomAgentConfig> config = + CustomAgentConfig.>builder().name("gs").store(store).build(); + Agent> agent = AgentActions.defineCustomAgent(registry, config, echoFn()); + + JsonNode out = + agent.runBidiJson(ctx, initJson(null), inputSourceWith(userInput("hi")), c -> {}); + String snapshotId = out.get("snapshotId").asText(); + + GetSnapshotRequest req = GetSnapshotRequest.builder().snapshotId(snapshotId).build(); + JsonNode snapJson = agent.getSnapshotDataAction().runJson(ctx, JsonUtils.toJsonNode(req), null); + assertNotNull(snapJson); + assertEquals(snapshotId, snapJson.get("snapshotId").asText()); + assertEquals("completed", snapJson.get("status").asText()); + } + + // ── abort companion ────────────────────────────────────────────────────────── + + @Test + void testAbortCompanion() throws Exception { + InMemorySessionStore> store = new InMemorySessionStore<>(); + CustomAgentConfig> config = + CustomAgentConfig.>builder().name("ab").store(store).build(); + Agent> agent = AgentActions.defineCustomAgent(registry, config, echoFn()); + + // Seed a PENDING snapshot directly in the store. + String now = Instant.now().toString(); + SessionState> state = + SessionState.>builder() + .sessionId("seeded-session") + .custom(new HashMap<>()) + .build(); + SessionSnapshot> pending = + SessionSnapshot.>builder() + .snapshotId("snap-pending") + .sessionId("seeded-session") + .createdAt(now) + .updatedAt(now) + .status(SnapshotStatus.PENDING) + .state(state) + .build(); + store.saveSnapshot("snap-pending", existing -> pending, SessionStoreOptions.empty()); + + // abort companion: PENDING -> ABORTED + AgentAbortRequest abortReq = AgentAbortRequest.builder().snapshotId("snap-pending").build(); + JsonNode resp1 = agent.abortAgentAction().runJson(ctx, JsonUtils.toJsonNode(abortReq), null); + assertEquals("aborted", resp1.get("status").asText()); + assertEquals("snap-pending", resp1.get("snapshotId").asText()); + + // calling again: terminal unchanged, still ABORTED + JsonNode resp2 = agent.abortAgentAction().runJson(ctx, JsonUtils.toJsonNode(abortReq), null); + assertEquals("aborted", resp2.get("status").asText()); + } + + // ── typed facades ──────────────────────────────────────────────────────────── + + @Test + void testTypedFacades() throws Exception { + InMemorySessionStore> store = new InMemorySessionStore<>(); + CustomAgentConfig> config = + CustomAgentConfig.>builder().name("facade").store(store).build(); + Agent> agent = AgentActions.defineCustomAgent(registry, config, echoFn()); + + JsonNode out = + agent.runBidiJson(ctx, initJson(null), inputSourceWith(userInput("hi")), c -> {}); + String snapshotId = out.get("snapshotId").asText(); + + // getSnapshotData typed facade + SessionSnapshot> snap = + agent.getSnapshotData(GetSnapshotRequest.builder().snapshotId(snapshotId).build()); + assertNotNull(snap); + assertEquals(snapshotId, snap.getSnapshotId()); + assertEquals(SnapshotStatus.COMPLETED, snap.getStatus()); + + // abort typed facade: COMPLETED is terminal -> stays COMPLETED + SnapshotStatus afterAbort = agent.abort(snapshotId); + assertEquals(SnapshotStatus.COMPLETED, afterAbort); + + // ref() + AgentRef ref = agent.ref(); + assertEquals("facade", ref.getName()); + + // store()/serverManaged() + assertSame(store, agent.store()); + assertTrue(agent.serverManaged()); + } + + // ── behavior 18: a user-defined SessionStore is genuinely exercised ───────── + + /** + * A minimal, independent {@link SessionStore} implementation backed by a {@link + * ConcurrentHashMap}, used to prove {@code defineCustomAgent} genuinely calls through to whatever + * store implementation the caller supplies (not just the built-in {@link InMemorySessionStore}). + */ + private static final class CountingSessionStore implements SessionStore> { + private final Map>> data = + new ConcurrentHashMap<>(); + private final AtomicInteger saveCount = new AtomicInteger(); + private final AtomicInteger getCount = new AtomicInteger(); + + @Override + public SessionSnapshot> getSnapshot(GetSnapshotOptions opts) { + getCount.incrementAndGet(); + if (opts.getSnapshotId() != null) { + return data.get(opts.getSnapshotId()); + } + if (opts.getSessionId() != null) { + // Single-snapshot-per-session in this minimal test double: return the newest. + return data.values().stream() + .filter(s -> opts.getSessionId().equals(s.getSessionId())) + .reduce((a, b) -> b) + .orElse(null); + } + return null; + } + + @Override + public String saveSnapshot( + String snapshotId, + SnapshotMutator> mutator, + SessionStoreOptions options) { + saveCount.incrementAndGet(); + SessionSnapshot> existing = + snapshotId != null ? data.get(snapshotId) : null; + SessionSnapshot> result = mutator.apply(existing); + if (result == null) { + return null; + } + String finalId = snapshotId != null ? snapshotId : java.util.UUID.randomUUID().toString(); + result.setSnapshotId(finalId); + if (result.getStatus() == null) { + result.setStatus(SnapshotStatus.COMPLETED); + } + data.put(finalId, result); + return finalId; + } + } + + @Test + void testCustomSessionStoreIsUsed() throws Exception { + CountingSessionStore customStore = new CountingSessionStore(); + CustomAgentConfig> config = + CustomAgentConfig.>builder() + .name("customStore") + .store(customStore) + .build(); + Agent> agent = AgentActions.defineCustomAgent(registry, config, echoFn()); + + // Turn 1: fresh session, no prior snapshot to read. + JsonNode out1 = + agent.runBidiJson(ctx, initJson(null), inputSourceWith(userInput("one")), c -> {}); + String sessionId = out1.get("sessionId").asText(); + String snapshotId1 = out1.get("snapshotId").asText(); + assertNotNull(sessionId); + assertNotNull(snapshotId1); + + assertTrue( + customStore.saveCount.get() > 0, "the custom store's saveSnapshot must have been called"); + assertEquals(1, customStore.saveCount.get()); + // The snapshot must genuinely be sitting in the custom store's own backing map. + assertTrue(customStore.data.containsKey(snapshotId1)); + + // Turn 2: resume by sessionId — this requires the custom store's getSnapshot to be called + // and to genuinely return the turn-1 snapshot for the runtime to resume from. + AgentInit> init2 = + AgentInit.>builder().sessionId(sessionId).build(); + JsonNode out2 = + agent.runBidiJson(ctx, initJson(init2), inputSourceWith(userInput("two")), c -> {}); + + assertTrue( + customStore.getCount.get() > 0, "the custom store's getSnapshot must have been called"); + assertEquals(2, customStore.saveCount.get(), "turn 2 must have saved a second snapshot"); + + String snapshotId2 = out2.get("snapshotId").asText(); + assertNotEquals(snapshotId1, snapshotId2); + + // Resume genuinely worked: history accumulated across both turns (2 user + 2 model msgs), + // sourced entirely from the custom store's own data. + SessionSnapshot> snap2 = customStore.data.get(snapshotId2); + assertNotNull(snap2); + assertEquals(4, snap2.getState().getMessages().size()); + assertEquals("echo: two", snap2.getState().getMessages().get(3).getContent().get(0).getText()); + } +} diff --git a/ai/src/test/java/com/google/genkit/ai/agent/DetachTest.java b/ai/src/test/java/com/google/genkit/ai/agent/DetachTest.java new file mode 100644 index 000000000..32846e555 --- /dev/null +++ b/ai/src/test/java/com/google/genkit/ai/agent/DetachTest.java @@ -0,0 +1,318 @@ +/* + * 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. + * + * SPDX-License-Identifier: Apache-2.0 + */ + +package com.google.genkit.ai.agent; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import com.fasterxml.jackson.databind.JsonNode; +import com.google.genkit.ai.Message; +import com.google.genkit.ai.Role; +import com.google.genkit.ai.agent.internal.AgentActions; +import com.google.genkit.ai.agent.internal.DetachController; +import com.google.genkit.core.ActionContext; +import com.google.genkit.core.BufferedInputSource; +import com.google.genkit.core.DefaultRegistry; +import com.google.genkit.core.JsonUtils; +import com.google.genkit.core.Registry; +import java.time.Instant; +import java.util.ArrayList; +import java.util.List; +import java.util.Map; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.TimeUnit; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; + +/** TDD tests for detach + heartbeat + background-finalize (Task 4.5b). */ +class DetachTest { + + private Registry registry; + private ActionContext ctx; + private long prevHeartbeat; + + @BeforeEach + void setUp() { + registry = new DefaultRegistry(); + ctx = new ActionContext(registry); + // Make heartbeat fire fast so the heartbeat test is deterministic. + prevHeartbeat = DetachController.setHeartbeatIntervalMillisForTest(50L); + } + + @AfterEach + void tearDown() { + DetachController.setHeartbeatIntervalMillisForTest(prevHeartbeat); + } + + // ── helpers ────────────────────────────────────────────────────────────────── + + /** + * An AgentFn that blocks on {@code gate} before echoing, and counts down {@code started} once it + * begins. Lets a test observe the PENDING phase, then release the turn to observe COMPLETED. + */ + private static AgentFn> gatedEchoFn( + CountDownLatch started, CountDownLatch gate) { + return (sess, fnCtx) -> { + started.countDown(); + if (!gate.await(5, TimeUnit.SECONDS)) { + throw new IllegalStateException("gate not released in time"); + } + List msgs = sess.getMessages(); + String userText = ""; + for (int i = msgs.size() - 1; i >= 0; i--) { + if (msgs.get(i).getRole() == Role.USER) { + userText = msgs.get(i).getText(); + break; + } + } + return AgentResult.builder() + .message(Message.model("echo: " + userText)) + .finishReason(AgentFinishReason.STOP) + .build(); + }; + } + + private static JsonNode initJson(AgentInit> init) { + return JsonUtils.toJsonNode(init != null ? init : new AgentInit>()); + } + + private static BufferedInputSource inputSourceWith(AgentInput... inputs) { + BufferedInputSource src = new BufferedInputSource<>(); + for (AgentInput in : inputs) { + src.offer(JsonUtils.toJsonNode(in)); + } + src.end(); + return src; + } + + private static AgentInput detachInput(String text) { + return AgentInput.builder().message(Message.user(text)).detach(true).build(); + } + + private SessionSnapshot> poll( + Agent> agent, String snapshotId, SnapshotStatus until) throws Exception { + long deadline = System.currentTimeMillis() + 5000; + SessionSnapshot> snap = null; + while (System.currentTimeMillis() < deadline) { + snap = agent.getSnapshotData(GetSnapshotRequest.builder().snapshotId(snapshotId).build()); + if (snap != null && snap.getStatus() == until) { + return snap; + } + Thread.sleep(20); + } + return snap; + } + + // ── detach returns DETACHED + pending snapshot, then finalizes to COMPLETED ─── + + @Test + void testDetachReturnsDetachedThenFinalizes() throws Exception { + InMemorySessionStore> store = new InMemorySessionStore<>(); + CustomAgentConfig> config = + CustomAgentConfig.>builder().name("dtc").store(store).build(); + + CountDownLatch started = new CountDownLatch(1); + CountDownLatch gate = new CountDownLatch(1); + Agent> agent = + AgentActions.defineCustomAgent(registry, config, gatedEchoFn(started, gate)); + + List chunks = new ArrayList<>(); + JsonNode out = + agent.runBidiJson(ctx, initJson(null), inputSourceWith(detachInput("hi")), chunks::add); + + // Immediate return: finishReason DETACHED + snapshotId, no inline state. + assertNotNull(out); + assertEquals("detached", out.get("finishReason").asText()); + String snapshotId = out.get("snapshotId").asText(); + assertNotNull(snapshotId); + assertFalse(snapshotId.isEmpty()); + + // The background turn has started but is gated → snapshot is PENDING. + assertTrue(started.await(5, TimeUnit.SECONDS), "background turn should have started"); + SessionSnapshot> pending = + agent.getSnapshotData(GetSnapshotRequest.builder().snapshotId(snapshotId).build()); + assertNotNull(pending); + assertEquals(SnapshotStatus.PENDING, pending.getStatus()); + + // No stream chunks for a detached run (streaming suppressed). + assertTrue(chunks.isEmpty(), "detached run must not emit stream chunks"); + + // Release the gate → background finalizes to COMPLETED with cumulative state. + gate.countDown(); + SessionSnapshot> done = poll(agent, snapshotId, SnapshotStatus.COMPLETED); + assertNotNull(done); + assertEquals(SnapshotStatus.COMPLETED, done.getStatus()); + assertEquals(AgentFinishReason.STOP, done.getFinishReason()); + // Cumulative state: user message + echoed model message. + assertNotNull(done.getState()); + assertEquals(2, done.getState().getMessages().size()); + assertEquals("echo: hi", done.getState().getMessages().get(1).getContent().get(0).getText()); + // Heartbeat cleared on terminal finalize. + assertTrue( + done.getHeartbeatAt() == null || done.getHeartbeatAt().isEmpty(), + "heartbeatAt should be cleared on finalize"); + } + + // ── heartbeat refreshes heartbeatAt while pending ───────────────────────────── + + @Test + void testHeartbeatRefreshesWhilePending() throws Exception { + InMemorySessionStore> store = new InMemorySessionStore<>(); + CustomAgentConfig> config = + CustomAgentConfig.>builder().name("hb").store(store).build(); + + CountDownLatch started = new CountDownLatch(1); + CountDownLatch gate = new CountDownLatch(1); + Agent> agent = + AgentActions.defineCustomAgent(registry, config, gatedEchoFn(started, gate)); + + JsonNode out = + agent.runBidiJson(ctx, initJson(null), inputSourceWith(detachInput("hi")), c -> {}); + String snapshotId = out.get("snapshotId").asText(); + assertTrue(started.await(5, TimeUnit.SECONDS)); + + // Observe the heartbeat advancing while pending (interval shrunk to 50ms in setUp). + String first = + agent + .getSnapshotData(GetSnapshotRequest.builder().snapshotId(snapshotId).build()) + .getHeartbeatAt(); + assertNotNull(first); + String later = first; + long deadline = System.currentTimeMillis() + 3000; + while (System.currentTimeMillis() < deadline) { + Thread.sleep(60); + later = + agent + .getSnapshotData(GetSnapshotRequest.builder().snapshotId(snapshotId).build()) + .getHeartbeatAt(); + if (later != null && !later.equals(first)) { + break; + } + } + assertNotNull(later); + assertTrue( + Instant.parse(later).compareTo(Instant.parse(first)) >= 0, + "heartbeatAt should advance while pending"); + assertFalse(later.equals(first), "heartbeatAt should be refreshed by the heartbeat task"); + + gate.countDown(); + poll(agent, snapshotId, SnapshotStatus.COMPLETED); + } + + // ── abort during pending: background finalize must NOT overwrite ABORTED ────── + + @Test + void testAbortDuringPendingNotOverwritten() throws Exception { + InMemorySessionStore> store = new InMemorySessionStore<>(); + CustomAgentConfig> config = + CustomAgentConfig.>builder().name("abt").store(store).build(); + + CountDownLatch started = new CountDownLatch(1); + CountDownLatch gate = new CountDownLatch(1); + Agent> agent = + AgentActions.defineCustomAgent(registry, config, gatedEchoFn(started, gate)); + + JsonNode out = + agent.runBidiJson(ctx, initJson(null), inputSourceWith(detachInput("hi")), c -> {}); + String snapshotId = out.get("snapshotId").asText(); + assertTrue(started.await(5, TimeUnit.SECONDS)); + + // Abort while still pending → ABORTED. + SnapshotStatus afterAbort = agent.abort(snapshotId); + assertEquals(SnapshotStatus.ABORTED, afterAbort); + + // Release the gate; the background finalize must NOT clobber the ABORTED row. + gate.countDown(); + Thread.sleep(500); + SessionSnapshot> snap = + agent.getSnapshotData(GetSnapshotRequest.builder().snapshotId(snapshotId).build()); + assertNotNull(snap); + assertEquals(SnapshotStatus.ABORTED, snap.getStatus()); + } + + // ── client-managed + detach: processed normally (graceful no-op for detach) ─── + + @Test + void testClientManagedDetachProcessedNormally() throws Exception { + CustomAgentConfig> config = + CustomAgentConfig.>builder().name("clidtc").build(); + CountDownLatch started = new CountDownLatch(1); + CountDownLatch gate = new CountDownLatch(0); // already open: no blocking for client-managed + Agent> agent = + AgentActions.defineCustomAgent(registry, config, gatedEchoFn(started, gate)); + + JsonNode out = + agent.runBidiJson(ctx, initJson(null), inputSourceWith(detachInput("yo")), c -> {}); + + // Client-managed: no store → detach N/A; processed normally with inline state, no DETACHED. + assertNotNull(out.get("state")); + assertFalse(out.get("state").isNull()); + assertFalse("detached".equals(out.get("finishReason").asText())); + } + + // ── behavior 20: a detached turn whose AgentFn throws finalizes to FAILED ───── + + @Test + void testDetachedTurnFailureTransitionsToFailed() throws Exception { + InMemorySessionStore> store = new InMemorySessionStore<>(); + CustomAgentConfig> config = + CustomAgentConfig.>builder().name("dtcFail").store(store).build(); + + CountDownLatch started = new CountDownLatch(1); + CountDownLatch gate = new CountDownLatch(1); + AgentFn> gatedThrowingFn = + (sess, fnCtx) -> { + started.countDown(); + if (!gate.await(5, TimeUnit.SECONDS)) { + throw new IllegalStateException("gate not released in time"); + } + throw new RuntimeException("detached-boom"); + }; + Agent> agent = + AgentActions.defineCustomAgent(registry, config, gatedThrowingFn); + + JsonNode out = + agent.runBidiJson(ctx, initJson(null), inputSourceWith(detachInput("hi")), c -> {}); + assertEquals("detached", out.get("finishReason").asText()); + String snapshotId = out.get("snapshotId").asText(); + assertTrue(started.await(5, TimeUnit.SECONDS), "background turn should have started"); + + SessionSnapshot> pending = + agent.getSnapshotData(GetSnapshotRequest.builder().snapshotId(snapshotId).build()); + assertNotNull(pending); + assertEquals(SnapshotStatus.PENDING, pending.getStatus()); + + // Release the gate → the AgentFn throws; the background finalize must transition to FAILED, + // never to COMPLETED. + gate.countDown(); + SessionSnapshot> failed = poll(agent, snapshotId, SnapshotStatus.FAILED); + assertNotNull(failed); + assertEquals(SnapshotStatus.FAILED, failed.getStatus()); + assertEquals(AgentFinishReason.FAILED, failed.getFinishReason()); + assertNotNull(failed.getError()); + assertEquals("detached-boom", failed.getError().getMessage()); + // Heartbeat cleared on terminal finalize, same as the successful path. + assertTrue( + failed.getHeartbeatAt() == null || failed.getHeartbeatAt().isEmpty(), + "heartbeatAt should be cleared on finalize"); + } +} diff --git a/ai/src/test/java/com/google/genkit/ai/agent/EnumSerdeTest.java b/ai/src/test/java/com/google/genkit/ai/agent/EnumSerdeTest.java new file mode 100644 index 000000000..0bf46ffed --- /dev/null +++ b/ai/src/test/java/com/google/genkit/ai/agent/EnumSerdeTest.java @@ -0,0 +1,119 @@ +/* + * 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. + * + * SPDX-License-Identifier: Apache-2.0 + */ + +package com.google.genkit.ai.agent; + +import static org.junit.jupiter.api.Assertions.*; + +import com.fasterxml.jackson.databind.ObjectMapper; +import com.google.genkit.core.JsonUtils; +import org.junit.jupiter.api.Test; + +/** Unit tests for SnapshotStatus and AgentFinishReason serialization/deserialization. */ +class EnumSerdeTest { + + private static final ObjectMapper objectMapper = JsonUtils.getObjectMapper(); + + // SnapshotStatus tests + + @Test + void testSnapshotStatusSerialize() throws Exception { + String json = objectMapper.writeValueAsString(SnapshotStatus.PENDING); + assertEquals("\"pending\"", json); + } + + @Test + void testSnapshotStatusDeserialize() throws Exception { + SnapshotStatus status = objectMapper.readValue("\"completed\"", SnapshotStatus.class); + assertEquals(SnapshotStatus.COMPLETED, status); + } + + @Test + void testSnapshotStatusRoundTrip() throws Exception { + SnapshotStatus[] values = { + SnapshotStatus.PENDING, + SnapshotStatus.COMPLETED, + SnapshotStatus.ABORTED, + SnapshotStatus.FAILED, + SnapshotStatus.EXPIRED + }; + for (SnapshotStatus status : values) { + String json = objectMapper.writeValueAsString(status); + SnapshotStatus deserialized = objectMapper.readValue(json, SnapshotStatus.class); + assertEquals(status, deserialized); + } + } + + @Test + void testSnapshotStatusFromValueOrCompletedWithNull() { + SnapshotStatus status = SnapshotStatus.fromValueOrCompleted(null); + assertEquals(SnapshotStatus.COMPLETED, status); + } + + @Test + void testSnapshotStatusFromValueOrCompletedWithEmpty() { + SnapshotStatus status = SnapshotStatus.fromValueOrCompleted(""); + assertEquals(SnapshotStatus.COMPLETED, status); + } + + @Test + void testSnapshotStatusFromValueOrCompletedWithAborted() { + SnapshotStatus status = SnapshotStatus.fromValueOrCompleted("aborted"); + assertEquals(SnapshotStatus.ABORTED, status); + } + + @Test + void testSnapshotStatusFromValueOrCompletedWithUnknown() { + assertThrows( + IllegalArgumentException.class, () -> SnapshotStatus.fromValueOrCompleted("unknown")); + } + + // AgentFinishReason tests + + @Test + void testAgentFinishReasonSerialize() throws Exception { + String json = objectMapper.writeValueAsString(AgentFinishReason.DETACHED); + assertEquals("\"detached\"", json); + } + + @Test + void testAgentFinishReasonDeserialize() throws Exception { + AgentFinishReason reason = objectMapper.readValue("\"interrupted\"", AgentFinishReason.class); + assertEquals(AgentFinishReason.INTERRUPTED, reason); + } + + @Test + void testAgentFinishReasonRoundTrip() throws Exception { + AgentFinishReason[] values = { + AgentFinishReason.STOP, + AgentFinishReason.LENGTH, + AgentFinishReason.BLOCKED, + AgentFinishReason.INTERRUPTED, + AgentFinishReason.OTHER, + AgentFinishReason.UNKNOWN, + AgentFinishReason.ABORTED, + AgentFinishReason.DETACHED, + AgentFinishReason.FAILED + }; + for (AgentFinishReason reason : values) { + String json = objectMapper.writeValueAsString(reason); + AgentFinishReason deserialized = objectMapper.readValue(json, AgentFinishReason.class); + assertEquals(reason, deserialized); + } + } +} diff --git a/ai/src/test/java/com/google/genkit/ai/agent/FileSessionStoreTest.java b/ai/src/test/java/com/google/genkit/ai/agent/FileSessionStoreTest.java new file mode 100644 index 000000000..e1df58165 --- /dev/null +++ b/ai/src/test/java/com/google/genkit/ai/agent/FileSessionStoreTest.java @@ -0,0 +1,401 @@ +/* + * 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. + * + * SPDX-License-Identifier: Apache-2.0 + */ + +package com.google.genkit.ai.agent; + +import static org.junit.jupiter.api.Assertions.*; + +import java.io.IOException; +import java.nio.file.Files; +import java.nio.file.Path; +import java.util.List; +import java.util.Map; +import java.util.concurrent.CopyOnWriteArrayList; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.TimeUnit; +import java.util.stream.Stream; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.io.TempDir; + +/** TDD tests for FileSessionStore (disk-backed session store). */ +class FileSessionStoreTest { + + @TempDir Path tempDir; + + private FileSessionStore> store; + + @BeforeEach + void setUp() { + store = new FileSessionStore<>(tempDir.toString()); + } + + // ── helpers ────────────────────────────────────────────────────────────────── + + private static SessionSnapshot> snapWithSession(String sessionId) { + return SessionSnapshot.>builder().sessionId(sessionId).build(); + } + + private static SessionSnapshot> snapWithSessionAndStatus( + String sessionId, SnapshotStatus status) { + SessionSnapshot> snap = + SessionSnapshot.>builder().sessionId(sessionId).build(); + snap.setStatus(status); + return snap; + } + + // ── file layout under global/ ───────────────────────────────────────────── + + @Test + void testSave_fileExistsUnderGlobalDir() throws IOException { + String id = + store.saveSnapshot( + null, existing -> snapWithSession("sess-file"), SessionStoreOptions.empty()); + + assertNotNull(id, "saveSnapshot must return a non-null id"); + Path snapshotFile = tempDir.resolve("global").resolve(id + ".json"); + assertTrue(Files.exists(snapshotFile), "snapshot file must exist under /global/.json"); + } + + @Test + void testGetSnapshot_bySnapshotId_roundTrip() { + String id = + store.saveSnapshot( + null, existing -> snapWithSession("sess-round-trip"), SessionStoreOptions.empty()); + + SessionSnapshot> fetched = + store.getSnapshot(GetSnapshotOptions.builder().snapshotId(id).build()); + + assertNotNull(fetched, "getSnapshot by id must return a snapshot"); + assertEquals(id, fetched.getSnapshotId()); + assertEquals("sess-round-trip", fetched.getSessionId()); + } + + // ── getSnapshot by sessionId: latest leaf via pointer + scan fallback ───── + + @Test + void testGetSnapshot_bySessionId_returnsLatestLeaf() { + String rootId = + store.saveSnapshot( + null, + existing -> { + SessionSnapshot> snap = snapWithSession("sess-chain"); + snap.setCreatedAt("2025-01-01T00:00:00Z"); + return snap; + }, + SessionStoreOptions.empty()); + + String childId = + store.saveSnapshot( + null, + existing -> { + SessionSnapshot> snap = snapWithSession("sess-chain"); + snap.setParentId(rootId); + snap.setCreatedAt("2025-01-02T00:00:00Z"); + return snap; + }, + SessionStoreOptions.empty()); + + SessionSnapshot> leaf = + store.getSnapshot(GetSnapshotOptions.builder().sessionId("sess-chain").build()); + assertNotNull(leaf, "getSnapshot by sessionId must return a snapshot"); + assertEquals(childId, leaf.getSnapshotId(), "must return the leaf (child)"); + } + + @Test + void testGetSnapshot_bySessionId_scanFallback_afterPointerDeleted() throws IOException { + String rootId = + store.saveSnapshot( + null, + existing -> { + SessionSnapshot> snap = snapWithSession("sess-scan"); + snap.setCreatedAt("2025-01-01T00:00:00Z"); + return snap; + }, + SessionStoreOptions.empty()); + + String childId = + store.saveSnapshot( + null, + existing -> { + SessionSnapshot> snap = snapWithSession("sess-scan"); + snap.setParentId(rootId); + snap.setCreatedAt("2025-01-02T00:00:00Z"); + return snap; + }, + SessionStoreOptions.empty()); + + // Delete the pointer file to force scan fallback + Path pointerFile = tempDir.resolve("global").resolve(".pointers").resolve("sess-scan.json"); + assertTrue(Files.exists(pointerFile), "pointer file must exist before deletion"); + Files.delete(pointerFile); + + // getSnapshot by sessionId should still work (scan fallback) and rewrite the pointer + SessionSnapshot> leaf = + store.getSnapshot(GetSnapshotOptions.builder().sessionId("sess-scan").build()); + assertNotNull(leaf, "scan fallback must find the snapshot"); + assertEquals(childId, leaf.getSnapshotId(), "scan fallback must return the leaf (child)"); + + // Pointer must be rewritten after scan + assertTrue(Files.exists(pointerFile), "pointer file must be rewritten after scan fallback"); + } + + // ── atomic write: no leftover .tmp files ───────────────────────────────── + + @Test + void testAtomicWrite_noTmpFilesLeft() throws IOException { + store.saveSnapshot( + null, existing -> snapWithSession("sess-atomic"), SessionStoreOptions.empty()); + + Path globalDir = tempDir.resolve("global"); + try (Stream files = Files.list(globalDir)) { + long tmpCount = files.filter(p -> p.toString().endsWith(".tmp")).count(); + assertEquals(0, tmpCount, "no .tmp files must remain after successful save"); + } + } + + // ── chain pruning keeps newest N ───────────────────────────────────────── + + @Test + void testChainPruning_keepsNewestN() throws IOException { + FileSessionStore> pruningStore = + FileSessionStore.>builder(tempDir.resolve("prunetest").toString()) + .maxPersistedChainLength(2) + .build(); + + // Save root + String id1 = + pruningStore.saveSnapshot( + null, + existing -> { + SessionSnapshot> snap = snapWithSession("sess-prune"); + snap.setCreatedAt("2025-01-01T00:00:00Z"); + return snap; + }, + SessionStoreOptions.empty()); + + // Save child1 of root + String id2 = + pruningStore.saveSnapshot( + null, + existing -> { + SessionSnapshot> snap = snapWithSession("sess-prune"); + snap.setParentId(id1); + snap.setCreatedAt("2025-01-02T00:00:00Z"); + return snap; + }, + SessionStoreOptions.empty()); + + // Save child2 of child1 — this should trigger pruning, dropping id1 + String id3 = + pruningStore.saveSnapshot( + null, + existing -> { + SessionSnapshot> snap = snapWithSession("sess-prune"); + snap.setParentId(id2); + snap.setCreatedAt("2025-01-03T00:00:00Z"); + return snap; + }, + SessionStoreOptions.empty()); + + Path pruneDir = tempDir.resolve("prunetest").resolve("global"); + // Count only .json files that are not in the .pointers sub-directory + try (Stream files = Files.list(pruneDir)) { + long snapshotCount = + files + .filter(p -> p.getFileName().toString().endsWith(".json")) + .filter(p -> !p.getParent().getFileName().toString().equals(".pointers")) + .count(); + assertEquals(2, snapshotCount, "only 2 snapshot files must remain after pruning (newest 2)"); + } + + // The oldest (id1) must be deleted + Path oldFile = pruneDir.resolve(id1 + ".json"); + assertFalse(Files.exists(oldFile), "oldest snapshot file must be pruned"); + + // The two newest must still exist + assertTrue(Files.exists(pruneDir.resolve(id2 + ".json")), "id2 must still exist"); + assertTrue(Files.exists(pruneDir.resolve(id3 + ".json")), "id3 must still exist"); + } + + // ── path safety: reject ../evil ────────────────────────────────────────── + + @Test + void testPathSafety_rejectsDotDotSnapshotId() { + assertThrows( + IllegalArgumentException.class, + () -> + store.saveSnapshot( + "../evil", existing -> snapWithSession("sess-evil"), SessionStoreOptions.empty()), + "saveSnapshot with '../evil' snapshotId must throw IllegalArgumentException"); + } + + @Test + void testPathSafety_rejectsSlashInSnapshotId() { + assertThrows( + IllegalArgumentException.class, + () -> + store.saveSnapshot( + "a/b", existing -> snapWithSession("sess-evil"), SessionStoreOptions.empty()), + "saveSnapshot with '/' in snapshotId must throw IllegalArgumentException"); + } + + @Test + void testPathSafety_rejectsDotDotGetSnapshot() { + assertThrows( + IllegalArgumentException.class, + () -> store.getSnapshot(GetSnapshotOptions.builder().snapshotId("../evil").build()), + "getSnapshot with '../evil' snapshotId must throw IllegalArgumentException"); + } + + @Test + void testPathSafety_rejectsDotDotInOnSnapshotStateChange() { + assertThrows( + IllegalArgumentException.class, + () -> store.onSnapshotStateChange("../evil", snap -> {}, SessionStoreOptions.empty()), + "onSnapshotStateChange with '../evil' snapshotId must throw IllegalArgumentException"); + } + + // ── subscriber fires on PENDING→ABORTED status change ───────────────────── + + @Test + void testSubscriber_firesOnStatusChange_pollingBased() throws Exception { + // Use a fast poll interval for testing + FileSessionStore> fastStore = + FileSessionStore.>builder(tempDir.resolve("subtest").toString()) + .snapshotWatchPollIntervalMs(100) + .build(); + + // Save a PENDING snapshot first + String id = + fastStore.saveSnapshot( + null, + existing -> snapWithSessionAndStatus("sess-sub", SnapshotStatus.PENDING), + SessionStoreOptions.empty()); + + List> received = new CopyOnWriteArrayList<>(); + CountDownLatch changeLatch = new CountDownLatch(2); // initial + change + + AutoCloseable sub = + fastStore.onSnapshotStateChange( + id, + snap -> { + received.add(snap); + changeLatch.countDown(); + }, + SessionStoreOptions.empty()); + + // Save again as ABORTED (same store, should trigger polling callback) + fastStore.saveSnapshot( + id, + existing -> snapWithSessionAndStatus("sess-sub", SnapshotStatus.ABORTED), + SessionStoreOptions.empty()); + + // Wait up to 3 seconds for both callbacks (immediate + change) + boolean completed = changeLatch.await(3, TimeUnit.SECONDS); + + sub.close(); + + assertTrue(completed, "subscriber must fire on PENDING→ABORTED status change within 3 seconds"); + assertFalse(received.isEmpty(), "received list must not be empty"); + // The last received snapshot must be ABORTED + SessionSnapshot last = received.get(received.size() - 1); + assertEquals( + SnapshotStatus.ABORTED, last.getStatus(), "subscriber must receive ABORTED status"); + } + + // ── status defaulting: null status → COMPLETED ─────────────────────────── + + @Test + void testSaveSnapshot_nullStatus_defaultsToCompleted() { + String id = + store.saveSnapshot( + null, + existing -> { + SessionSnapshot> snap = snapWithSession("sess-status"); + snap.setStatus(null); + return snap; + }, + SessionStoreOptions.empty()); + + SessionSnapshot> fetched = + store.getSnapshot(GetSnapshotOptions.builder().snapshotId(id).build()); + assertNotNull(fetched); + assertEquals( + SnapshotStatus.COMPLETED, fetched.getStatus(), "null status must default to COMPLETED"); + } + + // ── empty sessionId → INVALID_ARGUMENT ─────────────────────────────────── + + @Test + void testSaveSnapshot_emptySessionId_throwsInvalidArgument() { + com.google.genkit.core.GenkitException ex = + assertThrows( + com.google.genkit.core.GenkitException.class, + () -> + store.saveSnapshot( + null, + existing -> { + SessionSnapshot> snap = new SessionSnapshot<>(); + snap.setSessionId(""); + return snap; + }, + SessionStoreOptions.empty())); + assertEquals("INVALID_ARGUMENT", ex.getErrorCode()); + } + + // ── mutator returning null → no-op ─────────────────────────────────────── + + @Test + void testSaveSnapshot_mutatorReturnsNull_noOp() { + String id = + store.saveSnapshot( + null, existing -> snapWithSession("sess-noop"), SessionStoreOptions.empty()); + + String result = store.saveSnapshot(id, existing -> null, SessionStoreOptions.empty()); + assertNull(result, "saveSnapshot with null mutator result must return null"); + + // Original snapshot must still be present + SessionSnapshot> fetched = + store.getSnapshot(GetSnapshotOptions.builder().snapshotId(id).build()); + assertNotNull(fetched, "snapshot must still be present after no-op"); + assertEquals("sess-noop", fetched.getSessionId()); + } + + // ── behavior 19: custom prefix creates a subdirectory named after it ──── + + @Test + void testCustomPrefixCreatesSubdirectory() throws IOException { + FileSessionStore> prefixedStore = + FileSessionStore.>builder(tempDir.toString()) + .prefix("tenant-1") + .build(); + + String id = + prefixedStore.saveSnapshot( + null, existing -> snapWithSession("sess-tenant"), SessionStoreOptions.empty()); + + Path snapshotFile = tempDir.resolve("tenant-1").resolve(id + ".json"); + assertTrue( + Files.exists(snapshotFile), "snapshot file must be created under /tenant-1/.json"); + // Must NOT be written under the default "global" prefix. + Path defaultLocation = tempDir.resolve("global").resolve(id + ".json"); + assertFalse( + Files.exists(defaultLocation), "snapshot must not also exist under the default prefix"); + } +} diff --git a/ai/src/test/java/com/google/genkit/ai/agent/InMemorySessionStoreTest.java b/ai/src/test/java/com/google/genkit/ai/agent/InMemorySessionStoreTest.java new file mode 100644 index 000000000..b1987abcf --- /dev/null +++ b/ai/src/test/java/com/google/genkit/ai/agent/InMemorySessionStoreTest.java @@ -0,0 +1,565 @@ +/* + * 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. + * + * SPDX-License-Identifier: Apache-2.0 + */ + +package com.google.genkit.ai.agent; + +import static org.junit.jupiter.api.Assertions.*; + +import com.google.genkit.core.GenkitException; +import java.util.ArrayList; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicReference; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; + +/** TDD tests for InMemorySessionStore (agent package). */ +class InMemorySessionStoreTest { + + private InMemorySessionStore> store; + + @BeforeEach + void setUp() { + store = new InMemorySessionStore<>(); + } + + // ── helpers ────────────────────────────────────────────────────────────────── + + private static SessionSnapshot> snapWithSession(String sessionId) { + return SessionSnapshot.>builder().sessionId(sessionId).build(); + } + + private static SessionSnapshot> snapWithSessionAndCustom( + String sessionId, Map custom) { + SessionState> state = + SessionState.>builder().sessionId(sessionId).custom(custom).build(); + return SessionSnapshot.>builder().sessionId(sessionId).state(state).build(); + } + + // ── save/get round-trip ─────────────────────────────────────────────────────── + + @Test + void testSaveNewSnapshot_mintsUuid_andReturnsIt() { + String id = + store.saveSnapshot( + null, existing -> snapWithSession("sess-1"), SessionStoreOptions.empty()); + assertNotNull(id, "saveSnapshot must return a non-null id"); + assertFalse(id.isBlank(), "minted id must not be blank"); + } + + @Test + void testGetSnapshot_bySnapshotId_roundTrip() { + String id = + store.saveSnapshot( + null, existing -> snapWithSession("sess-1"), SessionStoreOptions.empty()); + + SessionSnapshot> fetched = + store.getSnapshot(GetSnapshotOptions.builder().snapshotId(id).build()); + + assertNotNull(fetched, "getSnapshot by id must return a snapshot"); + assertEquals(id, fetched.getSnapshotId()); + assertEquals("sess-1", fetched.getSessionId()); + } + + @Test + void testGetSnapshot_returnsDeepCopy_mutatingReturnedSnapshotDoesNotChangeStore() { + String id = + store.saveSnapshot( + null, existing -> snapWithSession("sess-1"), SessionStoreOptions.empty()); + + SessionSnapshot> fetched = + store.getSnapshot(GetSnapshotOptions.builder().snapshotId(id).build()); + assertNotNull(fetched); + + // Mutate the returned snapshot + fetched.setSessionId("MUTATED"); + + // Re-fetch from store — must still be original + SessionSnapshot> refetched = + store.getSnapshot(GetSnapshotOptions.builder().snapshotId(id).build()); + assertEquals( + "sess-1", refetched.getSessionId(), "store must not reflect mutation of returned snapshot"); + } + + @Test + void testSaveSnapshot_deepCopyOnStore_mutatingInputDoesNotChangeStore() { + Map custom = new HashMap<>(); + custom.put("key", "original"); + + String id = + store.saveSnapshot( + null, + existing -> snapWithSessionAndCustom("sess-2", custom), + SessionStoreOptions.empty()); + + // Mutate the custom map after save + custom.put("key", "mutated"); + + SessionSnapshot> fetched = + store.getSnapshot(GetSnapshotOptions.builder().snapshotId(id).build()); + assertNotNull(fetched); + assertNotNull(fetched.getState()); + // The stored value must reflect "original", not "mutated" + assertEquals( + "original", + fetched.getState().getCustom().get("key"), + "store must not reflect post-save mutation of mutator's returned snapshot"); + } + + // ── mint UUID ──────────────────────────────────────────────────────────────── + + @Test + void testSaveSnapshot_nullId_mintsUniqueUuids() { + String id1 = + store.saveSnapshot( + null, existing -> snapWithSession("sess-a"), SessionStoreOptions.empty()); + String id2 = + store.saveSnapshot( + null, existing -> snapWithSession("sess-b"), SessionStoreOptions.empty()); + assertNotEquals(id1, id2, "each save with null id must produce a unique id"); + } + + // ── getSnapshot by sessionId (latest leaf) ──────────────────────────────────── + + @Test + void testGetSnapshot_bySessionId_returnsLatestLeaf() { + // Save root snapshot + String rootId = + store.saveSnapshot( + null, existing -> snapWithSession("sess-chain"), SessionStoreOptions.empty()); + + // Save a second snapshot with parentId pointing to first + String childId = + store.saveSnapshot( + null, + existing -> { + SessionSnapshot> snap = snapWithSession("sess-chain"); + snap.setParentId(rootId); + snap.setCreatedAt("2025-01-02T00:00:00Z"); + return snap; + }, + SessionStoreOptions.empty()); + + // getSnapshot by sessionId must return the leaf (child) + SessionSnapshot> leaf = + store.getSnapshot(GetSnapshotOptions.builder().sessionId("sess-chain").build()); + assertNotNull(leaf, "getSnapshot by sessionId must return a snapshot"); + assertEquals(childId, leaf.getSnapshotId(), "must return the leaf (child)"); + } + + @Test + void testGetSnapshot_bySessionId_noMatch_returnsNull() { + SessionSnapshot> result = + store.getSnapshot(GetSnapshotOptions.builder().sessionId("nonexistent").build()); + assertNull(result, "getSnapshot by sessionId with no match must return null"); + } + + // ── preserve existing sessionId ─────────────────────────────────────────────── + + @Test + void testSaveSnapshot_preservesExistingSessionId_whenUpdatingById() { + String id = + store.saveSnapshot( + null, existing -> snapWithSession("sess-orig"), SessionStoreOptions.empty()); + + // Update with a different sessionId in the mutator — store should preserve original + store.saveSnapshot( + id, + existing -> { + assertNotNull(existing, "mutator should receive the existing snapshot"); + assertEquals("sess-orig", existing.getSessionId()); + // Return snapshot with a different sessionId — store must ignore this + return snapWithSession("sess-override"); + }, + SessionStoreOptions.empty()); + + SessionSnapshot> fetched = + store.getSnapshot(GetSnapshotOptions.builder().snapshotId(id).build()); + assertEquals( + "sess-orig", fetched.getSessionId(), "store must preserve existing sessionId on update"); + } + + // ── mutator returning null → no-op ──────────────────────────────────────────── + + @Test + void testSaveSnapshot_mutatorReturnsNull_noOp_returnsNull() { + String id = + store.saveSnapshot( + null, existing -> snapWithSession("sess-noop"), SessionStoreOptions.empty()); + + // Attempt update with mutator returning null + String result = store.saveSnapshot(id, existing -> null, SessionStoreOptions.empty()); + assertNull(result, "saveSnapshot with null mutator result must return null"); + + // Original snapshot must still be unchanged + SessionSnapshot> fetched = + store.getSnapshot(GetSnapshotOptions.builder().snapshotId(id).build()); + assertNotNull(fetched, "snapshot must still be present after no-op"); + assertEquals("sess-noop", fetched.getSessionId()); + } + + // ── empty sessionId → INVALID_ARGUMENT ─────────────────────────────────────── + + @Test + void testSaveSnapshot_emptySessionId_throwsInvalidArgument() { + GenkitException ex = + assertThrows( + GenkitException.class, + () -> + store.saveSnapshot( + null, + existing -> { + SessionSnapshot> snap = new SessionSnapshot<>(); + snap.setSessionId(""); // empty + return snap; + }, + SessionStoreOptions.empty())); + assertEquals( + "INVALID_ARGUMENT", + ex.getErrorCode(), + "empty sessionId must throw GenkitException with INVALID_ARGUMENT"); + } + + @Test + void testSaveSnapshot_nullSessionId_throwsInvalidArgument() { + GenkitException ex = + assertThrows( + GenkitException.class, + () -> + store.saveSnapshot( + null, + existing -> { + // Return snapshot with no sessionId at all + return new SessionSnapshot<>(); + }, + SessionStoreOptions.empty())); + assertEquals( + "INVALID_ARGUMENT", + ex.getErrorCode(), + "null sessionId must throw GenkitException with INVALID_ARGUMENT"); + } + + // ── null status → defaults to COMPLETED ────────────────────────────────────── + + @Test + void testSaveSnapshot_nullStatus_defaultsToCompleted() { + String id = + store.saveSnapshot( + null, + existing -> { + SessionSnapshot> snap = snapWithSession("sess-status"); + snap.setStatus(null); // explicitly null + return snap; + }, + SessionStoreOptions.empty()); + + SessionSnapshot> fetched = + store.getSnapshot(GetSnapshotOptions.builder().snapshotId(id).build()); + assertNotNull(fetched); + assertEquals( + SnapshotStatus.COMPLETED, + fetched.getStatus(), + "null status must be defaulted to COMPLETED"); + } + + // ── subscriber fires on status change ───────────────────────────────────────── + + @Test + void testSubscriber_firesOnStatusChange() throws Exception { + // Save an initial PENDING snapshot + String id = + store.saveSnapshot( + null, + existing -> { + SessionSnapshot> snap = snapWithSession("sess-sub"); + snap.setStatus(SnapshotStatus.PENDING); + return snap; + }, + SessionStoreOptions.empty()); + + List> received = new ArrayList<>(); + + // Subscribe — receives current snapshot immediately since it already exists + AutoCloseable sub = store.onSnapshotStateChange(id, received::add, SessionStoreOptions.empty()); + + // Immediate callback fires with current PENDING snapshot + assertEquals( + 1, + received.size(), + "subscriber must receive current snapshot immediately when existing snapshot is present"); + assertEquals( + SnapshotStatus.PENDING, + received.get(0).getStatus(), + "initial callback must have current PENDING status"); + + // Now save again, flipping status to ABORTED + store.saveSnapshot( + id, + existing -> { + SessionSnapshot> snap = snapWithSession("sess-sub"); + snap.setStatus(SnapshotStatus.ABORTED); + return snap; + }, + SessionStoreOptions.empty()); + + // Subscriber must have been called again with the updated snapshot + assertEquals(2, received.size(), "subscriber must fire again on subsequent status change"); + assertEquals( + SnapshotStatus.ABORTED, + received.get(received.size() - 1).getStatus(), + "subscriber must receive snapshot with new status ABORTED"); + + sub.close(); + } + + @Test + void testSubscriber_doesNotFireWhenStatusUnchanged() { + // Save an initial PENDING snapshot + String id = + store.saveSnapshot( + null, + existing -> { + SessionSnapshot> snap = snapWithSession("sess-no-fire"); + snap.setStatus(SnapshotStatus.PENDING); + return snap; + }, + SessionStoreOptions.empty()); + + List> received = new ArrayList<>(); + AutoCloseable sub = store.onSnapshotStateChange(id, received::add, SessionStoreOptions.empty()); + + // Immediate callback fires with current PENDING snapshot + assertEquals(1, received.size(), "immediate callback must fire"); + + // Save again with same status — subscriber should NOT fire (no status change) + store.saveSnapshot( + id, + existing -> { + SessionSnapshot> snap = snapWithSession("sess-no-fire"); + snap.setStatus(SnapshotStatus.PENDING); + return snap; + }, + SessionStoreOptions.empty()); + + // Still only 1 callback (the immediate one); no additional fire on unchanged status + assertEquals( + 1, + received.size(), + "subscriber must not fire when status is unchanged (only immediate callback)"); + + try { + sub.close(); + } catch (Exception e) { + // ignore + } + } + + @Test + void testSubscriber_unsubscribeStopsNotifications() { + String id = + store.saveSnapshot( + null, + existing -> { + SessionSnapshot> snap = snapWithSession("sess-unsub"); + snap.setStatus(SnapshotStatus.PENDING); + return snap; + }, + SessionStoreOptions.empty()); + + List> received = new ArrayList<>(); + AutoCloseable sub = store.onSnapshotStateChange(id, received::add, SessionStoreOptions.empty()); + + // Immediate callback fires with current PENDING snapshot + assertEquals(1, received.size(), "immediate callback must fire"); + + try { + sub.close(); + } catch (Exception e) { + fail("close() must not throw"); + } + + // Now save again — subscriber is closed, must not fire (no additional callback) + store.saveSnapshot( + id, + existing -> { + SessionSnapshot> snap = snapWithSession("sess-unsub"); + snap.setStatus(SnapshotStatus.ABORTED); + return snap; + }, + SessionStoreOptions.empty()); + + // Still only 1 callback (the immediate one before unsubscribe) + assertEquals(1, received.size(), "closed subscriber must not receive further notifications"); + } + + // ── first save fires subscriber (status: null→COMPLETED counts as creation) ─── + + @Test + void testSubscriber_firesOnFirstSave_whenSubscribedBeforeSave() { + // Subscribe to a snapshotId that doesn't exist yet + List> received = new ArrayList<>(); + AtomicReference subRef = new AtomicReference<>(); + + // We subscribe before saving — since snapshot doesn't exist yet, no immediate callback + // After first save (null→COMPLETED transition), subscriber fires + subRef.set( + store.onSnapshotStateChange("future-id", received::add, SessionStoreOptions.empty())); + + // Save with explicit id + store.saveSnapshot( + "future-id", + existing -> { + SessionSnapshot> snap = snapWithSession("sess-future"); + snap.setStatus(SnapshotStatus.COMPLETED); + return snap; + }, + SessionStoreOptions.empty()); + + // existing was null, so "null != COMPLETED" → subscriber fires + assertFalse(received.isEmpty(), "subscriber must fire on first save (null→COMPLETED)"); + + try { + subRef.get().close(); + } catch (Exception e) { + // ignore + } + } + + // ── getSnapshot with neither snapshotId nor sessionId → null ───────────────── + + @Test + void testGetSnapshot_neitherIdSet_returnsNull() { + SessionSnapshot> result = + store.getSnapshot(GetSnapshotOptions.builder().build()); + assertNull(result, "getSnapshot with no ids must return null"); + } + + // ── rejectBranching constructor ─────────────────────────────────────────────── + + @Test + void testRejectBranchingConstructor_canBeCreated() { + InMemorySessionStore> strictStore = new InMemorySessionStore<>(true); + assertNotNull(strictStore); + + // Single snapshot — no branching, should work fine + String id = + strictStore.saveSnapshot( + null, existing -> snapWithSession("sess-strict"), SessionStoreOptions.empty()); + assertNotNull(id); + } + + // ── subscribe-after-terminal: callback fires immediately with current terminal status ── + + @Test + void testSubscriber_subscribeAfterTerminal_firesImmediatelyWithCurrentStatus() throws Exception { + // Save a snapshot with terminal status ABORTED + String id = + store.saveSnapshot( + null, + existing -> { + SessionSnapshot> snap = snapWithSession("sess-terminal"); + snap.setStatus(SnapshotStatus.ABORTED); + return snap; + }, + SessionStoreOptions.empty()); + + // Verify snapshot is stored with ABORTED status + SessionSnapshot> stored = + store.getSnapshot(GetSnapshotOptions.builder().snapshotId(id).build()); + assertNotNull(stored); + assertEquals(SnapshotStatus.ABORTED, stored.getStatus()); + + // NOW subscribe (after the terminal save) + CountDownLatch latch = new CountDownLatch(1); + AtomicReference> received = new AtomicReference<>(); + + AutoCloseable sub = + store.onSnapshotStateChange( + id, + snap -> { + received.set(snap); + latch.countDown(); + }, + SessionStoreOptions.empty()); + + // Callback must fire immediately with the current ABORTED snapshot + boolean completed = latch.await(1, TimeUnit.SECONDS); + assertTrue(completed, "callback must fire immediately when snapshot already exists"); + assertNotNull(received.get(), "received snapshot must not be null"); + assertEquals( + SnapshotStatus.ABORTED, + received.get().getStatus(), + "immediate callback must have current ABORTED status"); + + try { + sub.close(); + } catch (Exception e) { + // ignore + } + } + + // ── behavior 17: rejectBranching throws when two leaves share a session ────── + + @Test + void testRejectBranchingThrowsWhenTwoLeaves() { + InMemorySessionStore> strictStore = new InMemorySessionStore<>(true); + + // Save a root snapshot for the session. + String rootId = + strictStore.saveSnapshot( + null, + existing -> { + SessionSnapshot> snap = snapWithSession("sess-branch"); + snap.setCreatedAt("2025-01-01T00:00:00Z"); + return snap; + }, + SessionStoreOptions.empty()); + + // Save two children of the SAME root — both are leaves (neither is referenced as a parentId). + strictStore.saveSnapshot( + null, + existing -> { + SessionSnapshot> snap = snapWithSession("sess-branch"); + snap.setParentId(rootId); + snap.setCreatedAt("2025-01-02T00:00:00Z"); + return snap; + }, + SessionStoreOptions.empty()); + strictStore.saveSnapshot( + null, + existing -> { + SessionSnapshot> snap = snapWithSession("sess-branch"); + snap.setParentId(rootId); + snap.setCreatedAt("2025-01-02T00:00:01Z"); + return snap; + }, + SessionStoreOptions.empty()); + + // getSnapshot by sessionId must throw FAILED_PRECONDITION: two leaves detected. + GenkitException ex = + assertThrows( + GenkitException.class, + () -> + strictStore.getSnapshot( + GetSnapshotOptions.builder().sessionId("sess-branch").build())); + assertEquals("FAILED_PRECONDITION", ex.getErrorCode()); + } +} diff --git a/ai/src/test/java/com/google/genkit/ai/agent/LeafSelectionTest.java b/ai/src/test/java/com/google/genkit/ai/agent/LeafSelectionTest.java new file mode 100644 index 000000000..c633639c3 --- /dev/null +++ b/ai/src/test/java/com/google/genkit/ai/agent/LeafSelectionTest.java @@ -0,0 +1,164 @@ +/* + * 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. + * + * SPDX-License-Identifier: Apache-2.0 + */ + +package com.google.genkit.ai.agent; + +import static org.junit.jupiter.api.Assertions.*; + +import com.google.genkit.ai.agent.internal.LeafSelection; +import com.google.genkit.core.GenkitException; +import java.util.Arrays; +import java.util.Collections; +import java.util.List; +import org.junit.jupiter.api.Test; + +/** TDD tests for LeafSelection.selectLeaf. */ +class LeafSelectionTest { + + // Helpers + private static SessionSnapshot snap(String id, String parentId, String createdAt) { + return SessionSnapshot.builder() + .snapshotId(id) + .parentId(parentId) + .createdAt(createdAt) + .build(); + } + + // ---- empty list → null ---- + + @Test + void testEmptyList_returnsNull() { + assertNull(LeafSelection.selectLeaf(Collections.emptyList(), false)); + } + + @Test + void testEmptyList_rejectBranching_returnsNull() { + assertNull(LeafSelection.selectLeaf(Collections.emptyList(), true)); + } + + // ---- single snapshot → itself ---- + + @Test + void testSingleSnapshot_returnsItself() { + SessionSnapshot a = snap("A", null, "2024-01-01T00:00:00Z"); + SessionSnapshot result = LeafSelection.selectLeaf(Collections.singletonList(a), false); + assertSame(a, result); + } + + @Test + void testSingleSnapshot_noCreatedAt_returnsItself() { + SessionSnapshot a = snap("A", null, null); + SessionSnapshot result = LeafSelection.selectLeaf(Collections.singletonList(a), false); + assertSame(a, result); + } + + // ---- linear chain A ← B ← C: selectLeaf → C ---- + + @Test + void testLinearChain_returnsLeaf() { + SessionSnapshot a = snap("A", null, "2024-01-01T00:00:00Z"); + SessionSnapshot b = snap("B", "A", "2024-01-02T00:00:00Z"); + SessionSnapshot c = snap("C", "B", "2024-01-03T00:00:00Z"); + + List> chain = Arrays.asList(a, b, c); + SessionSnapshot result = LeafSelection.selectLeaf(chain, false); + assertEquals("C", result.getSnapshotId()); + } + + @Test + void testLinearChain_orderedDifferently_returnsLeaf() { + // Same chain but snapshots in reverse order in the list + SessionSnapshot a = snap("A", null, "2024-01-01T00:00:00Z"); + SessionSnapshot b = snap("B", "A", "2024-01-02T00:00:00Z"); + SessionSnapshot c = snap("C", "B", "2024-01-03T00:00:00Z"); + + List> chain = Arrays.asList(c, b, a); + SessionSnapshot result = LeafSelection.selectLeaf(chain, false); + assertEquals("C", result.getSnapshotId()); + } + + // ---- two leaves: C1(parent=B), C2(parent=B), C2.createdAt later ---- + + @Test + void testTwoLeaves_rejectBranchingFalse_returnsMostRecent() { + SessionSnapshot a = snap("A", null, "2024-01-01T00:00:00Z"); + SessionSnapshot b = snap("B", "A", "2024-01-02T00:00:00Z"); + SessionSnapshot c1 = snap("C1", "B", "2024-01-03T00:00:00Z"); + SessionSnapshot c2 = snap("C2", "B", "2024-01-04T00:00:00Z"); + + List> snapshots = Arrays.asList(a, b, c1, c2); + SessionSnapshot result = LeafSelection.selectLeaf(snapshots, false); + assertEquals("C2", result.getSnapshotId()); + } + + @Test + void testTwoLeaves_rejectBranchingTrue_throws() { + SessionSnapshot a = snap("A", null, "2024-01-01T00:00:00Z"); + SessionSnapshot b = snap("B", "A", "2024-01-02T00:00:00Z"); + SessionSnapshot c1 = snap("C1", "B", "2024-01-03T00:00:00Z"); + SessionSnapshot c2 = snap("C2", "B", "2024-01-04T00:00:00Z"); + + List> snapshots = Arrays.asList(a, b, c1, c2); + GenkitException ex = + assertThrows(GenkitException.class, () -> LeafSelection.selectLeaf(snapshots, true)); + assertEquals("FAILED_PRECONDITION", ex.getErrorCode()); + } + + // ---- tie-break by snapshotId when createdAt equal ---- + + @Test + void testTwoLeaves_sameCreatedAt_tieBreakBySnapshotId() { + SessionSnapshot a = snap("A", null, "2024-01-01T00:00:00Z"); + SessionSnapshot b = snap("B", "A", "2024-01-02T00:00:00Z"); + // Same timestamp; lexicographically "Z" > "C" + SessionSnapshot c = snap("C", "B", "2024-01-03T00:00:00Z"); + SessionSnapshot z = snap("Z", "B", "2024-01-03T00:00:00Z"); + + List> snapshots = Arrays.asList(a, b, c, z); + SessionSnapshot result = LeafSelection.selectLeaf(snapshots, false); + assertEquals("Z", result.getSnapshotId()); + } + + // ---- null createdAt treated as earliest (epoch) ---- + + @Test + void testNullCreatedAt_treatedAsEarliest() { + SessionSnapshot a = snap("A", null, null); + SessionSnapshot b = snap("B", "A", "2024-01-02T00:00:00Z"); + SessionSnapshot c1 = snap("C1", "B", null); // null = earliest + SessionSnapshot c2 = snap("C2", "B", "2024-01-03T00:00:00Z"); + + List> snapshots = Arrays.asList(a, b, c1, c2); + SessionSnapshot result = LeafSelection.selectLeaf(snapshots, false); + assertEquals("C2", result.getSnapshotId()); + } + + // ---- zero leaves in non-empty list (cycle) → throws FAILED_PRECONDITION ---- + + @Test + void testCycle_noLeaves_throws() { + // A points to B, B points to A — neither is a leaf + SessionSnapshot a = snap("A", "B", "2024-01-01T00:00:00Z"); + SessionSnapshot b = snap("B", "A", "2024-01-02T00:00:00Z"); + + List> snapshots = Arrays.asList(a, b); + GenkitException ex = + assertThrows(GenkitException.class, () -> LeafSelection.selectLeaf(snapshots, false)); + assertEquals("FAILED_PRECONDITION", ex.getErrorCode()); + } +} diff --git a/ai/src/test/java/com/google/genkit/ai/agent/SessionResolverTest.java b/ai/src/test/java/com/google/genkit/ai/agent/SessionResolverTest.java new file mode 100644 index 000000000..cf01e82a8 --- /dev/null +++ b/ai/src/test/java/com/google/genkit/ai/agent/SessionResolverTest.java @@ -0,0 +1,246 @@ +/* + * 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. + * + * SPDX-License-Identifier: Apache-2.0 + */ + +package com.google.genkit.ai.agent; + +import static org.junit.jupiter.api.Assertions.*; + +import com.google.genkit.ai.agent.internal.SessionResolver; +import com.google.genkit.core.GenkitException; +import java.util.Map; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; + +/** TDD tests for SessionResolver. */ +class SessionResolverTest { + + private InMemorySessionStore> store; + private SessionStoreOptions opts; + + @BeforeEach + void setUp() { + store = new InMemorySessionStore<>(); + opts = SessionStoreOptions.empty(); + } + + // ── Helper ─────────────────────────────────────────────────────────────────── + + /** Saves a snapshot with the given sessionId and status into the store. Returns snapshotId. */ + private String seedSnapshot(String sessionId, SnapshotStatus status) { + SessionState> state = + SessionState.>builder().sessionId(sessionId).build(); + SessionSnapshot> snap = + SessionSnapshot.>builder() + .sessionId(sessionId) + .status(status) + .state(state) + .build(); + return store.saveSnapshot(null, existing -> snap, opts); + } + + // ── Client-managed ──────────────────────────────────────────────────────────── + + @Test + void clientManaged_withState_hydratesSession() { + SessionState> state = + SessionState.>builder().sessionId("client-session-1").build(); + AgentInit> init = + AgentInit.>builder().state(state).build(); + + SessionResolver.Resolution> result = + SessionResolver.resolve(null, false, init, opts); + + assertTrue(result.isOk()); + assertNotNull(result.session()); + assertEquals("client-session-1", result.session().sessionId()); + } + + @Test + void clientManaged_withSessionId_throwsFailedPrecondition() { + AgentInit> init = + AgentInit.>builder().sessionId("some-id").build(); + + GenkitException ex = + assertThrows(GenkitException.class, () -> SessionResolver.resolve(null, false, init, opts)); + assertEquals("FAILED_PRECONDITION", ex.getErrorCode()); + assertTrue(ex.getMessage().contains("sessionId")); + } + + @Test + void clientManaged_withSnapshotId_throwsFailedPrecondition() { + AgentInit> init = + AgentInit.>builder().snapshotId("snap-1").build(); + + GenkitException ex = + assertThrows(GenkitException.class, () -> SessionResolver.resolve(null, false, init, opts)); + assertEquals("FAILED_PRECONDITION", ex.getErrorCode()); + assertTrue(ex.getMessage().contains("snapshotId")); + } + + @Test + void clientManaged_noInit_freshSession() { + SessionResolver.Resolution> result = + SessionResolver.resolve(null, false, null, opts); + + assertTrue(result.isOk()); + assertNotNull(result.session()); + assertNotNull(result.session().sessionId()); + } + + // ── Server-managed: state → THROW ──────────────────────────────────────────── + + @Test + void serverManaged_withState_throwsFailedPrecondition() { + SessionState> state = + SessionState.>builder().sessionId("s1").build(); + AgentInit> init = + AgentInit.>builder().state(state).build(); + + GenkitException ex = + assertThrows(GenkitException.class, () -> SessionResolver.resolve(store, true, init, opts)); + assertEquals("FAILED_PRECONDITION", ex.getErrorCode()); + assertTrue(ex.getMessage().contains("state")); + } + + // ── Server-managed: snapshotId branch ──────────────────────────────────────── + + @Test + void serverManaged_snapshotId_completed_hydratesSession() { + String sid = "sess-snap-1"; + String snapId = seedSnapshot(sid, SnapshotStatus.COMPLETED); + + AgentInit> init = + AgentInit.>builder().snapshotId(snapId).build(); + + SessionResolver.Resolution> result = + SessionResolver.resolve(store, true, init, opts); + + assertTrue(result.isOk()); + assertEquals(sid, result.session().sessionId()); + } + + @Test + void serverManaged_unknownSnapshotId_throwsInvalidArgument() { + AgentInit> init = + AgentInit.>builder().snapshotId("nonexistent-snap").build(); + + GenkitException ex = + assertThrows(GenkitException.class, () -> SessionResolver.resolve(store, true, init, opts)); + assertEquals("INVALID_ARGUMENT", ex.getErrorCode()); + assertTrue(ex.getMessage().toLowerCase().contains("snapshot")); + } + + @Test + void serverManaged_nonCompletedSnapshot_throwsInvalidArgument() { + String snapId = seedSnapshot("sess-pending", SnapshotStatus.PENDING); + + AgentInit> init = + AgentInit.>builder().snapshotId(snapId).build(); + + GenkitException ex = + assertThrows(GenkitException.class, () -> SessionResolver.resolve(store, true, init, opts)); + assertEquals("INVALID_ARGUMENT", ex.getErrorCode()); + assertTrue(ex.getMessage().toLowerCase().contains("resumable")); + } + + @Test + void serverManaged_snapshotId_withMismatchedSessionId_throwsInvalidArgument() { + String snapId = seedSnapshot("sess-real", SnapshotStatus.COMPLETED); + + AgentInit> init = + AgentInit.>builder() + .snapshotId(snapId) + .sessionId("wrong-session-id") + .build(); + + GenkitException ex = + assertThrows(GenkitException.class, () -> SessionResolver.resolve(store, true, init, opts)); + assertEquals("INVALID_ARGUMENT", ex.getErrorCode()); + assertTrue(ex.getMessage().toLowerCase().contains("session")); + } + + // ── Server-managed: sessionId branch ───────────────────────────────────────── + + @Test + void serverManaged_sessionId_existingCompletedLeaf_hydratesSession() { + String sid = "sess-existing"; + seedSnapshot(sid, SnapshotStatus.COMPLETED); + + AgentInit> init = + AgentInit.>builder().sessionId(sid).build(); + + SessionResolver.Resolution> result = + SessionResolver.resolve(store, true, init, opts); + + assertTrue(result.isOk()); + assertEquals(sid, result.session().sessionId()); + } + + @Test + void serverManaged_sessionId_unknownSession_freshSessionBoundToId() { + String sid = "new-unknown-session"; + AgentInit> init = + AgentInit.>builder().sessionId(sid).build(); + + SessionResolver.Resolution> result = + SessionResolver.resolve(store, true, init, opts); + + assertTrue(result.isOk()); + assertEquals(sid, result.session().sessionId()); + } + + @Test + void serverManaged_sessionId_pendingLeaf_throwsFailedPrecondition() { + String sid = "sess-with-pending"; + seedSnapshot(sid, SnapshotStatus.PENDING); + + AgentInit> init = + AgentInit.>builder().sessionId(sid).build(); + + GenkitException ex = + assertThrows(GenkitException.class, () -> SessionResolver.resolve(store, true, init, opts)); + assertEquals("FAILED_PRECONDITION", ex.getErrorCode()); + assertTrue(ex.getMessage().toLowerCase().contains("resume")); + } + + // ── Server-managed: no init → fresh session ─────────────────────────────────── + + @Test + void serverManaged_noInit_freshSession() { + SessionResolver.Resolution> result = + SessionResolver.resolve(store, true, null, opts); + + assertTrue(result.isOk()); + assertNotNull(result.session()); + assertNotNull(result.session().sessionId()); + } + + // ── Resolution type ─────────────────────────────────────────────────────────── + + @Test + void resolution_failure_returnsFailureResolution() { + RuntimeError err = RuntimeError.builder().status("FAILED").message("something wrong").build(); + SessionResolver.Resolution> failure = + SessionResolver.Resolution.failure(err); + + assertFalse(failure.isOk()); + assertNull(failure.session()); + assertNotNull(failure.error()); + assertEquals("something wrong", failure.error().getMessage()); + } +} diff --git a/ai/src/test/java/com/google/genkit/ai/agent/SessionRunnerTest.java b/ai/src/test/java/com/google/genkit/ai/agent/SessionRunnerTest.java new file mode 100644 index 000000000..59a98a100 --- /dev/null +++ b/ai/src/test/java/com/google/genkit/ai/agent/SessionRunnerTest.java @@ -0,0 +1,392 @@ +/* + * 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. + * + * SPDX-License-Identifier: Apache-2.0 + */ + +package com.google.genkit.ai.agent; + +import static org.junit.jupiter.api.Assertions.*; + +import com.google.genkit.ai.Message; +import com.google.genkit.ai.Part; +import com.google.genkit.ai.Role; +import com.google.genkit.ai.agent.internal.AbortAwareMutator; +import com.google.genkit.core.GenkitException; +import java.util.HashMap; +import java.util.Map; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; + +/** TDD tests for SessionRunner (Task 4.1). */ +class SessionRunnerTest { + + private static final SessionStoreOptions OPTS = SessionStoreOptions.empty(); + + private InMemorySessionStore> store; + private Session> session; + + @BeforeEach + void setUp() { + store = new InMemorySessionStore<>(); + SessionState> initialState = + SessionState.>builder() + .sessionId("test-session-1") + .custom(new HashMap<>()) + .build(); + session = new Session<>(initialState); + } + + // ── helpers ───────────────────────────────────────────────────────────────── + + private static Message userMessage(String text) { + return new Message(Role.USER, java.util.List.of(Part.text(text))); + } + + private static Message modelMessage(String text) { + return new Message(Role.MODEL, java.util.List.of(Part.text(text))); + } + + // ── Two successful turns with parentId chaining ────────────────────────────── + + @Test + void testTwoSuccessfulTurns_parentIdChaining_andTurnIndex() { + SessionRunner> runner = new SessionRunner<>(session, store, OPTS); + assertEquals(0, runner.turnIndex()); + assertNull(runner.lastSnapshot()); + assertNull(runner.lastSnapshotId()); + + // Turn 1: user sends "hello", model replies "hi" + AgentInput input1 = AgentInput.builder().message(userMessage("hello")).build(); + runner.runTurn( + input1, + (inp, ctx) -> { + runner.addMessages(modelMessage("hi")); + return AgentFinishReason.STOP; + }); + + assertEquals(1, runner.turnIndex()); + assertNotNull(runner.lastSnapshot()); + String snap1Id = runner.lastSnapshotId(); + assertNotNull(snap1Id); + assertFalse(snap1Id.isEmpty()); + assertEquals(SnapshotStatus.COMPLETED, runner.lastSnapshot().getStatus()); + assertNull(runner.lastSnapshot().getParentId()); // first turn has no parent + + // Turn 2: user sends "world", model replies "there" + AgentInput input2 = AgentInput.builder().message(userMessage("world")).build(); + runner.runTurn( + input2, + (inp, ctx) -> { + runner.addMessages(modelMessage("there")); + return AgentFinishReason.STOP; + }); + + assertEquals(2, runner.turnIndex()); + String snap2Id = runner.lastSnapshotId(); + assertNotNull(snap2Id); + assertFalse(snap2Id.isEmpty()); + assertNotEquals(snap1Id, snap2Id); + + // Second snapshot's parentId must be the first snapshot's id + assertEquals(snap1Id, runner.lastSnapshot().getParentId()); + assertEquals(SnapshotStatus.COMPLETED, runner.lastSnapshot().getStatus()); + + // getSnapshot by sessionId returns the leaf (second) snapshot + GetSnapshotOptions getOpts = GetSnapshotOptions.builder().sessionId("test-session-1").build(); + SessionSnapshot leaf = store.getSnapshot(getOpts); + assertNotNull(leaf); + assertEquals(snap2Id, leaf.getSnapshotId()); + assertEquals(snap1Id, leaf.getParentId()); + + // Messages accumulate: user+model x2 = 4 messages + assertEquals(4, runner.getMessages().size()); + + // lastTurnFinishReason should be STOP (not FAILED) + assertEquals(AgentFinishReason.STOP, runner.lastTurnFinishReason()); + assertNull(runner.lastTurnError()); + } + + // ── Throwing turnBody → FAILED snapshot, no rethrow ───────────────────────── + + @Test + void testThrowingTurnBody_gracefulFailedSnapshot_noRethrow() { + SessionRunner> runner = new SessionRunner<>(session, store, OPTS); + + AgentInput input = AgentInput.builder().message(userMessage("hello")).build(); + + // Should NOT throw even though turnBody throws + assertDoesNotThrow( + () -> + runner.runTurn( + input, + (inp, ctx) -> { + throw new RuntimeException("simulated turn failure"); + })); + + // turnIndex advances even on failure + assertEquals(1, runner.turnIndex()); + + // FAILED snapshot was persisted + assertNotNull(runner.lastSnapshot()); + assertEquals(SnapshotStatus.FAILED, runner.lastSnapshot().getStatus()); + assertEquals(AgentFinishReason.FAILED, runner.lastTurnFinishReason()); + + // Error was recorded on the runner + assertNotNull(runner.lastTurnError()); + assertEquals("simulated turn failure", runner.lastTurnError().getMessage()); + + // Snapshot error was recorded + assertNotNull(runner.lastSnapshot().getError()); + + // Store also has the failed snapshot + String failedId = runner.lastSnapshotId(); + assertNotNull(failedId); + GetSnapshotOptions getOpts = GetSnapshotOptions.builder().snapshotId(failedId).build(); + SessionSnapshot stored = store.getSnapshot(getOpts); + assertNotNull(stored); + assertEquals(SnapshotStatus.FAILED, stored.getStatus()); + } + + // ── Client-managed (store=null) ────────────────────────────────────────────── + + @Test + void testClientManaged_nullStore_noException_stateReflected() { + SessionRunner> runner = new SessionRunner<>(session, null, OPTS); + + AgentInput input = AgentInput.builder().message(userMessage("hi")).build(); + + assertDoesNotThrow( + () -> + runner.runTurn( + input, + (inp, ctx) -> { + runner.addMessages(modelMessage("hello")); + return AgentFinishReason.STOP; + })); + + assertEquals(1, runner.turnIndex()); + // Client-managed: lastSnapshotId is "" (empty string) + assertEquals("", runner.lastSnapshotId()); + // But lastSnapshot is still tracked in memory + assertNotNull(runner.lastSnapshot()); + assertEquals(SnapshotStatus.COMPLETED, runner.lastSnapshot().getStatus()); + + // State reflects the added messages + assertEquals(2, runner.getMessages().size()); + } + + // ── Invalid input validation → INVALID_ARGUMENT thrown ─────────────────────── + + @Test + void testInvalidInput_toolRequestPart_throwsGenkitException() { + SessionRunner> runner = new SessionRunner<>(session, store, OPTS); + + Part toolRequestPart = new Part(); + toolRequestPart.setToolRequest(new com.google.genkit.ai.ToolRequest()); + Message badMsg = new Message(Role.USER, java.util.List.of(toolRequestPart)); + AgentInput input = AgentInput.builder().message(badMsg).build(); + + GenkitException ex = + assertThrows( + GenkitException.class, + () -> runner.runTurn(input, (inp, ctx) -> AgentFinishReason.STOP)); + assertEquals("INVALID_ARGUMENT", ex.getErrorCode()); + } + + @Test + void testInvalidInput_toolResponsePart_throwsGenkitException() { + SessionRunner> runner = new SessionRunner<>(session, store, OPTS); + + Part toolResponsePart = new Part(); + toolResponsePart.setToolResponse(new com.google.genkit.ai.ToolResponse()); + Message badMsg = new Message(Role.USER, java.util.List.of(toolResponsePart)); + AgentInput input = AgentInput.builder().message(badMsg).build(); + + GenkitException ex = + assertThrows( + GenkitException.class, + () -> runner.runTurn(input, (inp, ctx) -> AgentFinishReason.STOP)); + assertEquals("INVALID_ARGUMENT", ex.getErrorCode()); + } + + @Test + void testInvalidInput_nonUserRole_throwsGenkitException() { + SessionRunner> runner = new SessionRunner<>(session, store, OPTS); + + Message badMsg = modelMessage("should not be allowed"); + AgentInput input = AgentInput.builder().message(badMsg).build(); + + GenkitException ex = + assertThrows( + GenkitException.class, + () -> runner.runTurn(input, (inp, ctx) -> AgentFinishReason.STOP)); + assertEquals("INVALID_ARGUMENT", ex.getErrorCode()); + } + + // ── AbortAwareMutator unit test ────────────────────────────────────────────── + + @Test + void testAbortAwareMutator_existingAborted_returnsNull() { + SessionSnapshot> existing = + SessionSnapshot.>builder() + .snapshotId("snap-1") + .sessionId("test-session-1") + .status(SnapshotStatus.ABORTED) + .build(); + + SessionSnapshot> toWrite = + SessionSnapshot.>builder() + .snapshotId("snap-1") + .sessionId("test-session-1") + .status(SnapshotStatus.COMPLETED) + .build(); + + SnapshotMutator> inner = (e) -> toWrite; + SnapshotMutator> wrapped = AbortAwareMutator.wrap(inner); + + // When existing is ABORTED, wrapped mutator must return null (no-op) + assertNull(wrapped.apply(existing)); + } + + @Test + void testAbortAwareMutator_existingNotAborted_delegates() { + SessionSnapshot> existing = + SessionSnapshot.>builder() + .snapshotId("snap-1") + .sessionId("test-session-1") + .status(SnapshotStatus.COMPLETED) + .build(); + + SessionSnapshot> toWrite = + SessionSnapshot.>builder() + .snapshotId("snap-1") + .sessionId("test-session-1") + .status(SnapshotStatus.COMPLETED) + .build(); + + SnapshotMutator> inner = (e) -> toWrite; + SnapshotMutator> wrapped = AbortAwareMutator.wrap(inner); + + // When existing is not ABORTED, wrapped mutator delegates to inner + SessionSnapshot> result = wrapped.apply(existing); + assertNotNull(result); + assertSame(toWrite, result); + } + + @Test + void testAbortAwareMutator_existingNull_delegates() { + SessionSnapshot> toWrite = + SessionSnapshot.>builder() + .snapshotId("snap-1") + .sessionId("test-session-1") + .status(SnapshotStatus.COMPLETED) + .build(); + + SnapshotMutator> inner = (e) -> toWrite; + SnapshotMutator> wrapped = AbortAwareMutator.wrap(inner); + + // When existing is null (no prior snapshot), wrapped mutator delegates to inner + SessionSnapshot> result = wrapped.apply(null); + assertNotNull(result); + assertSame(toWrite, result); + } + + // ── InterruptedException → ABORTED (not FAILED) ────────────────────────────── + + @Test + void testInterruptedTurnBody_abortedStatus() { + SessionRunner> runner = new SessionRunner<>(session, store, OPTS); + + AgentInput input = AgentInput.builder().message(userMessage("hello")).build(); + + // Should NOT throw even though turnBody throws InterruptedException + assertDoesNotThrow( + () -> + runner.runTurn( + input, + (inp, ctx) -> { + throw new InterruptedException("aborted by user"); + })); + + assertEquals(1, runner.turnIndex()); + // For InterruptedException, status should be ABORTED, not FAILED + assertNotNull(runner.lastSnapshot()); + assertEquals(SnapshotStatus.ABORTED, runner.lastSnapshot().getStatus()); + assertEquals(AgentFinishReason.ABORTED, runner.lastTurnFinishReason()); + } + + // ── Message deep-copy: external mutation should not corrupt history ───────── + + @Test + void testMessageDeepCopy_externalMutationDoesNotAffectHistory() { + SessionRunner> runner = new SessionRunner<>(session, store, OPTS); + + // Create a user message and pass it via input + Message originalMsg = userMessage("original text"); + AgentInput input = AgentInput.builder().message(originalMsg).build(); + + runner.runTurn( + input, + (inp, ctx) -> { + runner.addMessages(modelMessage("response")); + return AgentFinishReason.STOP; + }); + + // Now mutate the original message object (caller's reference) + originalMsg.setContent(java.util.List.of(Part.text("mutated text"))); + + // Verify that the session's stored message is unchanged + java.util.List messages = runner.getMessages(); + assertEquals(2, messages.size()); + + Message storedUserMsg = messages.get(0); + assertEquals("original text", storedUserMsg.getText()); + assertNotEquals("mutated text", storedUserMsg.getText()); + } + + // ── behavior 16: sess.addMessages splices messages directly into history ──── + + @Test + void testAddMessagesManuallySplicedIntoHistory() { + SessionRunner> runner = new SessionRunner<>(session, store, OPTS); + + AgentInput input = AgentInput.builder().message(userMessage("hello")).build(); + runner.runTurn( + input, + (inp, ctx) -> { + // Directly splice extra messages into history from within the turn body, simulating + // an AgentFn that records e.g. tool-call bookkeeping messages via sess.addMessages(...). + runner.addMessages(modelMessage("spliced-1"), modelMessage("spliced-2")); + return AgentFinishReason.STOP; + }); + + java.util.List messages = runner.getMessages(); + // user "hello" + 2 spliced messages = 3 + assertEquals(3, messages.size()); + assertEquals("spliced-1", messages.get(1).getText()); + assertEquals("spliced-2", messages.get(2).getText()); + + // The spliced messages must also be present in the persisted snapshot after the turn. + GetSnapshotOptions getOpts = + GetSnapshotOptions.builder().snapshotId(runner.lastSnapshotId()).build(); + SessionSnapshot> stored = store.getSnapshot(getOpts); + assertNotNull(stored); + java.util.List storedMsgs = stored.getState().getMessages(); + assertEquals(3, storedMsgs.size()); + assertEquals("spliced-1", storedMsgs.get(1).getText()); + assertEquals("spliced-2", storedMsgs.get(2).getText()); + } +} diff --git a/ai/src/test/java/com/google/genkit/ai/agent/SessionTest.java b/ai/src/test/java/com/google/genkit/ai/agent/SessionTest.java new file mode 100644 index 000000000..dd8f540cf --- /dev/null +++ b/ai/src/test/java/com/google/genkit/ai/agent/SessionTest.java @@ -0,0 +1,398 @@ +/* + * 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. + * + * SPDX-License-Identifier: Apache-2.0 + */ + +package com.google.genkit.ai.agent; + +import static org.junit.jupiter.api.Assertions.*; + +import com.google.genkit.ai.Message; +import java.util.ArrayList; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.atomic.AtomicReference; +import org.junit.jupiter.api.Test; + +/** TDD tests for Session, ArtifactStore, and AgentSessionContext. */ +class SessionTest { + + // ---- updateCustom: version increment + listener ---- + + @Test + void testGetCustomReturnsDeepcopy() { + Map custom = new HashMap<>(); + custom.put("k", 1); + + SessionState> state = + SessionState.>builder().custom(custom).build(); + + Session> session = new Session<>(state); + + // Get the custom state and mutate it + Map returned = session.getCustom(); + returned.put("k", 999); + + // Fetch fresh and verify unchanged + Map fresh = session.getCustom(); + assertEquals(1, fresh.get("k"), "mutating returned custom must not alter internal state"); + + // Also verify via getState().getCustom() + SessionState> snapshot = session.getState(); + assertEquals( + 1, + snapshot.getCustom().get("k"), + "getState().getCustom() must also show unchanged internal state"); + } + + @Test + void testUpdateCustomIncrementsVersion() { + Map custom = new HashMap<>(); + custom.put("count", 0); + + SessionState> state = + SessionState.>builder().custom(custom).build(); + + Session> session = new Session<>(state); + long versionBefore = session.getVersion(); + + session.updateCustom( + c -> { + Map next = new HashMap<>(c); + next.put("count", 1); + return next; + }); + + assertEquals( + versionBefore + 1, session.getVersion(), "version must increment after updateCustom"); + assertEquals(1, session.getCustom().get("count"), "custom state must be updated"); + } + + @Test + void testUpdateCustomFiresOnCustomChangedListener() { + Map custom = new HashMap<>(); + custom.put("value", "initial"); + + SessionState> state = + SessionState.>builder().custom(custom).build(); + + Session> session = new Session<>(state); + + AtomicInteger callCount = new AtomicInteger(0); + session.setOnCustomChanged(() -> callCount.incrementAndGet()); + + session.updateCustom( + c -> { + Map next = new HashMap<>(c); + next.put("value", "changed"); + return next; + }); + + assertEquals(1, callCount.get(), "onCustomChanged must be called once"); + } + + // ---- addArtifacts: dedup by name ---- + + @Test + void testAddArtifactsDeduplicatesByName() { + SessionState> state = SessionState.>builder().build(); + Session> session = new Session<>(state); + + Artifact v1 = Artifact.builder().name("doc1").build(); + session.addArtifacts(v1); + assertEquals(1, session.getArtifacts().size()); + + Artifact v2 = Artifact.builder().name("doc1").build(); + session.addArtifacts(v2); + + List artifacts = session.getArtifacts(); + assertEquals(1, artifacts.size(), "dedup: same name must not add a second entry"); + assertSame( + v2.getName(), + artifacts.get(0).getName(), + "replaced artifact must have the same name reference"); + } + + @Test + void testAddArtifactsDifferentNameAppends() { + SessionState> state = SessionState.>builder().build(); + Session> session = new Session<>(state); + + session.addArtifacts(Artifact.builder().name("doc1").build()); + session.addArtifacts(Artifact.builder().name("doc2").build()); + + assertEquals(2, session.getArtifacts().size(), "different names must both be present"); + } + + @Test + void testAddArtifactsUpdateFiresListener() { + SessionState> state = SessionState.>builder().build(); + Session> session = new Session<>(state); + + AtomicInteger addCount = new AtomicInteger(0); + session.setOnArtifactChanged(a -> addCount.incrementAndGet()); + + Artifact v1 = Artifact.builder().name("doc1").build(); + session.addArtifacts(v1); + assertEquals(1, addCount.get(), "listener must fire on add"); + + Artifact v2 = Artifact.builder().name("doc1").build(); + session.addArtifacts(v2); + assertEquals(2, addCount.get(), "listener must fire on update/replace"); + } + + @Test + void testAddArtifactsNullNameAlwaysAppends() { + SessionState> state = SessionState.>builder().build(); + Session> session = new Session<>(state); + + session.addArtifacts(Artifact.builder().build()); // name is null + session.addArtifacts(Artifact.builder().build()); // name is null + + assertEquals(2, session.getArtifacts().size(), "null-named artifacts must always be appended"); + } + + // ---- getState deep copy ---- + + @Test + void testGetStateReturnsDeepCopy() { + List messages = new ArrayList<>(); + messages.add(Message.user("hello")); + + SessionState> state = + SessionState.>builder() + .sessionId("sess-deep") + .messages(messages) + .build(); + + Session> session = new Session<>(state); + + SessionState> snapshot = session.getState(); + + // Mutate the returned snapshot's messages list + if (snapshot.getMessages() != null) { + snapshot.setMessages(new ArrayList<>()); + } + + // Internal state must be unchanged + SessionState> snapshot2 = session.getState(); + assertNotNull( + snapshot2.getMessages(), "internal messages must not be cleared by external mutation"); + assertFalse( + snapshot2.getMessages().isEmpty(), + "internal messages must retain original entries after external mutation"); + } + + // ---- constructor sessionId minting ---- + + @Test + void testConstructorMintsSessionIdWhenAbsent() { + SessionState> state = + SessionState.>builder().build(); // no sessionId + + Session> session = new Session<>(state); + + assertNotNull(session.sessionId(), "sessionId must not be null when not provided"); + assertFalse(session.sessionId().isEmpty(), "sessionId must not be empty when not provided"); + } + + @Test + void testConstructorPreservesProvidedSessionId() { + SessionState> state = + SessionState.>builder().sessionId("my-session").build(); + + Session> session = new Session<>(state); + + assertEquals("my-session", session.sessionId(), "provided sessionId must be preserved"); + } + + // ---- getMessages returns a copy ---- + + @Test + void testGetMessagesReturnsCopy() { + List messages = new ArrayList<>(); + messages.add(Message.user("initial")); + + SessionState> state = + SessionState.>builder().messages(messages).build(); + + Session> session = new Session<>(state); + + List copy = session.getMessages(); + copy.add(Message.model("extra")); + + List copy2 = session.getMessages(); + assertEquals(1, copy2.size(), "mutating returned messages must not affect internal state"); + } + + // ---- AgentSessionContext ---- + + @Test + void testAgentSessionContextCurrentReturnsNullWhenNoSession() { + // Ensure clean state + assertNull( + AgentSessionContext.current(), + "current() must return null when no session is bound to the thread"); + } + + @Test + void testAgentSessionContextRunBindsSession() { + SessionState> state = + SessionState.>builder().sessionId("ctx-sess").build(); + + Session> session = new Session<>(state); + AtomicReference> captured = new AtomicReference<>(); + + AgentSessionContext.run(session, () -> captured.set(AgentSessionContext.current())); + + assertSame(session, captured.get(), "run() must bind session to the thread context"); + } + + @Test + void testAgentSessionContextCallBindsSession() { + SessionState> state = + SessionState.>builder().sessionId("ctx-call").build(); + + Session> session = new Session<>(state); + + Session captured = AgentSessionContext.call(session, () -> AgentSessionContext.current()); + + assertSame( + session, captured, "call() must bind session to the thread context and return value"); + } + + @Test + void testAgentSessionContextClearsAfterRun() { + SessionState> state = + SessionState.>builder().sessionId("ctx-clear").build(); + + Session> session = new Session<>(state); + AgentSessionContext.run(session, () -> {}); + + assertNull(AgentSessionContext.current(), "current() must be null after run() completes"); + } + + @Test + void testAgentSessionContextNestedRunRestoresPrior() { + SessionState> outerState = + SessionState.>builder().sessionId("outer").build(); + Session> outer = new Session<>(outerState); + + SessionState> innerState = + SessionState.>builder().sessionId("inner").build(); + Session> inner = new Session<>(innerState); + + AgentSessionContext.run( + outer, + () -> { + assertSame( + outer, AgentSessionContext.current(), "outer session must be bound in outer run"); + AgentSessionContext.run( + inner, + () -> { + assertSame( + inner, + AgentSessionContext.current(), + "inner session must be bound in nested run"); + }); + assertSame( + outer, + AgentSessionContext.current(), + "outer session must be restored after nested run"); + }); + + assertNull(AgentSessionContext.current(), "current must be null after outermost run completes"); + } + + @Test + void testAgentSessionContextNestedCallRestoresPrior() { + SessionState> outerState = + SessionState.>builder().sessionId("outer-call").build(); + Session> outer = new Session<>(outerState); + + SessionState> innerState = + SessionState.>builder().sessionId("inner-call").build(); + Session> inner = new Session<>(innerState); + + String result = + AgentSessionContext.call( + outer, + () -> { + assertSame( + outer, + AgentSessionContext.current(), + "outer session must be bound in outer call"); + String innerResult = + AgentSessionContext.call( + inner, + () -> { + assertSame( + inner, + AgentSessionContext.current(), + "inner session must be bound in nested call"); + return "inner-done"; + }); + assertEquals("inner-done", innerResult, "inner call must return its value"); + assertSame( + outer, + AgentSessionContext.current(), + "outer session must be restored after nested call"); + return "outer-done"; + }); + + assertEquals("outer-done", result, "outer call must return its value"); + assertNull( + AgentSessionContext.current(), "current must be null after outermost call completes"); + } + + @Test + void testCurrentArtifactStoreReturnsNullWhenNoSession() { + assertNull( + AgentSessionContext.currentArtifactStore(), + "currentArtifactStore() must return null when no session is bound"); + } + + @Test + void testCurrentArtifactStoreReturnsSessionAsArtifactStore() { + SessionState> state = + SessionState.>builder().sessionId("ctx-art").build(); + + Session> session = new Session<>(state); + + AtomicReference captured = new AtomicReference<>(); + AgentSessionContext.run( + session, () -> captured.set(AgentSessionContext.currentArtifactStore())); + + assertNotNull(captured.get(), "currentArtifactStore() must return non-null inside run()"); + assertSame(session, captured.get(), "Session must be the ArtifactStore"); + } + + // ---- ArtifactStore interface via Session ---- + + @Test + void testSessionImplementsArtifactStore() { + SessionState> state = SessionState.>builder().build(); + Session> session = new Session<>(state); + + assertTrue(session instanceof ArtifactStore, "Session must implement ArtifactStore"); + + ArtifactStore store = session; + store.addArtifacts(Artifact.builder().name("via-store").build()); + assertEquals(1, store.getArtifacts().size()); + } +} diff --git a/ai/src/test/java/com/google/genkit/ai/agent/StreamEmitterTest.java b/ai/src/test/java/com/google/genkit/ai/agent/StreamEmitterTest.java new file mode 100644 index 000000000..7b9170e42 --- /dev/null +++ b/ai/src/test/java/com/google/genkit/ai/agent/StreamEmitterTest.java @@ -0,0 +1,243 @@ +/* + * 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. + * + * SPDX-License-Identifier: Apache-2.0 + */ + +package com.google.genkit.ai.agent; + +import static org.junit.jupiter.api.Assertions.*; + +import com.fasterxml.jackson.databind.JsonNode; +import com.fasterxml.jackson.databind.ObjectMapper; +import com.google.genkit.ai.agent.internal.StreamEmitter; +import com.google.genkit.core.JsonUtils; +import java.util.ArrayList; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; + +/** + * TDD tests for StreamEmitter: customPatch streaming (whole-doc replace on first patch per turn, + * then incremental diffs) and artifact chunk emission. + */ +class StreamEmitterTest { + + private ObjectMapper mapper; + private List emitted; + private StreamEmitter> emitter; + + @BeforeEach + void setUp() { + mapper = JsonUtils.getObjectMapper(); + emitted = new ArrayList<>(); + emitter = new StreamEmitter<>(emitted::add, mapper); + } + + /** Helper to build a session with a given initial custom state. */ + private Session> sessionWith(Map custom) { + SessionState> state = + SessionState.>builder().custom(custom).build(); + return new Session<>(state); + } + + /** Helper to update the counter field in custom state. */ + private Map counterMap(int value) { + Map m = new HashMap<>(); + m.put("counter", value); + return m; + } + + // ────────────────────────────────────────────────────────────────────────── + // Test 1: Multi-custom-state updates in one turn + // - 3 updateCustom calls (counter 0→1, 1→2, 2→3) + // - 1st chunk: whole-doc replace with {counter:1} + // - 2nd chunk: incremental diff [{op:replace, path:"/counter", value:2}] + // - 3rd chunk: incremental diff [{op:replace, path:"/counter", value:3}] + // ────────────────────────────────────────────────────────────────────────── + @Test + void testThreeCustomUpdatesInOneTurn() { + Session> session = sessionWith(counterMap(0)); + emitter.attach(session); + emitter.beginTurn(); + + // Update 1: counter 0 → 1 + session.updateCustom( + s -> { + Map next = new HashMap<>(s); + next.put("counter", 1); + return next; + }); + + // Update 2: counter 1 → 2 + session.updateCustom( + s -> { + Map next = new HashMap<>(s); + next.put("counter", 2); + return next; + }); + + // Update 3: counter 2 → 3 + session.updateCustom( + s -> { + Map next = new HashMap<>(s); + next.put("counter", 3); + return next; + }); + + assertEquals(3, emitted.size(), "exactly 3 chunks must be emitted"); + + // Chunk 1: whole-doc replace at path "" with value {counter:1} + AgentStreamChunk chunk1 = emitted.get(0); + assertNotNull(chunk1.getCustomPatch(), "chunk1 must have a customPatch"); + assertTrue(chunk1.getCustomPatch().isArray(), "customPatch must be an array"); + assertEquals(1, chunk1.getCustomPatch().size(), "whole-doc replace: exactly 1 op"); + JsonNode op1 = chunk1.getCustomPatch().get(0); + assertEquals("replace", op1.path("op").asText(), "op must be replace"); + assertEquals("", op1.path("path").asText(), "path must be empty string (root)"); + assertEquals(1, op1.path("value").path("counter").asInt(), "value must be {counter:1}"); + + // Chunk 2: incremental diff [{op:replace, path:"/counter", value:2}] + AgentStreamChunk chunk2 = emitted.get(1); + assertNotNull(chunk2.getCustomPatch(), "chunk2 must have a customPatch"); + assertTrue(chunk2.getCustomPatch().isArray(), "customPatch must be an array"); + assertEquals(1, chunk2.getCustomPatch().size(), "incremental diff: exactly 1 op"); + JsonNode op2 = chunk2.getCustomPatch().get(0); + assertEquals("replace", op2.path("op").asText()); + assertEquals("/counter", op2.path("path").asText()); + assertEquals(2, op2.path("value").asInt()); + + // Chunk 3: incremental diff [{op:replace, path:"/counter", value:3}] + AgentStreamChunk chunk3 = emitted.get(2); + assertNotNull(chunk3.getCustomPatch(), "chunk3 must have a customPatch"); + assertEquals(1, chunk3.getCustomPatch().size(), "incremental diff: exactly 1 op"); + JsonNode op3 = chunk3.getCustomPatch().get(0); + assertEquals("replace", op3.path("op").asText()); + assertEquals("/counter", op3.path("path").asText()); + assertEquals(3, op3.path("value").asInt()); + } + + // ────────────────────────────────────────────────────────────────────────── + // Test 2: Second beginTurn() resets to whole-doc replace + // ────────────────────────────────────────────────────────────────────────── + @Test + void testSecondBeginTurnResetsToWholeDocReplace() { + Session> session = sessionWith(counterMap(0)); + emitter.attach(session); + + // Turn 1 + emitter.beginTurn(); + session.updateCustom( + s -> { + Map next = new HashMap<>(s); + next.put("counter", 1); + return next; + }); + // One chunk emitted (whole-doc replace), counter is now at 1 + + // Turn 2 + emitter.beginTurn(); + session.updateCustom( + s -> { + Map next = new HashMap<>(s); + next.put("counter", 2); + return next; + }); + + assertEquals(2, emitted.size(), "one chunk per turn"); + + // Chunk from turn 2 must be a whole-doc replace (not an incremental diff) + AgentStreamChunk chunk2 = emitted.get(1); + assertNotNull(chunk2.getCustomPatch()); + assertEquals(1, chunk2.getCustomPatch().size(), "turn-2 first patch: exactly 1 op"); + JsonNode op = chunk2.getCustomPatch().get(0); + assertEquals("replace", op.path("op").asText()); + assertEquals("", op.path("path").asText(), "second turn must start with whole-doc replace"); + assertEquals(2, op.path("value").path("counter").asInt()); + } + + // ────────────────────────────────────────────────────────────────────────── + // Test 3: Artifact chunk emission + // ────────────────────────────────────────────────────────────────────────── + @Test + void testArtifactChunkEmitted() { + Session> session = sessionWith(null); + emitter.attach(session); + emitter.beginTurn(); + + Artifact artifact = Artifact.builder().name("report").build(); + session.addArtifacts(artifact); + + assertEquals(1, emitted.size(), "exactly one chunk must be emitted for artifact"); + AgentStreamChunk chunk = emitted.get(0); + assertNotNull(chunk.getArtifact(), "chunk must have an artifact"); + assertEquals("report", chunk.getArtifact().getName()); + assertNull(chunk.getCustomPatch(), "artifact chunk must not have customPatch"); + } + + // ────────────────────────────────────────────────────────────────────────── + // Test 4: setSuppressed(true) → nothing emitted + // ────────────────────────────────────────────────────────────────────────── + @Test + void testSuppressedEmitsNothing() { + Session> session = sessionWith(counterMap(0)); + emitter.attach(session); + emitter.beginTurn(); + emitter.setSuppressed(true); + + // updateCustom must not emit + session.updateCustom( + s -> { + Map next = new HashMap<>(s); + next.put("counter", 99); + return next; + }); + + // addArtifacts must not emit + session.addArtifacts(Artifact.builder().name("ignored").build()); + + assertTrue(emitted.isEmpty(), "suppressed emitter must emit nothing"); + } + + // ────────────────────────────────────────────────────────────────────────── + // Test 5: Empty diff (no-op update) must not emit a chunk after the first + // ────────────────────────────────────────────────────────────────────────── + @Test + void testEmptyDiffDoesNotEmitChunk() { + Session> session = sessionWith(counterMap(5)); + emitter.attach(session); + emitter.beginTurn(); + + // First update: whole-doc replace emitted + session.updateCustom( + s -> { + Map next = new HashMap<>(s); + next.put("counter", 10); + return next; + }); + assertEquals(1, emitted.size(), "first update emits one chunk"); + + // Second update: same value as previous → empty diff → no chunk + session.updateCustom( + s -> { + Map next = new HashMap<>(s); + next.put("counter", 10); // same value + return next; + }); + assertEquals(1, emitted.size(), "no-op update must not emit a chunk"); + } +} diff --git a/ai/src/test/java/com/google/genkit/ai/agent/WireTypesSerdeTest.java b/ai/src/test/java/com/google/genkit/ai/agent/WireTypesSerdeTest.java new file mode 100644 index 000000000..a6e9601db --- /dev/null +++ b/ai/src/test/java/com/google/genkit/ai/agent/WireTypesSerdeTest.java @@ -0,0 +1,504 @@ +/* + * 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. + * + * SPDX-License-Identifier: Apache-2.0 + */ + +package com.google.genkit.ai.agent; + +import static org.junit.jupiter.api.Assertions.*; + +import com.fasterxml.jackson.core.type.TypeReference; +import com.fasterxml.jackson.databind.JsonNode; +import com.fasterxml.jackson.databind.ObjectMapper; +import com.google.genkit.ai.Message; +import com.google.genkit.ai.ModelResponseChunk; +import com.google.genkit.ai.Part; +import com.google.genkit.core.JsonUtils; +import java.util.Arrays; +import java.util.Collections; +import java.util.HashMap; +import java.util.Map; +import org.junit.jupiter.api.Test; + +/** TDD tests for agent wire data types serialization/deserialization. */ +class WireTypesSerdeTest { + + private static final ObjectMapper mapper = JsonUtils.getObjectMapper(); + + // ---- SessionState tests ---- + + @Test + void testSessionStateSerializesExactFieldNames() throws Exception { + SessionState> state = + SessionState.>builder() + .sessionId("sess-1") + .messages(Collections.singletonList(Message.user("hello"))) + .build(); + + JsonNode node = mapper.valueToTree(state); + assertTrue(node.has("sessionId"), "must have sessionId"); + assertTrue(node.has("messages"), "must have messages"); + assertFalse(node.has("custom"), "custom must be absent when null"); + assertFalse(node.has("artifacts"), "artifacts must be absent when null"); + } + + @Test + void testSessionStateWithCustomRoundTrip() throws Exception { + Map custom = new HashMap<>(); + custom.put("key", "value"); + + SessionState> state = + SessionState.>builder().sessionId("sess-2").custom(custom).build(); + + String json = mapper.writeValueAsString(state); + assertTrue(json.contains("\"custom\"")); + assertTrue(json.contains("\"sessionId\"")); + + SessionState> deserialized = + mapper.readValue(json, new TypeReference>>() {}); + assertEquals("sess-2", deserialized.getSessionId()); + assertEquals("value", deserialized.getCustom().get("key")); + } + + // ---- AgentInit tests ---- + + @Test + void testAgentInitSerializesExactFieldNames() throws Exception { + AgentInit> init = + AgentInit.>builder().snapshotId("snap-1").sessionId("sess-1").build(); + + JsonNode node = mapper.valueToTree(init); + assertTrue(node.has("snapshotId"), "must have snapshotId"); + assertTrue(node.has("sessionId"), "must have sessionId"); + assertFalse(node.has("state"), "state must be absent when null"); + } + + @Test + void testAgentInitRoundTrip() throws Exception { + SessionState> state = + SessionState.>builder().sessionId("sess-x").build(); + + AgentInit> init = + AgentInit.>builder() + .snapshotId("snap-rt") + .sessionId("sess-rt") + .state(state) + .build(); + + String json = mapper.writeValueAsString(init); + AgentInit> deserialized = + mapper.readValue(json, new TypeReference>>() {}); + assertEquals("snap-rt", deserialized.getSnapshotId()); + assertEquals("sess-rt", deserialized.getSessionId()); + assertNotNull(deserialized.getState()); + assertEquals("sess-x", deserialized.getState().getSessionId()); + } + + // ---- ToolResume tests ---- + + @Test + void testToolResumeSerializesExactFieldNames() throws Exception { + ToolResume resume = + ToolResume.builder().respond(Collections.singletonList(Part.text("response"))).build(); + + JsonNode node = mapper.valueToTree(resume); + assertTrue(node.has("respond"), "must have respond"); + assertFalse(node.has("restart"), "restart must be absent when null"); + } + + @Test + void testToolResumeRoundTrip() throws Exception { + ToolResume resume = + ToolResume.builder() + .respond(Arrays.asList(Part.text("resp1"), Part.text("resp2"))) + .restart(Collections.singletonList(Part.text("restart1"))) + .build(); + + String json = mapper.writeValueAsString(resume); + assertTrue(json.contains("\"respond\"")); + assertTrue(json.contains("\"restart\"")); + + ToolResume deserialized = mapper.readValue(json, ToolResume.class); + assertEquals(2, deserialized.getRespond().size()); + assertEquals(1, deserialized.getRestart().size()); + } + + // ---- AgentInput tests ---- + + @Test + void testAgentInputWithNoDetachOmitsKey() throws Exception { + AgentInput input = AgentInput.builder().message(Message.user("hello")).build(); + + JsonNode node = mapper.valueToTree(input); + assertTrue(node.has("message"), "must have message"); + assertFalse(node.has("detach"), "detach must be absent when null"); + assertFalse(node.has("resume"), "resume must be absent when null"); + } + + @Test + void testAgentInputWithDetachFalseOmitsKey() throws Exception { + AgentInput input = AgentInput.builder().detach(false).build(); + + JsonNode node = mapper.valueToTree(input); + assertFalse(node.has("detach"), "detach must be absent when false"); + } + + @Test + void testAgentInputWithDetachTrueIncludesKey() throws Exception { + AgentInput input = AgentInput.builder().detach(true).build(); + + JsonNode node = mapper.valueToTree(input); + assertTrue(node.has("detach"), "detach must be present when true"); + assertTrue(node.get("detach").asBoolean(), "detach must be true"); + } + + @Test + void testAgentInputWithDetachRoundTrip() throws Exception { + AgentInput input = AgentInput.builder().detach(true).build(); + + String json = mapper.writeValueAsString(input); + assertTrue(json.contains("\"detach\":true")); + + AgentInput deserialized = mapper.readValue(json, AgentInput.class); + assertTrue(deserialized.getDetach()); + } + + // ---- AgentOutput tests ---- + + @Test + void testAgentOutputSerializesExactFieldNames() throws Exception { + AgentOutput> output = + AgentOutput.>builder() + .sessionId("sess-out") + .finishReason(AgentFinishReason.STOP) + .build(); + + JsonNode node = mapper.valueToTree(output); + assertTrue(node.has("sessionId"), "must have sessionId"); + assertTrue(node.has("finishReason"), "must have finishReason"); + assertFalse(node.has("snapshotId"), "snapshotId must be absent when null"); + assertFalse(node.has("state"), "state must be absent when null"); + assertFalse(node.has("message"), "message must be absent when null"); + assertFalse(node.has("artifacts"), "artifacts must be absent when null"); + assertFalse(node.has("error"), "error must be absent when null"); + } + + @Test + void testAgentOutputRoundTrip() throws Exception { + AgentOutput> output = + AgentOutput.>builder() + .sessionId("sess-out-rt") + .snapshotId("snap-out-rt") + .finishReason(AgentFinishReason.STOP) + .message(Message.model("done")) + .build(); + + String json = mapper.writeValueAsString(output); + AgentOutput> deserialized = + mapper.readValue(json, new TypeReference>>() {}); + assertEquals("sess-out-rt", deserialized.getSessionId()); + assertEquals("snap-out-rt", deserialized.getSnapshotId()); + assertEquals(AgentFinishReason.STOP, deserialized.getFinishReason()); + } + + // ---- AgentResult tests ---- + + @Test + void testAgentResultSerializesExactFieldNames() throws Exception { + AgentResult result = AgentResult.builder().finishReason(AgentFinishReason.LENGTH).build(); + + JsonNode node = mapper.valueToTree(result); + assertTrue(node.has("finishReason"), "must have finishReason"); + assertFalse(node.has("message"), "message must be absent when null"); + assertFalse(node.has("artifacts"), "artifacts must be absent when null"); + } + + @Test + void testAgentResultRoundTrip() throws Exception { + AgentResult result = + AgentResult.builder() + .message(Message.model("result")) + .finishReason(AgentFinishReason.STOP) + .build(); + + String json = mapper.writeValueAsString(result); + AgentResult deserialized = mapper.readValue(json, AgentResult.class); + assertEquals(AgentFinishReason.STOP, deserialized.getFinishReason()); + assertNotNull(deserialized.getMessage()); + } + + // ---- TurnEnd tests ---- + + @Test + void testTurnEndSerializesExactFieldNames() throws Exception { + TurnEnd turnEnd = + TurnEnd.builder().snapshotId("snap-turn").finishReason(AgentFinishReason.STOP).build(); + + JsonNode node = mapper.valueToTree(turnEnd); + assertTrue(node.has("snapshotId"), "must have snapshotId"); + assertTrue(node.has("finishReason"), "must have finishReason"); + } + + @Test + void testTurnEndRoundTrip() throws Exception { + TurnEnd turnEnd = + TurnEnd.builder().snapshotId("snap-rt").finishReason(AgentFinishReason.INTERRUPTED).build(); + + String json = mapper.writeValueAsString(turnEnd); + TurnEnd deserialized = mapper.readValue(json, TurnEnd.class); + assertEquals("snap-rt", deserialized.getSnapshotId()); + assertEquals(AgentFinishReason.INTERRUPTED, deserialized.getFinishReason()); + } + + // ---- Artifact tests ---- + + @Test + void testArtifactAlwaysEmitsParts() throws Exception { + Artifact artifact = + Artifact.builder() + .name("my-artifact") + .parts(Collections.singletonList(Part.text("data"))) + .build(); + + JsonNode node = mapper.valueToTree(artifact); + assertTrue(node.has("name"), "must have name"); + assertTrue(node.has("parts"), "parts must always be present"); + assertFalse(node.has("metadata"), "metadata must be absent when null"); + } + + @Test + void testArtifactRoundTrip() throws Exception { + Map meta = new HashMap<>(); + meta.put("version", 1); + + Artifact artifact = + Artifact.builder() + .name("art-rt") + .parts(Arrays.asList(Part.text("p1"), Part.text("p2"))) + .metadata(meta) + .build(); + + String json = mapper.writeValueAsString(artifact); + assertTrue(json.contains("\"parts\"")); + assertTrue(json.contains("\"metadata\"")); + + Artifact deserialized = mapper.readValue(json, Artifact.class); + assertEquals("art-rt", deserialized.getName()); + assertEquals(2, deserialized.getParts().size()); + assertEquals(1, deserialized.getMetadata().get("version")); + } + + // ---- AgentStreamChunk tests ---- + + @Test + void testAgentStreamChunkSerializesExactFieldNames() throws Exception { + ModelResponseChunk modelChunk = ModelResponseChunk.text("chunk text"); + AgentStreamChunk chunk = AgentStreamChunk.builder().modelChunk(modelChunk).build(); + + JsonNode node = mapper.valueToTree(chunk); + assertTrue(node.has("modelChunk"), "must have modelChunk"); + assertFalse(node.has("customPatch"), "customPatch must be absent when null"); + assertFalse(node.has("artifact"), "artifact must be absent when null"); + assertFalse(node.has("turnEnd"), "turnEnd must be absent when null"); + } + + @Test + void testAgentStreamChunkWithCustomPatchRoundTrip() throws Exception { + // customPatch is a JSON array of patch ops + JsonNode patch = mapper.readTree("[{\"op\":\"add\",\"path\":\"/x\",\"value\":1}]"); + + AgentStreamChunk chunk = AgentStreamChunk.builder().customPatch(patch).build(); + + String json = mapper.writeValueAsString(chunk); + assertTrue(json.contains("\"customPatch\"")); + assertTrue(json.contains("\"op\"")); + + AgentStreamChunk deserialized = mapper.readValue(json, AgentStreamChunk.class); + assertNotNull(deserialized.getCustomPatch()); + assertTrue(deserialized.getCustomPatch().isArray()); + assertEquals(1, deserialized.getCustomPatch().size()); + assertEquals("add", deserialized.getCustomPatch().get(0).get("op").asText()); + } + + // ---- RuntimeError tests ---- + + @Test + void testRuntimeErrorSerializesExactFieldNames() throws Exception { + RuntimeError error = + RuntimeError.builder().status("500").message("something went wrong").build(); + + JsonNode node = mapper.valueToTree(error); + assertTrue(node.has("status"), "must have status"); + assertTrue(node.has("message"), "must have message"); + assertFalse(node.has("details"), "details must be absent when null"); + } + + @Test + void testRuntimeErrorRoundTrip() throws Exception { + Map details = new HashMap<>(); + details.put("code", 500); + + RuntimeError error = + RuntimeError.builder().status("500").message("error msg").details(details).build(); + + String json = mapper.writeValueAsString(error); + RuntimeError deserialized = mapper.readValue(json, RuntimeError.class); + assertEquals("500", deserialized.getStatus()); + assertEquals("error msg", deserialized.getMessage()); + assertNotNull(deserialized.getDetails()); + } + + // ---- SessionSnapshot tests ---- + + @Test + void testSessionSnapshotWithOnlyRequiredFieldsOmitsRest() throws Exception { + SessionSnapshot> snapshot = + SessionSnapshot.>builder() + .snapshotId("snap-min") + .createdAt("2025-01-01T00:00:00Z") + .build(); + + JsonNode node = mapper.valueToTree(snapshot); + assertTrue(node.has("snapshotId"), "must have snapshotId"); + assertTrue(node.has("createdAt"), "must have createdAt"); + assertFalse(node.has("sessionId"), "sessionId must be absent when null"); + assertFalse(node.has("parentId"), "parentId must be absent when null"); + assertFalse(node.has("updatedAt"), "updatedAt must be absent when null"); + assertFalse(node.has("heartbeatAt"), "heartbeatAt must be absent when null"); + assertFalse(node.has("status"), "status must be absent when null"); + assertFalse(node.has("finishReason"), "finishReason must be absent when null"); + assertFalse(node.has("error"), "error must be absent when null"); + assertFalse(node.has("state"), "state must be absent when null"); + } + + @Test + void testSessionSnapshotFullRoundTrip() throws Exception { + SessionSnapshot> snapshot = + SessionSnapshot.>builder() + .snapshotId("snap-full") + .sessionId("sess-full") + .parentId("parent-1") + .createdAt("2025-01-01T00:00:00Z") + .updatedAt("2025-01-02T00:00:00Z") + .heartbeatAt("2025-01-02T01:00:00Z") + .status(SnapshotStatus.COMPLETED) + .finishReason(AgentFinishReason.STOP) + .build(); + + String json = mapper.writeValueAsString(snapshot); + assertTrue(json.contains("\"snapshotId\"")); + assertTrue(json.contains("\"sessionId\"")); + assertTrue(json.contains("\"parentId\"")); + assertTrue(json.contains("\"createdAt\"")); + assertTrue(json.contains("\"updatedAt\"")); + assertTrue(json.contains("\"heartbeatAt\"")); + assertTrue(json.contains("\"status\"")); + assertTrue(json.contains("\"finishReason\"")); + + SessionSnapshot> deserialized = + mapper.readValue(json, new TypeReference>>() {}); + assertEquals("snap-full", deserialized.getSnapshotId()); + assertEquals("sess-full", deserialized.getSessionId()); + assertEquals(SnapshotStatus.COMPLETED, deserialized.getStatus()); + assertEquals(AgentFinishReason.STOP, deserialized.getFinishReason()); + } + + // ---- GetSnapshotRequest tests ---- + + @Test + void testGetSnapshotRequestRoundTrip() throws Exception { + GetSnapshotRequest req = + GetSnapshotRequest.builder().snapshotId("snap-req").sessionId("sess-req").build(); + + JsonNode node = mapper.valueToTree(req); + assertTrue(node.has("snapshotId"), "must have snapshotId"); + assertTrue(node.has("sessionId"), "must have sessionId"); + + String json = mapper.writeValueAsString(req); + GetSnapshotRequest deserialized = mapper.readValue(json, GetSnapshotRequest.class); + assertEquals("snap-req", deserialized.getSnapshotId()); + assertEquals("sess-req", deserialized.getSessionId()); + } + + // ---- AgentAbortRequest / AgentAbortResponse tests ---- + + @Test + void testAgentAbortRequestRoundTrip() throws Exception { + AgentAbortRequest req = AgentAbortRequest.builder().snapshotId("snap-abort").build(); + + JsonNode node = mapper.valueToTree(req); + assertTrue(node.has("snapshotId"), "must have snapshotId"); + + String json = mapper.writeValueAsString(req); + AgentAbortRequest deserialized = mapper.readValue(json, AgentAbortRequest.class); + assertEquals("snap-abort", deserialized.getSnapshotId()); + } + + @Test + void testAgentAbortResponseRoundTrip() throws Exception { + AgentAbortResponse resp = + AgentAbortResponse.builder() + .snapshotId("snap-abort-resp") + .status(SnapshotStatus.ABORTED) + .build(); + + JsonNode node = mapper.valueToTree(resp); + assertTrue(node.has("snapshotId"), "must have snapshotId"); + assertTrue(node.has("status"), "must have status"); + + String json = mapper.writeValueAsString(resp); + AgentAbortResponse deserialized = mapper.readValue(json, AgentAbortResponse.class); + assertEquals("snap-abort-resp", deserialized.getSnapshotId()); + assertEquals(SnapshotStatus.ABORTED, deserialized.getStatus()); + } + + // ---- AgentMetadata tests ---- + + @Test + void testAgentMetadataSerializesExactFieldNames() throws Exception { + AgentMetadata metadata = + AgentMetadata.builder().stateManagement("server").abortable(true).build(); + + JsonNode node = mapper.valueToTree(metadata); + assertTrue(node.has("stateManagement"), "must have stateManagement"); + assertTrue(node.has("abortable"), "must have abortable"); + assertFalse(node.has("stateSchema"), "stateSchema must be absent when null"); + } + + @Test + void testAgentMetadataRoundTrip() throws Exception { + Map schema = new HashMap<>(); + schema.put("type", "object"); + + AgentMetadata metadata = + AgentMetadata.builder() + .stateManagement("client") + .abortable(false) + .stateSchema(schema) + .build(); + + String json = mapper.writeValueAsString(metadata); + assertTrue(json.contains("\"stateManagement\"")); + assertTrue(json.contains("\"abortable\"")); + assertTrue(json.contains("\"stateSchema\"")); + + AgentMetadata deserialized = mapper.readValue(json, AgentMetadata.class); + assertEquals("client", deserialized.getStateManagement()); + assertFalse(deserialized.isAbortable()); + assertNotNull(deserialized.getStateSchema()); + assertEquals("object", deserialized.getStateSchema().get("type")); + } +} diff --git a/ai/src/test/java/com/google/genkit/ai/agent/internal/SnapshotShardingTest.java b/ai/src/test/java/com/google/genkit/ai/agent/internal/SnapshotShardingTest.java new file mode 100644 index 000000000..75b8537d1 --- /dev/null +++ b/ai/src/test/java/com/google/genkit/ai/agent/internal/SnapshotShardingTest.java @@ -0,0 +1,192 @@ +/* + * 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. + * + * SPDX-License-Identifier: Apache-2.0 + */ + +package com.google.genkit.ai.agent.internal; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import com.fasterxml.jackson.databind.JsonNode; +import com.fasterxml.jackson.databind.ObjectMapper; +import com.google.genkit.core.GenkitException; +import com.google.genkit.core.JsonUtils; +import com.google.genkit.core.jsonpatch.JsonPatch; +import java.nio.charset.StandardCharsets; +import java.util.ArrayList; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import org.junit.jupiter.api.Test; + +/** + * Unit tests for {@link SnapshotSharding} — the pure sharding / checkpoint-vs-diff / state + * reconstruction helpers shared by the Firestore, DynamoDB, and Cosmos session stores. These run + * with no backend client. + */ +class SnapshotShardingTest { + + private static final ObjectMapper MAPPER = JsonUtils.getObjectMapper(); + + @Test + void shardingRoundTripSmallString() { + String original = "hello world, this is a small JSON-ish payload {\"a\":1}"; + byte[] bytes = original.getBytes(StandardCharsets.UTF_8); + List shards = SnapshotSharding.shardString(original, 8); + String reassembled = SnapshotSharding.reassembleShards(shards); + assertEquals(original, reassembled); + int expectedCount = (bytes.length + 8 - 1) / 8; + assertEquals(expectedCount, shards.size()); + } + + @Test + void shardingRoundTripExactMultiple() { + String original = "abcdefghij"; // 10 bytes + List shards = SnapshotSharding.shardString(original, 5); + assertEquals(2, shards.size()); + assertEquals(original, SnapshotSharding.reassembleShards(shards)); + } + + @Test + void shardRoundTripAcrossMultibyteBoundary() { + // Split a string with multi-byte UTF-8 chars ("é" = 2 bytes) on a byte boundary that lands + // in the middle of a code point; byte-exact reassembly must recover the original. + String s = "aébcdéf"; + List shards = SnapshotSharding.shardString(s, 3); + assertEquals(s, SnapshotSharding.reassembleShards(shards)); + } + + @Test + void shardingRoundTripLargerThanSize() { + StringBuilder sb = new StringBuilder(); + for (int i = 0; i < 5000; i++) { + sb.append("x"); + } + String original = sb.toString(); + int shardSize = 512; + List shards = SnapshotSharding.shardString(original, shardSize); + int expected = (5000 + shardSize - 1) / shardSize; + assertEquals(expected, shards.size()); + assertEquals(original, SnapshotSharding.reassembleShards(shards)); + } + + @Test + void shardingSingleShardWhenSmallerThanSize() { + String original = "tiny"; + List shards = SnapshotSharding.shardString(original, 512 * 1024); + assertEquals(1, shards.size()); + assertEquals(original, SnapshotSharding.reassembleShards(shards)); + } + + @Test + void shardingEmptyStringProducesSingleEmptyShard() { + List shards = SnapshotSharding.shardString("", 8); + assertEquals(1, shards.size()); + assertEquals("", SnapshotSharding.reassembleShards(shards)); + } + + @Test + void decisionCheckpointWhenNoParent() { + assertTrue(SnapshotSharding.shouldCheckpoint(false, 0, 25, 10, 512 * 1024)); + } + + @Test + void decisionCheckpointWhenParentMissing() { + assertTrue(SnapshotSharding.shouldCheckpoint(false, 3, 25, 10, 512 * 1024)); + } + + @Test + void decisionDiffForNormalTurn() { + assertFalse(SnapshotSharding.shouldCheckpoint(true, 3, 25, 10, 512 * 1024)); + } + + @Test + void decisionCheckpointEveryInterval() { + assertTrue(SnapshotSharding.shouldCheckpoint(true, 25, 25, 10, 512 * 1024)); + } + + @Test + void decisionCheckpointWhenDiffExceedsShardSize() { + assertTrue(SnapshotSharding.shouldCheckpoint(true, 3, 25, 600, 512)); + } + + @Test + void reconstructStateFromCheckpointWithNoDiffs() throws Exception { + Map base = new HashMap<>(); + base.put("count", 1); + base.put("name", "alice"); + String checkpointJson = MAPPER.writeValueAsString(base); + + JsonNode result = SnapshotSharding.reconstructState(checkpointJson, new ArrayList<>()); + assertEquals(1, result.get("count").asInt()); + assertEquals("alice", result.get("name").asText()); + } + + @Test + void reconstructStateAppliesDiffsInOrder() throws Exception { + Map base = new HashMap<>(); + base.put("count", 1); + String checkpointJson = MAPPER.writeValueAsString(base); + + JsonNode s1 = MAPPER.valueToTree(Map.of("count", 1)); + JsonNode s2 = MAPPER.valueToTree(Map.of("count", 2)); + JsonNode s3 = MAPPER.valueToTree(Map.of("count", 3, "extra", "y")); + + String patch1 = MAPPER.writeValueAsString(JsonPatch.diff(s1, s2)); + String patch2 = MAPPER.writeValueAsString(JsonPatch.diff(s2, s3)); + + List diffs = List.of(patch1, patch2); + JsonNode result = SnapshotSharding.reconstructState(checkpointJson, diffs); + assertEquals(3, result.get("count").asInt()); + assertEquals("y", result.get("extra").asText()); + } + + @Test + void diffThenReconstructFullRoundTrip() throws Exception { + JsonNode v0 = MAPPER.valueToTree(Map.of("a", 1, "list", List.of(1, 2, 3))); + JsonNode v1 = MAPPER.valueToTree(Map.of("a", 2, "list", List.of(1, 2, 3))); + JsonNode v2 = MAPPER.valueToTree(Map.of("a", 2, "list", List.of(1, 2, 3, 4), "b", "hi")); + JsonNode v3 = MAPPER.valueToTree(Map.of("a", 2, "list", List.of(9), "b", "hi")); + + String checkpoint = MAPPER.writeValueAsString(v0); + List diffs = new ArrayList<>(); + diffs.add(MAPPER.writeValueAsString(JsonPatch.diff(v0, v1))); + diffs.add(MAPPER.writeValueAsString(JsonPatch.diff(v1, v2))); + diffs.add(MAPPER.writeValueAsString(JsonPatch.diff(v2, v3))); + + JsonNode result = SnapshotSharding.reconstructState(checkpoint, diffs); + assertEquals(v3, result); + } + + @Test + void validateIdRejectsSlash() { + assertThrows(GenkitException.class, () -> SnapshotSharding.validateId("foo/bar")); + } + + @Test + void validateIdRejectsNullOrEmpty() { + assertThrows(GenkitException.class, () -> SnapshotSharding.validateId(null)); + assertThrows(GenkitException.class, () -> SnapshotSharding.validateId("")); + } + + @Test + void validateIdAcceptsNormalId() { + SnapshotSharding.validateId("abc-123_DEF"); + } +} diff --git a/ai/src/test/java/com/google/genkit/ai/session/ChatOptionsTest.java b/ai/src/test/java/com/google/genkit/ai/session/ChatOptionsTest.java deleted file mode 100644 index 9c1933845..000000000 --- a/ai/src/test/java/com/google/genkit/ai/session/ChatOptionsTest.java +++ /dev/null @@ -1,285 +0,0 @@ -/* - * 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. - * - * SPDX-License-Identifier: Apache-2.0 - */ - -package com.google.genkit.ai.session; - -import static org.junit.jupiter.api.Assertions.assertEquals; -import static org.junit.jupiter.api.Assertions.assertNotNull; -import static org.junit.jupiter.api.Assertions.assertNull; -import static org.junit.jupiter.api.Assertions.assertSame; -import static org.junit.jupiter.api.Assertions.assertTrue; - -import com.google.genkit.ai.GenerationConfig; -import com.google.genkit.ai.OutputConfig; -import java.util.HashMap; -import java.util.Map; -import org.junit.jupiter.api.Test; - -/** Unit tests for ChatOptions. */ -class ChatOptionsTest { - - @Test - void testDefaultConstructor() { - ChatOptions options = new ChatOptions<>(); - - assertNull(options.getModel()); - assertNull(options.getSystem()); - assertNull(options.getConfig()); - assertNull(options.getTools()); - assertNull(options.getOutput()); - assertNull(options.getContext()); - assertNull(options.getMaxTurns()); - } - - @Test - void testSetAndGetModel() { - ChatOptions options = new ChatOptions<>(); - options.setModel("gpt-4"); - - assertEquals("gpt-4", options.getModel()); - } - - @Test - void testSetAndGetSystem() { - ChatOptions options = new ChatOptions<>(); - options.setSystem("You are a helpful assistant."); - - assertEquals("You are a helpful assistant.", options.getSystem()); - } - - @Test - void testSetAndGetConfig() { - ChatOptions options = new ChatOptions<>(); - GenerationConfig config = GenerationConfig.builder().temperature(0.7).build(); - options.setConfig(config); - - assertSame(config, options.getConfig()); - } - - @Test - void testSetAndGetOutput() { - ChatOptions options = new ChatOptions<>(); - OutputConfig output = new OutputConfig(); - options.setOutput(output); - - assertSame(output, options.getOutput()); - } - - @Test - void testSetAndGetContext() { - ChatOptions options = new ChatOptions<>(); - Map context = new HashMap<>(); - context.put("key1", "value1"); - context.put("key2", 42); - options.setContext(context); - - assertNotNull(options.getContext()); - assertEquals("value1", options.getContext().get("key1")); - assertEquals(42, options.getContext().get("key2")); - } - - @Test - void testSetAndGetMaxTurns() { - ChatOptions options = new ChatOptions<>(); - options.setMaxTurns(10); - - assertEquals(10, options.getMaxTurns()); - } - - @Test - void testBuilderEmpty() { - ChatOptions options = ChatOptions.builder().build(); - - assertNull(options.getModel()); - assertNull(options.getSystem()); - assertNull(options.getConfig()); - assertNull(options.getTools()); - assertNull(options.getOutput()); - assertNull(options.getContext()); - assertNull(options.getMaxTurns()); - } - - @Test - void testBuilderWithModel() { - ChatOptions options = ChatOptions.builder().model("claude-3").build(); - - assertEquals("claude-3", options.getModel()); - } - - @Test - void testBuilderWithSystem() { - ChatOptions options = - ChatOptions.builder().system("You are a coding assistant.").build(); - - assertEquals("You are a coding assistant.", options.getSystem()); - } - - @Test - void testBuilderWithConfig() { - GenerationConfig config = - GenerationConfig.builder().temperature(0.5).maxOutputTokens(100).build(); - - ChatOptions options = ChatOptions.builder().config(config).build(); - - assertSame(config, options.getConfig()); - } - - @Test - void testBuilderWithOutput() { - OutputConfig output = new OutputConfig(); - - ChatOptions options = ChatOptions.builder().output(output).build(); - - assertSame(output, options.getOutput()); - } - - @Test - void testBuilderWithContext() { - Map context = new HashMap<>(); - context.put("userId", "user123"); - - ChatOptions options = ChatOptions.builder().context(context).build(); - - assertNotNull(options.getContext()); - assertEquals("user123", options.getContext().get("userId")); - } - - @Test - void testBuilderWithMaxTurns() { - ChatOptions options = ChatOptions.builder().maxTurns(5).build(); - - assertEquals(5, options.getMaxTurns()); - } - - @Test - void testBuilderWithAllOptions() { - GenerationConfig config = GenerationConfig.builder().temperature(0.7).build(); - OutputConfig output = new OutputConfig(); - Map context = new HashMap<>(); - context.put("key", "value"); - - ChatOptions options = - ChatOptions.builder() - .model("gemini-pro") - .system("You are an expert programmer.") - .config(config) - .output(output) - .context(context) - .maxTurns(20) - .build(); - - assertEquals("gemini-pro", options.getModel()); - assertEquals("You are an expert programmer.", options.getSystem()); - assertSame(config, options.getConfig()); - assertSame(output, options.getOutput()); - assertEquals("value", options.getContext().get("key")); - assertEquals(20, options.getMaxTurns()); - } - - @Test - void testBuilderChaining() { - ChatOptions.Builder builder = ChatOptions.builder(); - - // Test that builder methods return the builder for chaining - assertSame(builder, builder.model("model")); - assertSame(builder, builder.system("system")); - assertSame(builder, builder.config(GenerationConfig.builder().build())); - assertSame(builder, builder.output(new OutputConfig())); - assertSame(builder, builder.context(new HashMap<>())); - assertSame(builder, builder.maxTurns(10)); - } - - @Test - void testLongSystemPrompt() { - String longPrompt = - "You are a helpful assistant. " - + "You should always be polite and professional. " - + "Never provide harmful or misleading information. " - + "If you don't know something, say so. " - + "Always cite your sources when possible."; - - ChatOptions options = ChatOptions.builder().system(longPrompt).build(); - - assertEquals(longPrompt, options.getSystem()); - } - - @Test - void testMultipleBuilds() { - ChatOptions.Builder builder = ChatOptions.builder().model("model1").maxTurns(5); - - ChatOptions options1 = builder.build(); - assertEquals("model1", options1.getModel()); - assertEquals(5, options1.getMaxTurns()); - - // Modify and build again - builder.model("model2").maxTurns(10); - ChatOptions options2 = builder.build(); - - assertEquals("model2", options2.getModel()); - assertEquals(10, options2.getMaxTurns()); - } - - @Test - void testWithComplexState() { - // Test with a custom state type - ChatOptions options = - ChatOptions.builder() - .model("test-model") - .system("Test system prompt") - .maxTurns(15) - .build(); - - assertEquals("test-model", options.getModel()); - assertEquals("Test system prompt", options.getSystem()); - assertEquals(15, options.getMaxTurns()); - } - - @Test - void testEmptyContext() { - ChatOptions options = ChatOptions.builder().context(new HashMap<>()).build(); - - assertNotNull(options.getContext()); - assertTrue(options.getContext().isEmpty()); - } - - @Test - void testZeroMaxTurns() { - ChatOptions options = ChatOptions.builder().maxTurns(0).build(); - - assertEquals(0, options.getMaxTurns()); - } - - /** Simple test state class. */ - static class TestState { - private final String name; - private final int value; - - TestState(String name, int value) { - this.name = name; - this.value = value; - } - - String getName() { - return name; - } - - int getValue() { - return value; - } - } -} diff --git a/ai/src/test/java/com/google/genkit/ai/session/InMemorySessionStoreTest.java b/ai/src/test/java/com/google/genkit/ai/session/InMemorySessionStoreTest.java deleted file mode 100644 index ff986ec7c..000000000 --- a/ai/src/test/java/com/google/genkit/ai/session/InMemorySessionStoreTest.java +++ /dev/null @@ -1,210 +0,0 @@ -/* - * 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. - * - * SPDX-License-Identifier: Apache-2.0 - */ - -package com.google.genkit.ai.session; - -import static org.junit.jupiter.api.Assertions.*; - -import com.google.genkit.ai.Message; -import java.util.ArrayList; -import java.util.List; -import java.util.concurrent.CompletableFuture; -import java.util.concurrent.ExecutionException; -import org.junit.jupiter.api.BeforeEach; -import org.junit.jupiter.api.Test; - -/** Unit tests for InMemorySessionStore. */ -class InMemorySessionStoreTest { - - private InMemorySessionStore store; - - @BeforeEach - void setUp() { - store = new InMemorySessionStore<>(); - } - - @Test - void testSaveAndGet() throws ExecutionException, InterruptedException { - SessionData data = new SessionData<>("session-1", "test-state"); - data.setThread("main", List.of(Message.user("Hello"))); - - store.save("session-1", data).get(); - - SessionData retrieved = store.get("session-1").get(); - - assertNotNull(retrieved); - assertEquals("session-1", retrieved.getId()); - assertEquals("test-state", retrieved.getState()); - assertEquals(1, retrieved.getThread("main").size()); - } - - @Test - void testGetNonexistentSession() throws ExecutionException, InterruptedException { - SessionData result = store.get("nonexistent").get(); - - assertNull(result); - } - - @Test - void testDelete() throws ExecutionException, InterruptedException { - SessionData data = new SessionData<>("to-delete", "state"); - store.save("to-delete", data).get(); - - assertTrue(store.exists("to-delete").get()); - - store.delete("to-delete").get(); - - assertFalse(store.exists("to-delete").get()); - assertNull(store.get("to-delete").get()); - } - - @Test - void testDeleteNonexistentSession() throws ExecutionException, InterruptedException { - // Should not throw - assertDoesNotThrow(() -> store.delete("nonexistent").get()); - } - - @Test - void testExists() throws ExecutionException, InterruptedException { - assertFalse(store.exists("new-session").get()); - - store.save("new-session", new SessionData<>("new-session")).get(); - - assertTrue(store.exists("new-session").get()); - } - - @Test - void testSize() throws ExecutionException, InterruptedException { - assertEquals(0, store.size()); - - store.save("session-1", new SessionData<>("session-1")).get(); - assertEquals(1, store.size()); - - store.save("session-2", new SessionData<>("session-2")).get(); - assertEquals(2, store.size()); - - store.delete("session-1").get(); - assertEquals(1, store.size()); - } - - @Test - void testClear() throws ExecutionException, InterruptedException { - store.save("session-1", new SessionData<>("session-1")).get(); - store.save("session-2", new SessionData<>("session-2")).get(); - store.save("session-3", new SessionData<>("session-3")).get(); - - assertEquals(3, store.size()); - - store.clear(); - - assertEquals(0, store.size()); - assertNull(store.get("session-1").get()); - } - - @Test - void testOverwriteSession() throws ExecutionException, InterruptedException { - SessionData original = new SessionData<>("session", "original-state"); - store.save("session", original).get(); - - SessionData updated = new SessionData<>("session", "updated-state"); - store.save("session", updated).get(); - - SessionData retrieved = store.get("session").get(); - assertEquals("updated-state", retrieved.getState()); - assertEquals(1, store.size()); - } - - @Test - void testMultipleSessions() throws ExecutionException, InterruptedException { - for (int i = 0; i < 10; i++) { - SessionData data = new SessionData<>("session-" + i, "state-" + i); - store.save("session-" + i, data).get(); - } - - assertEquals(10, store.size()); - - for (int i = 0; i < 10; i++) { - SessionData retrieved = store.get("session-" + i).get(); - assertNotNull(retrieved); - assertEquals("state-" + i, retrieved.getState()); - } - } - - @Test - void testWithComplexState() throws ExecutionException, InterruptedException { - // Create a store with complex state type - InMemorySessionStore> complexStore = new InMemorySessionStore<>(); - - List state = new ArrayList<>(); - state.add(1); - state.add(2); - state.add(3); - - SessionData> data = new SessionData<>("complex", state); - complexStore.save("complex", data).get(); - - SessionData> retrieved = complexStore.get("complex").get(); - assertNotNull(retrieved); - assertEquals(3, retrieved.getState().size()); - assertEquals(List.of(1, 2, 3), retrieved.getState()); - } - - @Test - void testAsyncOperations() throws ExecutionException, InterruptedException { - // Test that operations return proper CompletableFutures - CompletableFuture saveFuture = - store.save("async-session", new SessionData<>("async-session")); - assertNotNull(saveFuture); - saveFuture.get(); // Should complete without exception - - CompletableFuture> getFuture = store.get("async-session"); - assertNotNull(getFuture); - SessionData result = getFuture.get(); - assertNotNull(result); - - CompletableFuture existsFuture = store.exists("async-session"); - assertNotNull(existsFuture); - assertTrue(existsFuture.get()); - - CompletableFuture deleteFuture = store.delete("async-session"); - assertNotNull(deleteFuture); - deleteFuture.get(); // Should complete without exception - } - - @Test - void testSessionDataWithThreads() throws ExecutionException, InterruptedException { - SessionData data = new SessionData<>("threaded-session", "state"); - - List mainThread = new ArrayList<>(); - mainThread.add(Message.user("Hello")); - mainThread.add(Message.model("Hi there!")); - data.setThread("main", mainThread); - - List sideThread = new ArrayList<>(); - sideThread.add(Message.user("Different conversation")); - data.setThread("side", sideThread); - - store.save("threaded-session", data).get(); - - SessionData retrieved = store.get("threaded-session").get(); - assertNotNull(retrieved); - assertEquals(2, retrieved.getThread("main").size()); - assertEquals(1, retrieved.getThread("side").size()); - assertEquals("Hello", retrieved.getThread("main").get(0).getText()); - } -} diff --git a/ai/src/test/java/com/google/genkit/ai/session/SessionContextTest.java b/ai/src/test/java/com/google/genkit/ai/session/SessionContextTest.java deleted file mode 100644 index 742751199..000000000 --- a/ai/src/test/java/com/google/genkit/ai/session/SessionContextTest.java +++ /dev/null @@ -1,263 +0,0 @@ -/* - * 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. - * - * SPDX-License-Identifier: Apache-2.0 - */ - -package com.google.genkit.ai.session; - -import static org.junit.jupiter.api.Assertions.*; -import static org.mockito.Mockito.*; - -import com.google.genkit.core.Registry; -import java.util.concurrent.CountDownLatch; -import java.util.concurrent.ExecutorService; -import java.util.concurrent.Executors; -import java.util.concurrent.TimeUnit; -import java.util.concurrent.atomic.AtomicReference; -import org.junit.jupiter.api.AfterEach; -import org.junit.jupiter.api.BeforeEach; -import org.junit.jupiter.api.Test; - -/** Unit tests for SessionContext. */ -class SessionContextTest { - - private Registry mockRegistry; - private InMemorySessionStore store; - - @BeforeEach - void setUp() { - // Ensure clean state before each test - SessionContext.clearSession(); - mockRegistry = mock(Registry.class); - store = new InMemorySessionStore<>(); - } - - @AfterEach - void tearDown() { - // Clean up after each test - SessionContext.clearSession(); - } - - /** Helper to create a test session. */ - private Session createTestSession(String id) { - SessionData sessionData = new SessionData<>(id, "test-state"); - return new Session( - mockRegistry, - store, - sessionData, - () -> null, // We don't need actual - // chat for - // context - // tests - null // No agent registry needed for these tests - ); - } - - @Test - void testCurrentSessionThrowsWhenNotSet() { - assertThrows(SessionContext.SessionException.class, () -> SessionContext.currentSession()); - } - - @Test - void testGetCurrentSessionReturnsNullWhenNotSet() { - assertNull(SessionContext.getCurrentSession()); - } - - @Test - void testHasSessionReturnsFalseWhenNotSet() { - assertFalse(SessionContext.hasSession()); - } - - @Test - void testSetAndGetSession() { - Session session = createTestSession("test-id"); - - SessionContext.setSession(session); - - assertTrue(SessionContext.hasSession()); - assertSame(session, SessionContext.currentSession()); - assertSame(session, SessionContext.getCurrentSession()); - - SessionContext.clearSession(); - } - - @Test - void testRunWithSession() throws Exception { - Session session = createTestSession("test-session-id"); - - AtomicReference> capturedSession = new AtomicReference<>(); - - String result = - SessionContext.runWithSession( - session, - () -> { - capturedSession.set(SessionContext.currentSession()); - return "test-result"; - }); - - assertEquals("test-result", result); - assertSame(session, capturedSession.get()); - // Session should be cleared after runWithSession - assertFalse(SessionContext.hasSession()); - } - - @Test - void testRunWithSessionRestoresPreviousSession() throws Exception { - Session outerSession = createTestSession("outer"); - Session innerSession = createTestSession("inner"); - - SessionContext.setSession(outerSession); - - String result = - SessionContext.runWithSession( - innerSession, - () -> { - assertSame(innerSession, SessionContext.currentSession()); - return "done"; - }); - - // Outer session should be restored - assertSame(outerSession, SessionContext.currentSession()); - - SessionContext.clearSession(); - } - - @Test - void testRunWithSessionHandlesException() { - Session session = createTestSession("error-session"); - - assertThrows( - RuntimeException.class, - () -> { - SessionContext.runWithSession( - session, - () -> { - throw new RuntimeException("Test exception"); - }); - }); - - // Session should be cleared even after exception - assertFalse(SessionContext.hasSession()); - } - - @Test - void testThreadIsolation() throws InterruptedException { - CountDownLatch latch = new CountDownLatch(2); - AtomicReference thread1SessionId = new AtomicReference<>(); - AtomicReference thread2SessionId = new AtomicReference<>(); - - Session session1 = createTestSession("session-1"); - Session session2 = createTestSession("session-2"); - - ExecutorService executor = Executors.newFixedThreadPool(2); - - executor.submit( - () -> { - SessionContext.setSession(session1); - try { - Thread.sleep(50); // Allow time for overlap - } catch (InterruptedException e) { - Thread.currentThread().interrupt(); - } - thread1SessionId.set(SessionContext.currentSession().getId()); - latch.countDown(); - }); - - executor.submit( - () -> { - SessionContext.setSession(session2); - try { - Thread.sleep(50); // Allow time for overlap - } catch (InterruptedException e) { - Thread.currentThread().interrupt(); - } - thread2SessionId.set(SessionContext.currentSession().getId()); - latch.countDown(); - }); - - assertTrue(latch.await(1, TimeUnit.SECONDS)); - executor.shutdown(); - - // Each thread should have its own session - assertEquals("session-1", thread1SessionId.get()); - assertEquals("session-2", thread2SessionId.get()); - } - - @Test - void testClearSession() { - Session session = createTestSession("clear-test"); - - SessionContext.setSession(session); - assertTrue(SessionContext.hasSession()); - - SessionContext.clearSession(); - assertFalse(SessionContext.hasSession()); - } - - @Test - void testNestedRunWithSession() throws Exception { - Session session1 = createTestSession("session-1"); - Session session2 = createTestSession("session-2"); - Session session3 = createTestSession("session-3"); - - AtomicReference level1 = new AtomicReference<>(); - AtomicReference level2 = new AtomicReference<>(); - AtomicReference level3 = new AtomicReference<>(); - AtomicReference afterLevel2 = new AtomicReference<>(); - - SessionContext.runWithSession( - session1, - () -> { - level1.set(SessionContext.currentSession().getId()); - - SessionContext.runWithSession( - session2, - () -> { - level2.set(SessionContext.currentSession().getId()); - - SessionContext.runWithSession( - session3, - () -> { - level3.set(SessionContext.currentSession().getId()); - return null; - }); - - afterLevel2.set(SessionContext.currentSession().getId()); - return null; - }); - - return null; - }); - - assertEquals("session-1", level1.get()); - assertEquals("session-2", level2.get()); - assertEquals("session-3", level3.get()); - assertEquals("session-2", afterLevel2.get()); - } - - @Test - void testRunWithSessionWithNullSession() throws Exception { - String result = - SessionContext.runWithSession( - null, - () -> { - assertFalse(SessionContext.hasSession()); - return "result"; - }); - - assertEquals("result", result); - } -} diff --git a/ai/src/test/java/com/google/genkit/ai/session/SessionDataTest.java b/ai/src/test/java/com/google/genkit/ai/session/SessionDataTest.java deleted file mode 100644 index 84e2f004c..000000000 --- a/ai/src/test/java/com/google/genkit/ai/session/SessionDataTest.java +++ /dev/null @@ -1,216 +0,0 @@ -/* - * 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. - * - * SPDX-License-Identifier: Apache-2.0 - */ - -package com.google.genkit.ai.session; - -import static org.junit.jupiter.api.Assertions.*; - -import com.google.genkit.ai.Message; -import java.util.ArrayList; -import java.util.HashMap; -import java.util.List; -import java.util.Map; -import org.junit.jupiter.api.Test; - -/** Unit tests for SessionData. */ -class SessionDataTest { - - @Test - void testDefaultConstructor() { - SessionData data = new SessionData<>(); - - assertNull(data.getId()); - assertNull(data.getState()); - assertNotNull(data.getThreads()); - assertTrue(data.getThreads().isEmpty()); - } - - @Test - void testConstructorWithId() { - SessionData data = new SessionData<>("test-session-123"); - - assertEquals("test-session-123", data.getId()); - assertNull(data.getState()); - assertNotNull(data.getThreads()); - assertTrue(data.getThreads().isEmpty()); - } - - @Test - void testConstructorWithIdAndState() { - SessionData data = new SessionData<>("test-session", "initial-state"); - - assertEquals("test-session", data.getId()); - assertEquals("initial-state", data.getState()); - assertNotNull(data.getThreads()); - assertTrue(data.getThreads().isEmpty()); - } - - @Test - void testSetAndGetId() { - SessionData data = new SessionData<>(); - data.setId("my-session-id"); - - assertEquals("my-session-id", data.getId()); - } - - @Test - void testSetAndGetState() { - SessionData data = new SessionData<>(); - data.setState(42); - - assertEquals(42, data.getState()); - } - - @Test - void testSetAndGetThreads() { - SessionData data = new SessionData<>(); - - Map> threads = new HashMap<>(); - List messages = new ArrayList<>(); - messages.add(Message.user("Hello")); - threads.put("main", messages); - - data.setThreads(threads); - - assertEquals(1, data.getThreads().size()); - assertTrue(data.getThreads().containsKey("main")); - assertEquals(1, data.getThreads().get("main").size()); - } - - @Test - void testGetThread() { - SessionData data = new SessionData<>(); - - List messages = new ArrayList<>(); - messages.add(Message.user("Test message")); - data.setThread("test-thread", messages); - - List retrieved = data.getThread("test-thread"); - - assertNotNull(retrieved); - assertEquals(1, retrieved.size()); - assertEquals("Test message", retrieved.get(0).getText()); - } - - @Test - void testGetThreadReturnsNullForNonexistent() { - SessionData data = new SessionData<>(); - - assertNull(data.getThread("nonexistent")); - } - - @Test - void testGetOrCreateThread() { - SessionData data = new SessionData<>(); - - // First call should create the thread - List thread1 = data.getOrCreateThread("new-thread"); - assertNotNull(thread1); - assertTrue(thread1.isEmpty()); - - // Add a message - thread1.add(Message.user("Hello")); - - // Second call should return the same thread - List thread2 = data.getOrCreateThread("new-thread"); - assertEquals(1, thread2.size()); - assertEquals("Hello", thread2.get(0).getText()); - } - - @Test - void testSetThread() { - SessionData data = new SessionData<>(); - - List messages = new ArrayList<>(); - messages.add(Message.user("First")); - messages.add(Message.model("Second")); - - data.setThread("conversation", messages); - - List retrieved = data.getThread("conversation"); - assertEquals(2, retrieved.size()); - assertEquals("First", retrieved.get(0).getText()); - assertEquals("Second", retrieved.get(1).getText()); - } - - @Test - void testSetThreadCreatesDefensiveCopy() { - SessionData data = new SessionData<>(); - - List messages = new ArrayList<>(); - messages.add(Message.user("Original")); - data.setThread("thread", messages); - - // Modify original list - messages.add(Message.model("Added after")); - - // The stored thread should not be affected - List retrieved = data.getThread("thread"); - assertEquals(1, retrieved.size()); - } - - @Test - void testBuilder() { - SessionData data = - SessionData.builder().id("builder-session").state("builder-state").build(); - - assertEquals("builder-session", data.getId()); - assertEquals("builder-state", data.getState()); - assertNotNull(data.getThreads()); - } - - @Test - void testBuilderWithThreads() { - Map> threads = new HashMap<>(); - List mainThread = new ArrayList<>(); - mainThread.add(Message.user("Hello")); - threads.put("main", mainThread); - - SessionData data = SessionData.builder().id("session").threads(threads).build(); - - assertEquals(1, data.getThreads().size()); - assertTrue(data.getThreads().containsKey("main")); - } - - @Test - void testBuilderAddThread() { - List messages = new ArrayList<>(); - messages.add(Message.system("System prompt")); - - SessionData data = - SessionData.builder().id("session").thread("custom", messages).build(); - - assertNotNull(data.getThread("custom")); - assertEquals(1, data.getThread("custom").size()); - } - - @Test - void testWithComplexState() { - // Test with a complex state object - Map complexState = new HashMap<>(); - complexState.put("userName", "Alice"); - complexState.put("preferences", Map.of("theme", "dark", "language", "en")); - complexState.put("messageCount", 5); - - SessionData> data = new SessionData<>("complex-session", complexState); - - assertEquals("complex-session", data.getId()); - assertEquals("Alice", data.getState().get("userName")); - assertEquals(5, data.getState().get("messageCount")); - } -} diff --git a/ai/src/test/java/com/google/genkit/ai/session/SessionOptionsTest.java b/ai/src/test/java/com/google/genkit/ai/session/SessionOptionsTest.java deleted file mode 100644 index 27f552775..000000000 --- a/ai/src/test/java/com/google/genkit/ai/session/SessionOptionsTest.java +++ /dev/null @@ -1,156 +0,0 @@ -/* - * 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. - * - * SPDX-License-Identifier: Apache-2.0 - */ - -package com.google.genkit.ai.session; - -import static org.junit.jupiter.api.Assertions.*; - -import org.junit.jupiter.api.Test; - -/** Unit tests for SessionOptions. */ -class SessionOptionsTest { - - @Test - void testDefaultConstructor() { - SessionOptions options = new SessionOptions<>(); - - assertNull(options.getStore()); - assertNull(options.getInitialState()); - assertNull(options.getSessionId()); - } - - @Test - void testSetAndGetStore() { - SessionOptions options = new SessionOptions<>(); - InMemorySessionStore store = new InMemorySessionStore<>(); - - options.setStore(store); - - assertSame(store, options.getStore()); - } - - @Test - void testSetAndGetInitialState() { - SessionOptions options = new SessionOptions<>(); - options.setInitialState(42); - - assertEquals(42, options.getInitialState()); - } - - @Test - void testSetAndGetSessionId() { - SessionOptions options = new SessionOptions<>(); - options.setSessionId("custom-session-id"); - - assertEquals("custom-session-id", options.getSessionId()); - } - - @Test - void testBuilderEmpty() { - SessionOptions options = SessionOptions.builder().build(); - - assertNull(options.getStore()); - assertNull(options.getInitialState()); - assertNull(options.getSessionId()); - } - - @Test - void testBuilderWithStore() { - InMemorySessionStore store = new InMemorySessionStore<>(); - - SessionOptions options = SessionOptions.builder().store(store).build(); - - assertSame(store, options.getStore()); - } - - @Test - void testBuilderWithInitialState() { - SessionOptions options = - SessionOptions.builder().initialState("initial-value").build(); - - assertEquals("initial-value", options.getInitialState()); - } - - @Test - void testBuilderWithSessionId() { - SessionOptions options = - SessionOptions.builder().sessionId("my-session-123").build(); - - assertEquals("my-session-123", options.getSessionId()); - } - - @Test - void testBuilderWithAllOptions() { - InMemorySessionStore store = new InMemorySessionStore<>(); - - SessionOptions options = - SessionOptions.builder() - .store(store) - .initialState("test-state") - .sessionId("session-456") - .build(); - - assertSame(store, options.getStore()); - assertEquals("test-state", options.getInitialState()); - assertEquals("session-456", options.getSessionId()); - } - - @Test - void testBuilderChaining() { - SessionOptions.Builder builder = SessionOptions.builder(); - - // Test that builder methods return the builder for chaining - assertSame(builder, builder.store(new InMemorySessionStore<>())); - assertSame(builder, builder.initialState("state")); - assertSame(builder, builder.sessionId("id")); - } - - @Test - void testWithComplexState() { - // Test with a custom state class - TestState state = new TestState("Alice", 25); - - SessionOptions options = - SessionOptions.builder() - .initialState(state) - .sessionId("complex-state-session") - .build(); - - assertEquals("Alice", options.getInitialState().getName()); - assertEquals(25, options.getInitialState().getAge()); - } - - /** Simple test state class. */ - static class TestState { - private final String name; - private final int age; - - TestState(String name, int age) { - this.name = name; - this.age = age; - } - - String getName() { - return name; - } - - int getAge() { - return age; - } - } -} diff --git a/core/src/main/java/com/google/genkit/core/ActionContext.java b/core/src/main/java/com/google/genkit/core/ActionContext.java index 0baa380a1..39b63ba6d 100644 --- a/core/src/main/java/com/google/genkit/core/ActionContext.java +++ b/core/src/main/java/com/google/genkit/core/ActionContext.java @@ -19,6 +19,7 @@ package com.google.genkit.core; import com.google.genkit.core.tracing.SpanContext; +import java.util.Map; /** * ActionContext provides context for action execution including tracing and flow information. It is @@ -32,6 +33,9 @@ public class ActionContext { private final Registry registry; private final String sessionId; private final String threadName; + private final Map context; + private final Object resumed; + private final Object originalInput; /** * Creates a new ActionContext. @@ -42,6 +46,7 @@ public class ActionContext { * @param registry the Genkit registry * @param sessionId the session ID for multi-turn conversations * @param threadName the thread name for grouping related requests + * @param context the request-scoped user context (e.g. {@code {"auth": {...}}}), may be null */ public ActionContext( SpanContext spanContext, @@ -49,13 +54,66 @@ public ActionContext( String spanPath, Registry registry, String sessionId, - String threadName) { + String threadName, + Map context) { + this(spanContext, flowName, spanPath, registry, sessionId, threadName, context, null, null); + } + + /** + * Creates a new ActionContext including resume-awareness fields. + * + * @param spanContext the tracing span context, may be null + * @param flowName the name of the enclosing flow, may be null + * @param spanPath the current span path for tracing + * @param registry the Genkit registry + * @param sessionId the session ID for multi-turn conversations + * @param threadName the thread name for grouping related requests + * @param context the request-scoped user context (e.g. {@code {"auth": {...}}}), may be null + * @param resumed the resume metadata attached when this action is being re-invoked after an + * interrupt/restart (mirrors JS {@code ToolRunOptions.resumed} / Go {@code + * ToolContext.Resumed}); {@code null} on a normal (non-resumed) invocation + * @param originalInput the tool request's original input (before any restart-replaced input); + * mirrors Go {@code ToolContext.OriginalInput}; {@code null} when not resuming + */ + public ActionContext( + SpanContext spanContext, + String flowName, + String spanPath, + Registry registry, + String sessionId, + String threadName, + Map context, + Object resumed, + Object originalInput) { this.spanContext = spanContext; this.flowName = flowName; this.spanPath = spanPath; this.registry = registry; this.sessionId = sessionId; this.threadName = threadName; + this.context = context; + this.resumed = resumed; + this.originalInput = originalInput; + } + + /** + * Creates a new ActionContext. + * + * @param spanContext the tracing span context, may be null + * @param flowName the name of the enclosing flow, may be null + * @param spanPath the current span path for tracing + * @param registry the Genkit registry + * @param sessionId the session ID for multi-turn conversations + * @param threadName the thread name for grouping related requests + */ + public ActionContext( + SpanContext spanContext, + String flowName, + String spanPath, + Registry registry, + String sessionId, + String threadName) { + this(spanContext, flowName, spanPath, registry, sessionId, threadName, null); } /** @@ -145,6 +203,54 @@ public String getThreadName() { return threadName; } + /** + * Returns the request-scoped user context. + * + *

This is the {@code context} object injected by callers (such as the Dev UI "Execution + * context" panel or the reflection/serving layers), e.g. {@code {"auth": {"user": "alice"}}}. It + * is threaded through the run so tools and flows can read it via {@link #getContext()}. + * + * @return the user context map, or null if not set + */ + public Map getContext() { + return context; + } + + /** + * Returns the resume metadata attached when this action is being re-invoked after an + * interrupt/restart, or {@code null} on a normal invocation. + * + *

Mirrors JS {@code ToolRunOptions.resumed} and Go {@code ToolContext.Resumed}: a + * restart-aware tool can inspect this to distinguish a fresh call from a resumed one and read the + * client-supplied approval/confirmation payload. + * + * @return the resumed metadata value, or {@code null} if this is not a resumed invocation + */ + public Object getResumed() { + return resumed; + } + + /** + * Returns {@code true} if this action is being re-invoked after an interrupt/restart (i.e. {@link + * #getResumed()} is non-null). + * + * @return whether this is a resumed invocation + */ + public boolean isResumed() { + return resumed != null; + } + + /** + * Returns the tool request's original input (before any restart-replaced input), mirroring Go + * {@code ToolContext.OriginalInput}. {@code null} when not resuming or when the original input + * was not preserved. + * + * @return the original input, or {@code null} + */ + public Object getOriginalInput() { + return originalInput; + } + /** * Creates a new ActionContext with a different flow name. * @@ -153,7 +259,15 @@ public String getThreadName() { */ public ActionContext withFlowName(String flowName) { return new ActionContext( - this.spanContext, flowName, this.spanPath, this.registry, this.sessionId, this.threadName); + this.spanContext, + flowName, + this.spanPath, + this.registry, + this.sessionId, + this.threadName, + this.context, + this.resumed, + this.originalInput); } /** @@ -164,7 +278,15 @@ public ActionContext withFlowName(String flowName) { */ public ActionContext withSpanContext(SpanContext spanContext) { return new ActionContext( - spanContext, this.flowName, this.spanPath, this.registry, this.sessionId, this.threadName); + spanContext, + this.flowName, + this.spanPath, + this.registry, + this.sessionId, + this.threadName, + this.context, + this.resumed, + this.originalInput); } /** @@ -175,7 +297,15 @@ public ActionContext withSpanContext(SpanContext spanContext) { */ public ActionContext withSpanPath(String spanPath) { return new ActionContext( - this.spanContext, this.flowName, spanPath, this.registry, this.sessionId, this.threadName); + this.spanContext, + this.flowName, + spanPath, + this.registry, + this.sessionId, + this.threadName, + this.context, + this.resumed, + this.originalInput); } /** @@ -186,7 +316,15 @@ public ActionContext withSpanPath(String spanPath) { */ public ActionContext withSessionId(String sessionId) { return new ActionContext( - this.spanContext, this.flowName, this.spanPath, this.registry, sessionId, this.threadName); + this.spanContext, + this.flowName, + this.spanPath, + this.registry, + sessionId, + this.threadName, + this.context, + this.resumed, + this.originalInput); } /** @@ -197,7 +335,55 @@ public ActionContext withSessionId(String sessionId) { */ public ActionContext withThreadName(String threadName) { return new ActionContext( - this.spanContext, this.flowName, this.spanPath, this.registry, this.sessionId, threadName); + this.spanContext, + this.flowName, + this.spanPath, + this.registry, + this.sessionId, + threadName, + this.context, + this.resumed, + this.originalInput); + } + + /** + * Creates a new ActionContext with the given request-scoped user context. + * + * @param context the user context map (e.g. {@code {"auth": {...}}}), may be null + * @return a new ActionContext with the updated user context + */ + public ActionContext withContext(Map context) { + return new ActionContext( + this.spanContext, + this.flowName, + this.spanPath, + this.registry, + this.sessionId, + this.threadName, + context, + this.resumed, + this.originalInput); + } + + /** + * Creates a new ActionContext carrying resume-awareness for a restarted tool call. + * + * @param resumed the resume metadata value (from the restart part's {@code metadata.resumed}); + * may be {@code null} + * @param originalInput the tool request's original input; may be {@code null} + * @return a new ActionContext with the resume-awareness fields set + */ + public ActionContext withResumed(Object resumed, Object originalInput) { + return new ActionContext( + this.spanContext, + this.flowName, + this.spanPath, + this.registry, + this.sessionId, + this.threadName, + this.context, + resumed, + originalInput); } /** @@ -217,6 +403,7 @@ public static class Builder { private Registry registry; private String sessionId; private String threadName; + private Map context; public Builder spanContext(SpanContext spanContext) { this.spanContext = spanContext; @@ -248,11 +435,17 @@ public Builder threadName(String threadName) { return this; } + public Builder context(Map context) { + this.context = context; + return this; + } + public ActionContext build() { if (registry == null) { throw new IllegalStateException("registry is required"); } - return new ActionContext(spanContext, flowName, spanPath, registry, sessionId, threadName); + return new ActionContext( + spanContext, flowName, spanPath, registry, sessionId, threadName, context); } } } diff --git a/core/src/main/java/com/google/genkit/core/ActionType.java b/core/src/main/java/com/google/genkit/core/ActionType.java index 0d66034cf..a4a3e4687 100644 --- a/core/src/main/java/com/google/genkit/core/ActionType.java +++ b/core/src/main/java/com/google/genkit/core/ActionType.java @@ -76,7 +76,16 @@ public enum ActionType { CHECK_OPERATION("check-operation"), /** An action for cancelling operations. */ - CANCEL_OPERATION("cancel-operation"); + CANCEL_OPERATION("cancel-operation"), + + /** An agent action representing a bidirectional agent interaction. */ + AGENT("agent"), + + /** An agent-snapshot action for capturing agent state. */ + AGENT_SNAPSHOT("agent-snapshot"), + + /** An agent-abort action for stopping agent execution. */ + AGENT_ABORT("agent-abort"); private final String value; diff --git a/core/src/main/java/com/google/genkit/core/BidiAction.java b/core/src/main/java/com/google/genkit/core/BidiAction.java new file mode 100644 index 000000000..9581f6ce4 --- /dev/null +++ b/core/src/main/java/com/google/genkit/core/BidiAction.java @@ -0,0 +1,128 @@ +/* + * 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. + * + * SPDX-License-Identifier: Apache-2.0 + */ + +package com.google.genkit.core; + +import com.fasterxml.jackson.databind.JsonNode; +import java.util.function.Consumer; + +/** + * A bidirectional streaming action. + * + *

A {@code BidiAction} accepts a one-time session-init value ({@code Init}), a stream of + * per-turn inputs ({@link InputSource}{@code }), streams chunk values of type {@code S} via a + * callback, and returns a final output of type {@code O}. + * + *

Callers that only need a single input can use the inherited {@link Action#run} methods, which + * wrap the single input in a {@link BufferedInputSource} and invoke the bidi handler transparently. + * + * @param per-turn input type + * @param final output type + * @param stream chunk type + * @param session-init type (pass {@code Void} / {@code null} when not needed) + */ +public interface BidiAction extends Action { + + /** + * Runs the action in bidirectional-streaming mode. + * + * @param ctx the action context + * @param init the one-time session-init value; may be {@code null} + * @param inputs the stream of per-turn inputs + * @param streamCallback callback invoked for each emitted chunk; may be {@code null} + * @return the final output + * @throws GenkitException if execution fails + */ + O runBidi(ActionContext ctx, Init init, InputSource inputs, Consumer streamCallback) + throws GenkitException; + + /** + * Runs the action in bidirectional-streaming mode with JSON-typed arguments. + * + *

Deserializes {@code init} and each element from {@code inputs} before handing them to the + * typed handler, and serializes all chunks and the final result back to {@link JsonNode}. + * + * @param ctx the action context + * @param init the one-time session-init as a {@link JsonNode}; may be {@code null} + * @param inputs the stream of per-turn inputs as {@link JsonNode} values + * @param streamCallback callback invoked for each emitted chunk serialized to {@link JsonNode}; + * may be {@code null} + * @return the final output serialized to {@link JsonNode} + * @throws GenkitException if execution fails + */ + JsonNode runBidiJson( + ActionContext ctx, + JsonNode init, + InputSource inputs, + Consumer streamCallback) + throws GenkitException; + + /** + * Runs the action in bidirectional-streaming mode with JSON-typed arguments and telemetry. + * + *

Behaves like {@link #runBidiJson} (deserializing {@code init} and each element from {@code + * inputs}, serializing all chunks and the final result back to {@link JsonNode}) but additionally + * creates a new tracing span and captures the resulting trace/span IDs, returning them in an + * {@link ActionRunResult}. Unlike the unary {@link Action#runJsonWithTelemetry}, this threads the + * real {@code init} and the full stream of {@code inputs} through to the handler. + * + * @param ctx the action context + * @param init the one-time session-init as a {@link JsonNode}; may be {@code null} + * @param inputs the stream of per-turn inputs as {@link JsonNode} values + * @param streamCallback callback invoked for each emitted chunk serialized to {@link JsonNode}; + * may be {@code null} + * @return the final output serialized to {@link JsonNode} together with the captured trace/span + * IDs + * @throws GenkitException if execution fails + */ + ActionRunResult runBidiJsonWithTelemetry( + ActionContext ctx, + JsonNode init, + InputSource inputs, + Consumer streamCallback) + throws GenkitException; + + // ------------------------------------------------------------------------- + // Nested functional interface for the handler + // ------------------------------------------------------------------------- + + /** + * Functional interface that implements the bidirectional streaming logic. + * + * @param per-turn input type + * @param final output type + * @param stream chunk type + * @param session-init type + */ + @FunctionalInterface + interface BidiHandler { + + /** + * Handles one invocation of the bidirectional action. + * + * @param ctx the action context + * @param init the one-time session-init value; may be {@code null} + * @param inputs the stream of per-turn inputs + * @param streamCallback callback for emitting chunks + * @return the final output + * @throws Exception if handling fails + */ + O handle(ActionContext ctx, Init init, InputSource inputs, Consumer streamCallback) + throws Exception; + } +} diff --git a/core/src/main/java/com/google/genkit/core/BidiActionImpl.java b/core/src/main/java/com/google/genkit/core/BidiActionImpl.java new file mode 100644 index 000000000..84fb3758c --- /dev/null +++ b/core/src/main/java/com/google/genkit/core/BidiActionImpl.java @@ -0,0 +1,523 @@ +/* + * 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. + * + * SPDX-License-Identifier: Apache-2.0 + */ + +package com.google.genkit.core; + +import com.fasterxml.jackson.databind.JsonNode; +import com.fasterxml.jackson.databind.ObjectMapper; +import com.google.genkit.core.tracing.SpanMetadata; +import com.google.genkit.core.tracing.Tracer; +import java.util.HashMap; +import java.util.Map; +import java.util.Optional; +import java.util.function.Consumer; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +/** + * Concrete implementation of {@link BidiAction} built from a name, type parameters, optional + * metadata, and a {@link BidiAction.BidiHandler}. + * + *

Key invariants: + * + *

    + *
  • {@link #getType()} always returns {@link ActionType#AGENT}. + *
  • {@link #getMetadata()} always contains {@code bidi=true}; the caller's metadata map is + * never mutated. + *
  • {@link #getDesc()} returns an {@link ActionDesc} with key {@code /agent/}. + *
  • The inherited unary {@link Action#run} methods adapt to bidi by wrapping the single input + * in a {@link BufferedInputSource}. + *
+ * + * @param per-turn input type + * @param final output type + * @param stream chunk type + * @param session-init type + */ +public final class BidiActionImpl implements BidiAction { + + private static final Logger logger = LoggerFactory.getLogger(BidiActionImpl.class); + private static final ObjectMapper MAPPER = JsonUtils.getObjectMapper(); + + private final ActionDesc desc; + private final BidiAction.BidiHandler handler; + private final Class inputClass; + private final Class outputClass; + private final Class streamClass; + private final Class initClass; + private final Map metadata; // the merged metadata (includes bidi=true) + + // ------------------------------------------------------------------------- + // Private constructor – use Builder + // ------------------------------------------------------------------------- + + private BidiActionImpl( + String name, + Class inputClass, + Class outputClass, + Class streamClass, + Class initClass, + Map inputSchema, + Map outputSchema, + Map callerMetadata, + BidiAction.BidiHandler handler) { + + this.inputClass = inputClass; + this.outputClass = outputClass; + this.streamClass = streamClass; + this.initClass = initClass; + this.handler = handler; + + // Merge caller metadata + bidi=true without mutating the caller's map + Map merged = new HashMap<>(); + if (callerMetadata != null) { + merged.putAll(callerMetadata); + } + merged.put("bidi", Boolean.TRUE); + this.metadata = merged; + + // Extract description from metadata if present (mirrors ActionDef convention) + String description = null; + if (merged.get("description") instanceof String) { + description = (String) merged.get("description"); + } + + // Generate schemas when not explicitly provided + Map actualInputSchema = inputSchema; + if (actualInputSchema == null && inputClass != null && inputClass != Void.class) { + actualInputSchema = SchemaUtils.inferSchema(inputClass); + } + Map actualOutputSchema = outputSchema; + if (actualOutputSchema == null && outputClass != null && outputClass != Void.class) { + actualOutputSchema = SchemaUtils.inferSchema(outputClass); + } + + this.desc = + ActionDesc.builder() + .type(ActionType.AGENT) + .name(name) + .description(description) + .inputSchema(actualInputSchema) + .outputSchema(actualOutputSchema) + .metadata(merged) + .build(); + } + + // ------------------------------------------------------------------------- + // Action interface + // ------------------------------------------------------------------------- + + @Override + public String getName() { + return desc.getName(); + } + + @Override + public ActionType getType() { + return ActionType.AGENT; + } + + @Override + public ActionDesc getDesc() { + return desc; + } + + @Override + public Map getInputSchema() { + return desc.getInputSchema(); + } + + @Override + public Map getOutputSchema() { + return desc.getOutputSchema(); + } + + @Override + public Map getMetadata() { + return metadata; + } + + // ------------------------------------------------------------------------- + // Unary adaptation + // ------------------------------------------------------------------------- + + /** + * Unary adaptation: wraps {@code input} in a {@link BufferedInputSource} (with {@code null} init) + * and delegates to the bidi handler. + */ + @Override + public O run(ActionContext ctx, I input) throws GenkitException { + return run(ctx, input, null); + } + + /** + * Unary adaptation with streaming: wraps {@code input} in a {@link BufferedInputSource} (with + * {@code null} init) and delegates to the bidi handler. + */ + @Override + public O run(ActionContext ctx, I input, Consumer streamCallback) throws GenkitException { + logger.debug("BidiActionImpl.run (unary): name={}, input={}", getName(), input); + + SpanMetadata spanMetadata = + SpanMetadata.builder() + .name(desc.getName()) + .type(ActionType.AGENT.getValue()) + .subtype(ActionType.AGENT.getValue()) + .build(); + + return Tracer.runInNewSpan( + ctx, + spanMetadata, + input, + (spanCtx, in) -> { + BufferedInputSource source = new BufferedInputSource<>(); + if (in != null) { + source.offer(in); + } + source.end(); + try { + return handler.handle(ctx.withSpanContext(spanCtx), null, source, streamCallback); + } catch (GenkitException ge) { + throw ge; + } catch (Exception e) { + throw new GenkitException("BidiAction execution failed: " + e.getMessage(), e); + } + }); + } + + @Override + public JsonNode runJson(ActionContext ctx, JsonNode input, Consumer streamCallback) + throws GenkitException { + try { + I typedInput = null; + if (inputClass != null && inputClass != Void.class && input != null) { + typedInput = MAPPER.treeToValue(input, inputClass); + } + + Consumer typedCallback = buildTypedCallback(streamCallback); + O result = run(ctx, typedInput, typedCallback); + return result != null ? MAPPER.valueToTree(result) : null; + } catch (GenkitException ge) { + throw ge; + } catch (Exception e) { + throw new GenkitException("JSON BidiAction execution failed: " + e.getMessage(), e); + } + } + + @Override + public ActionRunResult runJsonWithTelemetry( + ActionContext ctx, JsonNode input, Consumer streamCallback) throws GenkitException { + + final String[] capturedTraceInfo = new String[2]; // [traceId, spanId] + + SpanMetadata spanMetadata = + SpanMetadata.builder() + .name(desc.getName()) + .type(ActionType.AGENT.getValue()) + .subtype(ActionType.AGENT.getValue()) + .build(); + + try { + I typedInput = null; + if (inputClass != null && inputClass != Void.class && input != null) { + typedInput = MAPPER.treeToValue(input, inputClass); + } + final I finalInput = typedInput; + final Consumer typedCallback = buildTypedCallback(streamCallback); + + O result = + Tracer.runInNewSpan( + ctx, + spanMetadata, + finalInput, + (spanCtx, in) -> { + capturedTraceInfo[0] = spanCtx.getTraceId(); + capturedTraceInfo[1] = spanCtx.getSpanId(); + + BufferedInputSource source = new BufferedInputSource<>(); + if (in != null) { + source.offer(in); + } + source.end(); + try { + return handler.handle(ctx.withSpanContext(spanCtx), null, source, typedCallback); + } catch (GenkitException ge) { + throw ge; + } catch (Exception e) { + throw new GenkitException("BidiAction execution failed: " + e.getMessage(), e); + } + }); + + JsonNode jsonResult = result != null ? MAPPER.valueToTree(result) : null; + return new ActionRunResult<>(jsonResult, capturedTraceInfo[0], capturedTraceInfo[1]); + } catch (GenkitException ge) { + throw ge; + } catch (Exception e) { + throw new GenkitException("JSON BidiAction execution failed: " + e.getMessage(), e); + } + } + + // ------------------------------------------------------------------------- + // BidiAction interface + // ------------------------------------------------------------------------- + + @Override + public O runBidi(ActionContext ctx, Init init, InputSource inputs, Consumer streamCallback) + throws GenkitException { + logger.debug("BidiActionImpl.runBidi: name={}", getName()); + try { + return handler.handle(ctx, init, inputs, streamCallback); + } catch (GenkitException ge) { + throw ge; + } catch (Exception e) { + throw new GenkitException("BidiAction execution failed: " + e.getMessage(), e); + } + } + + @Override + public JsonNode runBidiJson( + ActionContext ctx, + JsonNode init, + InputSource inputs, + Consumer streamCallback) + throws GenkitException { + try { + Init typedInit = deserializeInit(init); + InputSource typedInputs = adaptInputs(inputs); + Consumer typedCallback = buildTypedCallback(streamCallback); + + // Invoke handler + O result = handler.handle(ctx, typedInit, typedInputs, typedCallback); + return result != null ? MAPPER.valueToTree(result) : null; + } catch (GenkitException ge) { + throw ge; + } catch (Exception e) { + throw new GenkitException("JSON BidiAction execution failed: " + e.getMessage(), e); + } + } + + @Override + public ActionRunResult runBidiJsonWithTelemetry( + ActionContext ctx, + JsonNode init, + InputSource inputs, + Consumer streamCallback) + throws GenkitException { + + final String[] capturedTraceInfo = new String[2]; // [traceId, spanId] + + SpanMetadata spanMetadata = + SpanMetadata.builder() + .name(desc.getName()) + .type(ActionType.AGENT.getValue()) + .subtype(ActionType.AGENT.getValue()) + .build(); + + try { + // Deserialize the real init and adapt the full input stream up front so the handler + // receives the client-managed session state, not a null init / single-input adaptation. + final Init typedInit = deserializeInit(init); + final InputSource typedInputs = adaptInputs(inputs); + final Consumer typedCallback = buildTypedCallback(streamCallback); + + O result = + Tracer.runInNewSpan( + ctx, + spanMetadata, + null, + (spanCtx, in) -> { + capturedTraceInfo[0] = spanCtx.getTraceId(); + capturedTraceInfo[1] = spanCtx.getSpanId(); + try { + return handler.handle( + ctx.withSpanContext(spanCtx), typedInit, typedInputs, typedCallback); + } catch (GenkitException ge) { + throw ge; + } catch (Exception e) { + throw new GenkitException("BidiAction execution failed: " + e.getMessage(), e); + } + }); + + JsonNode jsonResult = result != null ? MAPPER.valueToTree(result) : null; + return new ActionRunResult<>(jsonResult, capturedTraceInfo[0], capturedTraceInfo[1]); + } catch (GenkitException ge) { + throw ge; + } catch (Exception e) { + throw new GenkitException("JSON BidiAction execution failed: " + e.getMessage(), e); + } + } + + // ------------------------------------------------------------------------- + // Registerable + // ------------------------------------------------------------------------- + + @Override + public void register(Registry registry) { + registry.registerAction(desc.getKey(), this); + } + + // ------------------------------------------------------------------------- + // Helpers + // ------------------------------------------------------------------------- + + /** Deserializes a JSON {@code init} value to the typed {@code Init}, or returns null. */ + private Init deserializeInit(JsonNode init) throws Exception { + if (initClass != null && initClass != Void.class && init != null) { + return MAPPER.treeToValue(init, initClass); + } + return null; + } + + /** Adapts an {@code InputSource} to a typed {@code InputSource}. */ + private InputSource adaptInputs(InputSource inputs) { + final InputSource jsonInputs = inputs; + return new InputSource() { + @Override + public Optional next() throws InterruptedException { + Optional jsonNext = jsonInputs.next(); + if (!jsonNext.isPresent()) { + return Optional.empty(); + } + try { + I typed = MAPPER.treeToValue(jsonNext.get(), inputClass); + return Optional.of(typed); + } catch (Exception e) { + throw new RuntimeException("Failed to deserialize input JsonNode", e); + } + } + + @Override + public void close() { + jsonInputs.close(); + } + }; + } + + /** Adapts a {@code Consumer} to a typed {@code Consumer}, or returns null. */ + private Consumer buildTypedCallback(Consumer streamCallback) { + if (streamCallback == null) { + return null; + } + return chunk -> { + try { + JsonNode jsonChunk = MAPPER.valueToTree(chunk); + streamCallback.accept(jsonChunk); + } catch (Exception e) { + throw new RuntimeException("Failed to serialize stream chunk", e); + } + }; + } + + // ------------------------------------------------------------------------- + // Builder + // ------------------------------------------------------------------------- + + /** + * Creates a new builder for {@code BidiActionImpl}. + * + * @param per-turn input type + * @param final output type + * @param stream chunk type + * @param session-init type + * @return a new builder + */ + public static Builder builder() { + return new Builder<>(); + } + + /** Builder for {@link BidiActionImpl}. */ + public static final class Builder { + + private String name; + private Class inputClass; + private Class outputClass; + private Class streamClass; + private Class initClass; + private Map inputSchema; + private Map outputSchema; + private Map metadata; + private BidiAction.BidiHandler handler; + + private Builder() {} + + public Builder name(String name) { + this.name = name; + return this; + } + + public Builder inputClass(Class inputClass) { + this.inputClass = inputClass; + return this; + } + + public Builder outputClass(Class outputClass) { + this.outputClass = outputClass; + return this; + } + + public Builder streamClass(Class streamClass) { + this.streamClass = streamClass; + return this; + } + + public Builder initClass(Class initClass) { + this.initClass = initClass; + return this; + } + + public Builder inputSchema(Map inputSchema) { + this.inputSchema = inputSchema; + return this; + } + + public Builder outputSchema(Map outputSchema) { + this.outputSchema = outputSchema; + return this; + } + + public Builder metadata(Map metadata) { + this.metadata = metadata; + return this; + } + + public Builder handler(BidiAction.BidiHandler handler) { + this.handler = handler; + return this; + } + + /** Builds the {@link BidiActionImpl}. */ + public BidiActionImpl build() { + if (name == null || name.isEmpty()) { + throw new IllegalStateException("name is required"); + } + if (handler == null) { + throw new IllegalStateException("handler is required"); + } + return new BidiActionImpl<>( + name, + inputClass, + outputClass, + streamClass, + initClass, + inputSchema, + outputSchema, + metadata, + handler); + } + } +} diff --git a/core/src/main/java/com/google/genkit/core/BufferedInputSource.java b/core/src/main/java/com/google/genkit/core/BufferedInputSource.java new file mode 100644 index 000000000..53db45cfc --- /dev/null +++ b/core/src/main/java/com/google/genkit/core/BufferedInputSource.java @@ -0,0 +1,93 @@ +/* + * 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. + * + * SPDX-License-Identifier: Apache-2.0 + */ + +package com.google.genkit.core; + +import java.util.Optional; +import java.util.concurrent.LinkedBlockingQueue; + +/** + * A thread-safe {@link InputSource} that is fed by a producer via {@link #offer} and signals + * end-of-stream via {@link #end}. + * + *

Internally uses a {@link LinkedBlockingQueue} of {@code Optional} values. The end-of-stream + * sentinel is {@link Optional#empty()}. Once the sentinel has been consumed, every subsequent call + * to {@link #next()} returns empty immediately. + * + *

Thread-safety: safe for one producer thread calling {@link #offer}/{@link #end} and one + * consumer thread calling {@link #next}. + * + * @param the type of each input element + */ +public final class BufferedInputSource implements InputSource { + + // Sentinel value placed in the queue by end() to signal end-of-stream. + private static final Object END_SENTINEL = new Object(); + + private final LinkedBlockingQueue queue = new LinkedBlockingQueue<>(); + private volatile boolean ended = false; + + /** Creates a new {@code BufferedInputSource}. */ + public BufferedInputSource() {} + + /** + * Enqueues one input for the consumer. Must not be called after {@link #end()}. + * + * @param input the input to enqueue; must not be {@code null} + */ + public void offer(I input) { + if (input == null) { + throw new NullPointerException("input must not be null"); + } + queue.offer(input); + } + + /** + * Signals end-of-stream. After this call, {@link #next()} will return {@link Optional#empty()} + * once all previously enqueued inputs have been consumed. This method is idempotent. + */ + public void end() { + queue.offer(END_SENTINEL); + } + + /** + * {@inheritDoc} + * + *

Blocks until an input is available or end-of-stream is signalled. Once end-of-stream is + * reached, all subsequent calls return {@link Optional#empty()} without blocking. + */ + @Override + @SuppressWarnings("unchecked") + public Optional next() throws InterruptedException { + if (ended) { + return Optional.empty(); + } + Object item = queue.take(); + if (item == END_SENTINEL) { + ended = true; + return Optional.empty(); + } + return Optional.of((I) item); + } + + /** No-op; the queue needs no explicit resource release. */ + @Override + public void close() { + // Nothing to close; end() is handled by the sentinel pattern. + } +} diff --git a/core/src/main/java/com/google/genkit/core/InputSource.java b/core/src/main/java/com/google/genkit/core/InputSource.java new file mode 100644 index 000000000..449eefa71 --- /dev/null +++ b/core/src/main/java/com/google/genkit/core/InputSource.java @@ -0,0 +1,46 @@ +/* + * 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. + * + * SPDX-License-Identifier: Apache-2.0 + */ + +package com.google.genkit.core; + +import java.util.Optional; + +/** + * A blocking pull source of inputs for a bidirectional action. + * + *

The consumer calls {@link #next()} repeatedly to obtain successive inputs. When the stream is + * exhausted, {@link #next()} returns {@link Optional#empty()}. After that, every subsequent call + * also returns empty. The {@link #close()} method releases any resources held by this source. + * + * @param the type of each input element + */ +public interface InputSource extends AutoCloseable { + + /** + * Blocks until the next input is available or the stream ends. + * + * @return an {@link Optional} containing the next input, or {@link Optional#empty()} when the + * stream has ended + * @throws InterruptedException if the calling thread is interrupted while waiting + */ + Optional next() throws InterruptedException; + + /** Releases any resources held by this source. Implementations must be idempotent. */ + @Override + void close(); +} diff --git a/core/src/main/java/com/google/genkit/core/jsonpatch/JsonPatch.java b/core/src/main/java/com/google/genkit/core/jsonpatch/JsonPatch.java new file mode 100644 index 000000000..577181ed0 --- /dev/null +++ b/core/src/main/java/com/google/genkit/core/jsonpatch/JsonPatch.java @@ -0,0 +1,458 @@ +/* + * 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. + * + * SPDX-License-Identifier: Apache-2.0 + */ + +package com.google.genkit.core.jsonpatch; + +import com.fasterxml.jackson.databind.JsonNode; +import com.fasterxml.jackson.databind.node.ArrayNode; +import com.fasterxml.jackson.databind.node.NullNode; +import com.fasterxml.jackson.databind.node.ObjectNode; +import com.google.genkit.core.JsonUtils; +import java.util.ArrayList; +import java.util.LinkedHashSet; +import java.util.List; +import java.util.Set; + +/** + * A tiny, dependency-free RFC 6902 (JSON Patch) implementation operating on Jackson {@link + * JsonNode}. + * + *

Genkit uses JSON Patch to stream incremental changes to a session's custom state ({@code + * AgentStreamChunk.customPatch}). The {@link #diff} helper only emits {@code add} / {@code remove} + * / {@code replace} operations (a valid RFC 6902 subset; {@code move} / {@code copy} are + * optimisations we deliberately skip from diff output), while {@link #apply} understands the full + * operation set for interoperability. + * + *

Ported from {@code js/ai/src/json-patch.ts} in the Genkit upstream repository. + */ +public final class JsonPatch { + + /** Reference tokens that could be used for prototype-pollution-style attacks. */ + private static final Set FORBIDDEN_TOKENS = new LinkedHashSet<>(); + + static { + FORBIDDEN_TOKENS.add("__proto__"); + FORBIDDEN_TOKENS.add("prototype"); + FORBIDDEN_TOKENS.add("constructor"); + } + + private JsonPatch() {} + + // ──────────────────────────────────────────────────────────────────────────── + // Public API + // ──────────────────────────────────────────────────────────────────────────── + + /** + * Applies an RFC-6902 patch (a JSON array of ops) to {@code document}, returning the new + * document. + * + *

The input document is not mutated; a deep copy is patched and returned. + * + * @param document the document to patch (may be {@code null} / {@link NullNode}) + * @param patch an {@link ArrayNode} of RFC-6902 operation objects + * @return the patched document + * @throws IllegalArgumentException if a {@code test} op fails or the patch is malformed + */ + public static JsonNode apply(JsonNode document, JsonNode patch) { + JsonNode doc = deepCopy(document); + for (JsonNode op : patch) { + doc = applyOperation(doc, op); + } + return doc; + } + + /** + * Computes a minimal RFC-6902 patch that transforms {@code from} into {@code to}. + * + *

Only {@code add}, {@code remove}, and {@code replace} ops are emitted. Arrays that differ + * are replaced as a single atomic unit (one {@code replace} op for the whole array), matching the + * JS reference implementation in {@code js/ai/src/json-patch.ts}. + * + * @param from the source document + * @param to the target document + * @return an {@link ArrayNode} of RFC-6902 operation objects + */ + public static JsonNode diff(JsonNode from, JsonNode to) { + ArrayNode patch = JsonUtils.getObjectMapper().createArrayNode(); + diffRecursive(from, to, "", patch); + return patch; + } + + /** + * Returns a whole-document replace patch: {@code [{"op":"replace","path":"","value":value}]}. + * + * @param value the value to place at the document root + * @return a single-element {@link ArrayNode} + */ + public static JsonNode wholeDocumentReplace(JsonNode value) { + ArrayNode patch = JsonUtils.getObjectMapper().createArrayNode(); + ObjectNode op = JsonUtils.getObjectMapper().createObjectNode(); + op.put("op", "replace"); + op.put("path", ""); + op.set("value", deepCopy(value)); + patch.add(op); + return patch; + } + + // ──────────────────────────────────────────────────────────────────────────── + // apply internals + // ──────────────────────────────────────────────────────────────────────────── + + private static JsonNode applyOperation(JsonNode doc, JsonNode op) { + String opName = op.path("op").asText(); + String path = op.path("path").asText(); + List tokens = parsePointer(path); + + if (tokens.isEmpty()) { + // Root-level operations + switch (opName) { + case "add": + case "replace": + return deepCopy(op.get("value")); + case "remove": + return NullNode.getInstance(); + case "test": + { + JsonNode expected = op.get("value"); + if (!deepEqual(doc, expected)) { + throw new IllegalArgumentException("JSON Patch 'test' failed at root."); + } + return doc; + } + case "move": + case "copy": + { + String fromPath = op.path("from").asText(); + List fromTokens = parsePointer(fromPath); + return deepCopy(getValue(doc, fromTokens)); + } + default: + throw new IllegalArgumentException("Unsupported JSON Patch op: " + opName); + } + } + + // Lenient: initialize missing root container for add/replace + if ((doc == null || doc.isNull() || doc.isMissingNode()) + && ("add".equals(opName) || "replace".equals(opName))) { + doc = JsonUtils.getObjectMapper().createObjectNode(); + } + + switch (opName) { + case "add": + setValue(doc, tokens, deepCopy(op.get("value")), true); + return doc; + case "replace": + setValue(doc, tokens, deepCopy(op.get("value")), false); + return doc; + case "remove": + removeValue(doc, tokens); + return doc; + case "test": + { + JsonNode actual = getValue(doc, tokens); + JsonNode expected = op.get("value"); + if (!deepEqual(actual, expected)) { + throw new IllegalArgumentException("JSON Patch 'test' failed at \"" + path + "\"."); + } + return doc; + } + case "move": + { + String fromPath = op.path("from").asText(); + List fromTokens = parsePointer(fromPath); + JsonNode value = deepCopy(getValue(doc, fromTokens)); + removeValue(doc, fromTokens); + setValue(doc, tokens, value, true); + return doc; + } + case "copy": + { + String fromPath = op.path("from").asText(); + List fromTokens = parsePointer(fromPath); + JsonNode value = deepCopy(getValue(doc, fromTokens)); + setValue(doc, tokens, value, true); + return doc; + } + default: + throw new IllegalArgumentException("Unsupported JSON Patch op: " + opName); + } + } + + /** Reads the value at {@code tokens}, returning {@link NullNode} for any missing segment. */ + private static JsonNode getValue(JsonNode doc, List tokens) { + JsonNode cur = doc; + for (String token : tokens) { + if (cur == null || cur.isNull() || cur.isMissingNode()) { + return NullNode.getInstance(); + } + if (cur.isArray()) { + int idx = parseIndex(token); + if (idx < 0 || idx >= cur.size()) { + return NullNode.getInstance(); + } + cur = cur.get(idx); + } else if (cur.isObject()) { + cur = cur.path(token); + if (cur.isMissingNode()) { + return NullNode.getInstance(); + } + } else { + return NullNode.getInstance(); + } + } + return cur != null ? cur : NullNode.getInstance(); + } + + /** + * Sets the value at {@code tokens}. When {@code isAdd} is true and the parent is an array, the + * special {@code -} token appends and a numeric token inserts at that index. + */ + private static void setValue(JsonNode doc, List tokens, JsonNode value, boolean isAdd) { + JsonNode parent = ensureParent(doc, tokens); + if (parent == null || parent.isNull() || parent.isMissingNode()) { + return; // lenient: nothing to set onto + } + String last = tokens.get(tokens.size() - 1); + if (parent.isArray()) { + ArrayNode arr = (ArrayNode) parent; + if ("-".equals(last)) { + arr.add(value); + return; + } + int idx = parseIndex(last); + if (idx < 0) { + return; + } + if (isAdd) { + arr.insert(idx, value); + } else { + arr.set(idx, value); + } + return; + } + if (parent.isObject()) { + ((ObjectNode) parent).set(last, value); + } + } + + /** Removes the value at {@code tokens}. Missing members are a no-op. */ + private static void removeValue(JsonNode doc, List tokens) { + List parentTokens = tokens.subList(0, tokens.size() - 1); + JsonNode parent = getValue(doc, parentTokens); + if (parent == null || parent.isNull() || parent.isMissingNode()) { + return; + } + String last = tokens.get(tokens.size() - 1); + if (parent.isArray()) { + int idx = parseIndex(last); + if (idx >= 0 && idx < parent.size()) { + ((ArrayNode) parent).remove(idx); + } + } else if (parent.isObject()) { + ((ObjectNode) parent).remove(last); + } + } + + /** + * Walks to the parent container of {@code tokens}, lazily creating intermediate object nodes for + * missing segments (lenient apply behaviour). + */ + private static JsonNode ensureParent(JsonNode doc, List tokens) { + JsonNode cur = doc; + for (int i = 0; i < tokens.size() - 1; i++) { + String token = tokens.get(i); + if (cur == null || cur.isNull() || cur.isMissingNode()) { + return null; + } + JsonNode next; + if (cur.isArray()) { + int idx = parseIndex(token); + if (idx < 0) { + return null; + } + next = cur.get(idx); + } else { + next = cur.path(token); + } + if (next == null || next.isMissingNode() || next.isNull() || !next.isContainerNode()) { + // Create an intermediate object container + ObjectNode created = JsonUtils.getObjectMapper().createObjectNode(); + if (cur.isArray()) { + int idx = parseIndex(token); + ((ArrayNode) cur).set(idx, created); + } else { + ((ObjectNode) cur).set(token, created); + } + cur = created; + } else { + cur = next; + } + } + return cur; + } + + // ──────────────────────────────────────────────────────────────────────────── + // diff internals + // ──────────────────────────────────────────────────────────────────────────── + + private static void diffRecursive(JsonNode from, JsonNode to, String pointer, ArrayNode patch) { + if (deepEqual(from, to)) { + return; + } + + // Both plain objects → recurse member-by-member + if (isObject(from) && isObject(to)) { + Set keys = new LinkedHashSet<>(); + from.fieldNames().forEachRemaining(keys::add); + to.fieldNames().forEachRemaining(keys::add); + + for (String key : keys) { + String childPointer = pointer + "/" + escapeToken(key); + boolean inFrom = from.has(key); + boolean inTo = to.has(key); + if (inFrom && !inTo) { + ObjectNode op = JsonUtils.getObjectMapper().createObjectNode(); + op.put("op", "remove"); + op.put("path", childPointer); + patch.add(op); + } else if (!inFrom && inTo) { + ObjectNode op = JsonUtils.getObjectMapper().createObjectNode(); + op.put("op", "add"); + op.put("path", childPointer); + op.set("value", deepCopy(to.get(key))); + patch.add(op); + } else if (inFrom && inTo) { + diffRecursive(from.get(key), to.get(key), childPointer, patch); + } + } + return; + } + + // Both arrays → treat the array as a single atomic value; emit one replace op + if (from != null && from.isArray() && to != null && to.isArray()) { + // deepEqual already returned false above, so they differ + ObjectNode op = JsonUtils.getObjectMapper().createObjectNode(); + op.put("op", "replace"); + op.put("path", pointer); + op.set("value", deepCopy(to)); + patch.add(op); + return; + } + + // Type mismatch or differing primitives → replace at this location + ObjectNode op = JsonUtils.getObjectMapper().createObjectNode(); + op.put("op", "replace"); + op.put("path", pointer); + op.set("value", deepCopy(to)); + patch.add(op); + } + + // ──────────────────────────────────────────────────────────────────────────── + // Pointer helpers + // ──────────────────────────────────────────────────────────────────────────── + + /** + * Parses a JSON Pointer string (RFC-6901) into its reference tokens. + * + *

The root pointer ({@code ""}) returns an empty list. Forbidden tokens ({@code __proto__}, + * {@code prototype}, {@code constructor}) are rejected. + */ + private static List parsePointer(String pointer) { + if (pointer == null || pointer.isEmpty()) { + return new ArrayList<>(); + } + if (pointer.charAt(0) != '/') { + throw new IllegalArgumentException( + "Invalid JSON Pointer: \"" + pointer + "\" must start with \"/\"."); + } + String[] parts = pointer.substring(1).split("/", -1); + List tokens = new ArrayList<>(parts.length); + for (String part : parts) { + String token = unescapeToken(part); + if (FORBIDDEN_TOKENS.contains(token)) { + throw new IllegalArgumentException( + "Invalid JSON Pointer: \"" + + pointer + + "\" contains forbidden token \"" + + token + + "\"."); + } + tokens.add(token); + } + return tokens; + } + + /** + * Escapes a single reference token per RFC-6901 ({@code ~} → {@code ~0}, {@code /} → {@code ~1}). + */ + private static String escapeToken(String token) { + return token.replace("~", "~0").replace("/", "~1"); + } + + /** + * Unescapes a single reference token per RFC-6901 ({@code ~1} → {@code /}, {@code ~0} → {@code + * ~}). + */ + private static String unescapeToken(String token) { + // Order matters: unescape ~1 before ~0 + return token.replace("~1", "/").replace("~0", "~"); + } + + // ──────────────────────────────────────────────────────────────────────────── + // JsonNode helpers + // ──────────────────────────────────────────────────────────────────────────── + + private static boolean isObject(JsonNode node) { + return node != null && node.isObject(); + } + + /** + * Deep structural equality for {@link JsonNode} values, mirroring JSON semantics (missing nodes + * equal null). + */ + private static boolean deepEqual(JsonNode a, JsonNode b) { + // Normalize nulls/missing + boolean aNullish = a == null || a.isNull() || a.isMissingNode(); + boolean bNullish = b == null || b.isNull() || b.isMissingNode(); + if (aNullish && bNullish) { + return true; + } + if (aNullish || bNullish) { + return false; + } + return a.equals(b); + } + + /** Returns a deep copy of {@code node}, or {@link NullNode} if {@code node} is null. */ + private static JsonNode deepCopy(JsonNode node) { + if (node == null) { + return NullNode.getInstance(); + } + return node.deepCopy(); + } + + /** Parses an array index token. Returns -1 if the token is not a non-negative integer. */ + private static int parseIndex(String token) { + try { + int idx = Integer.parseInt(token); + return idx >= 0 ? idx : -1; + } catch (NumberFormatException e) { + return -1; + } + } +} diff --git a/core/src/main/java/com/google/genkit/core/jsonpatch/package-info.java b/core/src/main/java/com/google/genkit/core/jsonpatch/package-info.java new file mode 100644 index 000000000..b47b52a8b --- /dev/null +++ b/core/src/main/java/com/google/genkit/core/jsonpatch/package-info.java @@ -0,0 +1,52 @@ +/* + * 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. + * + * SPDX-License-Identifier: Apache-2.0 + */ + +/** + * RFC 6902 (JSON Patch) support for Genkit Java. + * + *

This package provides a hand-rolled, dependency-free RFC 6902 JSON Patch implementation + * operating on Jackson {@link com.fasterxml.jackson.databind.JsonNode}. It is intentionally + * self-contained so that the diff output shape is controlled exactly — no external JSON-patch + * library is used. + * + *

Genkit uses JSON Patch to stream incremental changes to a session's custom state ({@code + * AgentStreamChunk.customPatch}). The first patch of every agent turn is a whole-document replace + * at the root pointer ({@code ""}); subsequent patches are incremental diffs. + * + *

Example usage: + * + * 

{@code
+ * JsonNode from = mapper.readTree("{\"counter\":1}");
+ * JsonNode to   = mapper.readTree("{\"counter\":2}");
+ *
+ * // Compute a minimal patch
+ * JsonNode patch = JsonPatch.diff(from, to);
+ * // → [{"op":"replace","path":"/counter","value":2}]
+ *
+ * // Apply the patch (does not mutate 'from')
+ * JsonNode result = JsonPatch.apply(from, patch);
+ * // → {"counter":2}
+ *
+ * // Whole-document replace (first patch of a turn)
+ * JsonNode firstPatch = JsonPatch.wholeDocumentReplace(to);
+ * // → [{"op":"replace","path":"","value":{"counter":2}}]
+ * }
+ * + * @see com.google.genkit.core.jsonpatch.JsonPatch + */ +package com.google.genkit.core.jsonpatch; diff --git a/core/src/test/java/com/google/genkit/core/ActionContextTest.java b/core/src/test/java/com/google/genkit/core/ActionContextTest.java index 8d7b1046a..369514a8b 100644 --- a/core/src/test/java/com/google/genkit/core/ActionContextTest.java +++ b/core/src/test/java/com/google/genkit/core/ActionContextTest.java @@ -20,6 +20,8 @@ import static org.junit.jupiter.api.Assertions.*; +import com.google.genkit.core.tracing.SpanContext; +import java.util.Map; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; @@ -168,4 +170,76 @@ void testNullSpanContext() { assertNull(context.getSpanContext()); assertEquals("testFlow", context.getFlowName()); } + + // ── user-context: builder + getter + withContext ───────────────────────────── + + @Test + void testBuilderSetsContext() { + Map userContext = Map.of("auth", Map.of("user", "alice")); + + ActionContext context = ActionContext.builder().registry(registry).context(userContext).build(); + + assertSame(userContext, context.getContext()); + assertEquals("alice", nestedUser(context)); + } + + @Test + void testWithContext() { + Map userContext = Map.of("auth", Map.of("user", "bob")); + + ActionContext context = new ActionContext(registry).withContext(userContext); + + assertSame(userContext, context.getContext()); + } + + @Test + void testContextDefaultsToNull() { + ActionContext context = new ActionContext(registry); + assertNull(context.getContext()); + + ActionContext built = ActionContext.builder().registry(registry).build(); + assertNull(built.getContext()); + } + + @Test + void testContextSurvivesWithSpanContext() { + Map userContext = Map.of("auth", Map.of("user", "alice")); + ActionContext context = new ActionContext(registry).withContext(userContext); + + // The run path (BidiActionImpl / ReflectionServerV2) calls withSpanContext — context MUST + // survive it, otherwise tools never see the injected execution context. + SpanContext spanCtx = new SpanContext("trace-1", "span-1", null); + ActionContext withSpan = context.withSpanContext(spanCtx); + + assertSame(spanCtx, withSpan.getSpanContext()); + assertSame(userContext, withSpan.getContext()); + assertEquals("alice", nestedUser(withSpan)); + } + + @Test + void testContextSurvivesAllWithers() { + Map userContext = Map.of("auth", Map.of("user", "alice")); + + ActionContext context = + new ActionContext(registry) + .withContext(userContext) + .withSpanContext(new SpanContext("t", "s", null)) + .withFlowName("flow") + .withSpanPath("/flow/flow") + .withSessionId("session-1") + .withThreadName("thread-1"); + + assertSame(userContext, context.getContext()); + assertEquals("flow", context.getFlowName()); + assertEquals("/flow/flow", context.getSpanPath()); + assertEquals("session-1", context.getSessionId()); + assertEquals("thread-1", context.getThreadName()); + assertEquals("alice", nestedUser(context)); + } + + @SuppressWarnings("unchecked") + private static String nestedUser(ActionContext context) { + Map auth = (Map) context.getContext().get("auth"); + return (String) auth.get("user"); + } } diff --git a/core/src/test/java/com/google/genkit/core/ActionTypeTest.java b/core/src/test/java/com/google/genkit/core/ActionTypeTest.java index b3c1a2cdd..4a30db036 100644 --- a/core/src/test/java/com/google/genkit/core/ActionTypeTest.java +++ b/core/src/test/java/com/google/genkit/core/ActionTypeTest.java @@ -175,4 +175,32 @@ void testEnumValueOf() { assertEquals(ActionType.FLOW, ActionType.valueOf("FLOW")); assertEquals(ActionType.MODEL, ActionType.valueOf("MODEL")); } + + @Test + void testAgentType() { + assertEquals("agent", ActionType.AGENT.getValue()); + assertEquals("agent", ActionType.AGENT.toString()); + assertEquals("/agent/weatherAgent", ActionType.AGENT.keyFromName("weatherAgent")); + } + + @Test + void testAgentSnapshotType() { + assertEquals("agent-snapshot", ActionType.AGENT_SNAPSHOT.getValue()); + assertEquals("agent-snapshot", ActionType.AGENT_SNAPSHOT.toString()); + assertEquals("/agent-snapshot/x", ActionType.AGENT_SNAPSHOT.keyFromName("x")); + } + + @Test + void testAgentAbortType() { + assertEquals("agent-abort", ActionType.AGENT_ABORT.getValue()); + assertEquals("agent-abort", ActionType.AGENT_ABORT.toString()); + assertEquals("/agent-abort/x", ActionType.AGENT_ABORT.keyFromName("x")); + } + + @Test + void testAgentTypeFromValue() { + assertEquals(ActionType.AGENT, ActionType.fromValue("agent")); + assertEquals(ActionType.AGENT_SNAPSHOT, ActionType.fromValue("agent-snapshot")); + assertEquals(ActionType.AGENT_ABORT, ActionType.fromValue("agent-abort")); + } } diff --git a/core/src/test/java/com/google/genkit/core/BidiActionImplTest.java b/core/src/test/java/com/google/genkit/core/BidiActionImplTest.java new file mode 100644 index 000000000..f5d207207 --- /dev/null +++ b/core/src/test/java/com/google/genkit/core/BidiActionImplTest.java @@ -0,0 +1,339 @@ +/* + * 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. + * + * SPDX-License-Identifier: Apache-2.0 + */ + +package com.google.genkit.core; + +import static org.junit.jupiter.api.Assertions.*; + +import com.fasterxml.jackson.databind.JsonNode; +import com.fasterxml.jackson.databind.ObjectMapper; +import java.util.ArrayList; +import java.util.List; +import java.util.Map; +import java.util.Optional; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.TimeUnit; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; + +/** Unit tests for BidiActionImpl. */ +class BidiActionImplTest { + + private Registry registry; + private ActionContext ctx; + private static final ObjectMapper MAPPER = JsonUtils.getObjectMapper(); + + @BeforeEach + void setUp() { + registry = new DefaultRegistry(); + ctx = new ActionContext(registry); + } + + // ------------------------------------------------------------------------- + // Multi-input bidi test: feed 3 inputs from a separate thread, assert + // runBidi returns 3 and 3 chunks were observed. + // ------------------------------------------------------------------------- + @Test + void testMultiInputBidi() throws Exception { + // Handler: drain inputs, count them, emit one chunk per input, return count + BidiActionImpl action = + BidiActionImpl.builder() + .name("countInputs") + .inputClass(Integer.class) + .outputClass(Integer.class) + .streamClass(Integer.class) + .initClass(Void.class) + .handler( + (handlerCtx, init, inputs, cb) -> { + int count = 0; + Optional next; + while ((next = inputs.next()).isPresent()) { + count++; + cb.accept(next.get()); + } + return count; + }) + .build(); + + List chunks = new ArrayList<>(); + BufferedInputSource source = new BufferedInputSource<>(); + + // Feed inputs from a separate thread + ExecutorService exec = Executors.newSingleThreadExecutor(); + CountDownLatch latch = new CountDownLatch(1); + exec.submit( + () -> { + try { + latch.await(); // Wait until runBidi starts + source.offer(10); + source.offer(20); + source.offer(30); + source.end(); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + } + }); + + // Use a latch to signal producer after runBidi is called + latch.countDown(); + + Integer result = action.runBidi(ctx, null, source, chunks::add); + + exec.shutdown(); + assertTrue(exec.awaitTermination(5, TimeUnit.SECONDS)); + + assertEquals(3, result); + assertEquals(3, chunks.size()); + assertEquals(List.of(10, 20, 30), chunks); + } + + // ------------------------------------------------------------------------- + // Type + metadata assertions + // ------------------------------------------------------------------------- + @Test + void testTypeAndMetadata() { + BidiActionImpl action = + BidiActionImpl.builder() + .name("myAgent") + .inputClass(Integer.class) + .outputClass(Integer.class) + .streamClass(Integer.class) + .initClass(Void.class) + .metadata(Map.of("custom", "value")) + .handler((handlerCtx, init, inputs, cb) -> 0) + .build(); + + assertEquals(ActionType.AGENT, action.getType()); + assertEquals(Boolean.TRUE, action.getMetadata().get("bidi")); + // custom metadata must also be present + assertEquals("value", action.getMetadata().get("custom")); + // desc key must be /agent/ + assertEquals("/agent/myAgent", action.getDesc().getKey()); + } + + // ------------------------------------------------------------------------- + // Caller's metadata map must NOT be mutated + // ------------------------------------------------------------------------- + @Test + void testMetadataNotMutated() { + Map callerMeta = new java.util.HashMap<>(); + callerMeta.put("x", "y"); + + BidiActionImpl action = + BidiActionImpl.builder() + .name("myAgent2") + .inputClass(Integer.class) + .outputClass(Integer.class) + .streamClass(Integer.class) + .initClass(Void.class) + .metadata(callerMeta) + .handler((handlerCtx, init, inputs, cb) -> 0) + .build(); + + // Original map must not contain "bidi" + assertFalse(callerMeta.containsKey("bidi")); + // Action's metadata must contain "bidi" + assertEquals(Boolean.TRUE, action.getMetadata().get("bidi")); + } + + // ------------------------------------------------------------------------- + // Unary adaptation: action.run(ctx, singleInput, chunkCollector) must work + // ------------------------------------------------------------------------- + @Test + void testUnaryAdaptation() throws Exception { + BidiActionImpl action = + BidiActionImpl.builder() + .name("unaryAgent") + .inputClass(Integer.class) + .outputClass(Integer.class) + .streamClass(Integer.class) + .initClass(Void.class) + .handler( + (handlerCtx, init, inputs, cb) -> { + int count = 0; + Optional next; + while ((next = inputs.next()).isPresent()) { + count++; + cb.accept(next.get()); + } + return count; + }) + .build(); + + List chunks = new ArrayList<>(); + Integer result = action.run(ctx, 42, chunks::add); + + // Should have seen exactly 1 input → count = 1 + assertEquals(1, result); + assertEquals(1, chunks.size()); + assertEquals(42, chunks.get(0)); + } + + // ------------------------------------------------------------------------- + // Unary adaptation with null input + // ------------------------------------------------------------------------- + @Test + void testUnaryAdaptationNullInput() throws Exception { + BidiActionImpl action = + BidiActionImpl.builder() + .name("unaryNullAgent") + .inputClass(Integer.class) + .outputClass(Integer.class) + .streamClass(Integer.class) + .initClass(Void.class) + .handler( + (handlerCtx, init, inputs, cb) -> { + int count = 0; + Optional next; + while ((next = inputs.next()).isPresent()) { + count++; + } + return count; + }) + .build(); + + // null input → 0 inputs seen + Integer result = action.run(ctx, null, null); + assertEquals(0, result); + } + + // ------------------------------------------------------------------------- + // runBidiJson round-trip: feed JsonNode inputs, assert JSON output + // ------------------------------------------------------------------------- + @Test + void testRunBidiJsonRoundTrip() throws Exception { + BidiActionImpl action = + BidiActionImpl.builder() + .name("jsonAgent") + .inputClass(Integer.class) + .outputClass(Integer.class) + .streamClass(Integer.class) + .initClass(Void.class) + .handler( + (handlerCtx, init, inputs, cb) -> { + int count = 0; + Optional next; + while ((next = inputs.next()).isPresent()) { + count++; + cb.accept(next.get() * 2); // emit doubled value as chunk + } + return count; + }) + .build(); + + BufferedInputSource jsonSource = new BufferedInputSource<>(); + jsonSource.offer(MAPPER.valueToTree(5)); + jsonSource.offer(MAPPER.valueToTree(10)); + jsonSource.end(); + + List jsonChunks = new ArrayList<>(); + JsonNode jsonResult = action.runBidiJson(ctx, null, jsonSource, jsonChunks::add); + + // 2 inputs → count = 2 + assertNotNull(jsonResult); + assertEquals(2, jsonResult.asInt()); + assertEquals(2, jsonChunks.size()); + assertEquals(10, jsonChunks.get(0).asInt()); // 5 * 2 + assertEquals(20, jsonChunks.get(1).asInt()); // 10 * 2 + } + + // ------------------------------------------------------------------------- + // BufferedInputSource: offer/end/next semantics + // ------------------------------------------------------------------------- + @Test + void testBufferedInputSourceBasic() throws InterruptedException { + BufferedInputSource source = new BufferedInputSource<>(); + source.offer("hello"); + source.offer("world"); + source.end(); + + assertEquals(Optional.of("hello"), source.next()); + assertEquals(Optional.of("world"), source.next()); + assertEquals(Optional.empty(), source.next()); + // After end, subsequent calls must also return empty + assertEquals(Optional.empty(), source.next()); + } + + // ------------------------------------------------------------------------- + // Registration: register(registry) must register under /agent/ + // ------------------------------------------------------------------------- + @Test + void testRegistration() { + BidiActionImpl action = + BidiActionImpl.builder() + .name("registeredAgent") + .inputClass(Integer.class) + .outputClass(Integer.class) + .streamClass(Integer.class) + .initClass(Void.class) + .handler((handlerCtx, init, inputs, cb) -> 0) + .build(); + + action.register(registry); + + Action found = registry.lookupAction("/agent/registeredAgent"); + assertNotNull(found); + assertSame(action, found); + } + + // ------------------------------------------------------------------------- + // runBidiJsonWithTelemetry must thread the REAL init to the handler — the V1 + // reflection server relies on this for agent multi-turn (the Dev UI sends the + // prior turn's session state/snapshotId in init each turn). The inherited + // unary runJsonWithTelemetry adaptation passes a null init, which is exactly + // why agents could not be resumed over V1 before the fix. + // ------------------------------------------------------------------------- + @Test + void testRunBidiJsonWithTelemetryThreadsInit() throws Exception { + // Handler echoes back the init it received: {"initSeen": } + BidiActionImpl action = + BidiActionImpl.builder() + .name("echoInit") + .inputClass(JsonNode.class) + .outputClass(JsonNode.class) + .streamClass(JsonNode.class) + .initClass(JsonNode.class) + .handler( + (handlerCtx, init, inputs, cb) -> { + inputs.next(); // drain the single input + var out = MAPPER.createObjectNode(); + out.set("initSeen", init == null ? MAPPER.nullNode() : init); + return out; + }) + .build(); + + JsonNode initJson = MAPPER.readTree("{\"state\":{\"messages\":[{\"role\":\"user\"}]}}"); + + // Bidi-with-telemetry: init MUST reach the handler. + BufferedInputSource inputs = new BufferedInputSource<>(); + inputs.offer(MAPPER.readTree("{\"message\":{\"role\":\"user\"}}")); + inputs.end(); + ActionRunResult bidiResult = + action.runBidiJsonWithTelemetry(ctx, initJson, inputs, null); + assertEquals(initJson, bidiResult.getResult().get("initSeen")); + + // Unary path drops init (handler sees null) — documents why V1 must not use it for agents. + JsonNode unary = + action + .runJsonWithTelemetry(ctx, MAPPER.readTree("{\"message\":{\"role\":\"user\"}}"), null) + .getResult(); + assertTrue(unary.get("initSeen").isNull()); + } +} diff --git a/core/src/test/java/com/google/genkit/core/jsonpatch/JsonPatchTest.java b/core/src/test/java/com/google/genkit/core/jsonpatch/JsonPatchTest.java new file mode 100644 index 000000000..a37ed8c54 --- /dev/null +++ b/core/src/test/java/com/google/genkit/core/jsonpatch/JsonPatchTest.java @@ -0,0 +1,445 @@ +/* + * 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. + * + * SPDX-License-Identifier: Apache-2.0 + */ + +package com.google.genkit.core.jsonpatch; + +import static org.junit.jupiter.api.Assertions.*; + +import com.fasterxml.jackson.databind.JsonNode; +import com.fasterxml.jackson.databind.ObjectMapper; +import com.fasterxml.jackson.databind.node.ObjectNode; +import com.google.genkit.core.JsonUtils; +import org.junit.jupiter.api.Test; + +/** Unit tests for JsonPatch (RFC-6902). */ +class JsonPatchTest { + + private final ObjectMapper mapper = JsonUtils.getObjectMapper(); + + // ────────────────────────────────────────────────────────────────────────── + // (a) apply: whole-document replace at path "" + // ────────────────────────────────────────────────────────────────────────── + + @Test + void applyWholeDocumentReplace() throws Exception { + JsonNode doc = mapper.readTree("{\"x\":9}"); + JsonNode patch = mapper.readTree("[{\"op\":\"replace\",\"path\":\"\",\"value\":{\"a\":1}}]"); + + JsonNode result = JsonPatch.apply(doc, patch); + + assertEquals(mapper.readTree("{\"a\":1}"), result); + } + + // ────────────────────────────────────────────────────────────────────────── + // (b) apply: replace a specific field + // ────────────────────────────────────────────────────────────────────────── + + @Test + void applyReplaceField() throws Exception { + JsonNode doc = mapper.readTree("{\"counter\":1}"); + JsonNode patch = mapper.readTree("[{\"op\":\"replace\",\"path\":\"/counter\",\"value\":2}]"); + + JsonNode result = JsonPatch.apply(doc, patch); + + assertEquals(mapper.readTree("{\"counter\":2}"), result); + } + + // ────────────────────────────────────────────────────────────────────────── + // (c) apply: add/remove object key; add into array at index and at "-"; remove array element + // ────────────────────────────────────────────────────────────────────────── + + @Test + void applyAddNewObjectKey() throws Exception { + JsonNode doc = mapper.readTree("{\"a\":1}"); + JsonNode patch = mapper.readTree("[{\"op\":\"add\",\"path\":\"/b\",\"value\":2}]"); + + JsonNode result = JsonPatch.apply(doc, patch); + + assertEquals(mapper.readTree("{\"a\":1,\"b\":2}"), result); + } + + @Test + void applyRemoveObjectKey() throws Exception { + JsonNode doc = mapper.readTree("{\"a\":1,\"b\":2}"); + JsonNode patch = mapper.readTree("[{\"op\":\"remove\",\"path\":\"/b\"}]"); + + JsonNode result = JsonPatch.apply(doc, patch); + + assertEquals(mapper.readTree("{\"a\":1}"), result); + } + + @Test + void applyAddIntoArrayAtIndex() throws Exception { + JsonNode doc = mapper.readTree("{\"items\":[1,3]}"); + JsonNode patch = mapper.readTree("[{\"op\":\"add\",\"path\":\"/items/1\",\"value\":2}]"); + + JsonNode result = JsonPatch.apply(doc, patch); + + assertEquals(mapper.readTree("{\"items\":[1,2,3]}"), result); + } + + @Test + void applyAddIntoArrayAtEnd() throws Exception { + JsonNode doc = mapper.readTree("{\"items\":[1]}"); + JsonNode patch = mapper.readTree("[{\"op\":\"add\",\"path\":\"/items/-\",\"value\":2}]"); + + JsonNode result = JsonPatch.apply(doc, patch); + + assertEquals(mapper.readTree("{\"items\":[1,2]}"), result); + } + + @Test + void applyRemoveArrayElement() throws Exception { + JsonNode doc = mapper.readTree("{\"items\":[1,2,3]}"); + JsonNode patch = mapper.readTree("[{\"op\":\"remove\",\"path\":\"/items/1\"}]"); + + JsonNode result = JsonPatch.apply(doc, patch); + + assertEquals(mapper.readTree("{\"items\":[1,3]}"), result); + } + + // ────────────────────────────────────────────────────────────────────────── + // (d) pointer escaping: ~1 → /, ~0 → ~ + // ────────────────────────────────────────────────────────────────────────── + + @Test + void applyPointerEscapingSlash() throws Exception { + JsonNode doc = mapper.readTree("{}"); + JsonNode patch = mapper.readTree("[{\"op\":\"add\",\"path\":\"/a~1b\",\"value\":1}]"); + + JsonNode result = JsonPatch.apply(doc, patch); + + assertEquals(mapper.readTree("{\"a/b\":1}"), result); + } + + @Test + void applyPointerEscapingTilde() throws Exception { + JsonNode doc = mapper.readTree("{}"); + JsonNode patch = mapper.readTree("[{\"op\":\"add\",\"path\":\"/c~0d\",\"value\":2}]"); + + JsonNode result = JsonPatch.apply(doc, patch); + + assertEquals(mapper.readTree("{\"c~d\":2}"), result); + } + + // ────────────────────────────────────────────────────────────────────────── + // (e) test op: success leaves doc unchanged; failure raises + // ────────────────────────────────────────────────────────────────────────── + + @Test + void applyTestOpSuccess() throws Exception { + JsonNode doc = mapper.readTree("{\"a\":1}"); + JsonNode patch = mapper.readTree("[{\"op\":\"test\",\"path\":\"/a\",\"value\":1}]"); + + JsonNode result = JsonPatch.apply(doc, patch); + + assertEquals(mapper.readTree("{\"a\":1}"), result); + } + + @Test + void applyTestOpFailureThrows() throws Exception { + JsonNode doc = mapper.readTree("{\"a\":1}"); + JsonNode patch = mapper.readTree("[{\"op\":\"test\",\"path\":\"/a\",\"value\":2}]"); + + assertThrows(IllegalArgumentException.class, () -> JsonPatch.apply(doc, patch)); + } + + @Test + void applyTestOpRootSuccess() throws Exception { + JsonNode doc = mapper.readTree("{\"a\":1}"); + JsonNode patch = mapper.readTree("[{\"op\":\"test\",\"path\":\"\",\"value\":{\"a\":1}}]"); + + JsonNode result = JsonPatch.apply(doc, patch); + + assertEquals(mapper.readTree("{\"a\":1}"), result); + } + + @Test + void applyTestOpRootFailureThrows() throws Exception { + JsonNode doc = mapper.readTree("{\"a\":1}"); + JsonNode patch = mapper.readTree("[{\"op\":\"test\",\"path\":\"\",\"value\":{\"b\":2}}]"); + + assertThrows(IllegalArgumentException.class, () -> JsonPatch.apply(doc, patch)); + } + + // ────────────────────────────────────────────────────────────────────────── + // (f) diff: simple replace + // ────────────────────────────────────────────────────────────────────────── + + @Test + void diffSimpleReplace() throws Exception { + JsonNode from = mapper.readTree("{\"counter\":1}"); + JsonNode to = mapper.readTree("{\"counter\":2}"); + + JsonNode result = JsonPatch.diff(from, to); + + assertEquals( + mapper.readTree("[{\"op\":\"replace\",\"path\":\"/counter\",\"value\":2}]"), result); + } + + @Test + void diffAddMember() throws Exception { + JsonNode from = mapper.readTree("{\"a\":1}"); + JsonNode to = mapper.readTree("{\"a\":1,\"b\":2}"); + + JsonNode result = JsonPatch.diff(from, to); + + assertEquals(mapper.readTree("[{\"op\":\"add\",\"path\":\"/b\",\"value\":2}]"), result); + } + + @Test + void diffRemoveMember() throws Exception { + JsonNode from = mapper.readTree("{\"a\":1,\"b\":2}"); + JsonNode to = mapper.readTree("{\"a\":1}"); + + JsonNode result = JsonPatch.diff(from, to); + + assertEquals(mapper.readTree("[{\"op\":\"remove\",\"path\":\"/b\"}]"), result); + } + + @Test + void diffEqualValues() throws Exception { + JsonNode from = mapper.readTree("{\"a\":1}"); + JsonNode to = mapper.readTree("{\"a\":1}"); + + JsonNode result = JsonPatch.diff(from, to); + + assertEquals(mapper.readTree("[]"), result); + } + + // ────────────────────────────────────────────────────────────────────────── + // (g) round-trip: apply(from, diff(from, to)) equals to + // ────────────────────────────────────────────────────────────────────────── + + @Test + void roundTripObjectAddRemoveReplace() throws Exception { + JsonNode from = mapper.readTree("{\"a\":1,\"b\":2}"); + JsonNode to = mapper.readTree("{\"a\":3,\"c\":4}"); + + assertRoundTrip(from, to); + } + + @Test + void roundTripNestedObject() throws Exception { + JsonNode from = mapper.readTree("{\"nested\":{\"x\":1}}"); + JsonNode to = mapper.readTree("{\"nested\":{\"x\":2,\"y\":3}}"); + + assertRoundTrip(from, to); + } + + @Test + void roundTripArrayChange() throws Exception { + JsonNode from = mapper.readTree("{\"items\":[1,2]}"); + JsonNode to = mapper.readTree("{\"items\":[1,2,3]}"); + + assertRoundTrip(from, to); + } + + @Test + void roundTripArrayShrink() throws Exception { + JsonNode from = mapper.readTree("{\"items\":[1,2,3]}"); + JsonNode to = mapper.readTree("{\"items\":[1]}"); + + assertRoundTrip(from, to); + } + + @Test + void roundTripComplexMutation() throws Exception { + JsonNode from = mapper.readTree("{\"status\":\"a\",\"items\":[1,2],\"nested\":{\"x\":1}}"); + JsonNode to = + mapper.readTree("{\"status\":\"b\",\"items\":[1,2,3],\"nested\":{\"x\":1,\"y\":2}}"); + + assertRoundTrip(from, to); + } + + @Test + void roundTripTypeChange() throws Exception { + // object → array: triggers whole-document replace + JsonNode from = mapper.readTree("{\"a\":1}"); + JsonNode to = mapper.readTree("[1,2]"); + + assertRoundTrip(from, to); + } + + // ────────────────────────────────────────────────────────────────────────── + // (h) wholeDocumentReplace + // ────────────────────────────────────────────────────────────────────────── + + @Test + void wholeDocumentReplaceShape() throws Exception { + JsonNode value = mapper.readTree("{\"a\":1}"); + + JsonNode result = JsonPatch.wholeDocumentReplace(value); + + assertEquals( + mapper.readTree("[{\"op\":\"replace\",\"path\":\"\",\"value\":{\"a\":1}}]"), result); + } + + // ────────────────────────────────────────────────────────────────────────── + // (i) apply does not mutate input document + // ────────────────────────────────────────────────────────────────────────── + + @Test + void applyDoesNotMutateInput() throws Exception { + JsonNode doc = mapper.readTree("{\"a\":1}"); + String originalJson = doc.toString(); + + JsonNode patch = mapper.readTree("[{\"op\":\"replace\",\"path\":\"/a\",\"value\":2}]"); + JsonPatch.apply(doc, patch); + + // Original document must be unchanged + assertEquals(originalJson, doc.toString()); + assertEquals(1, doc.get("a").asInt()); + } + + @Test + void applyDoesNotMutateInputOnWholeDocReplace() throws Exception { + JsonNode doc = mapper.readTree("{\"x\":9}"); + String originalJson = doc.toString(); + + JsonNode patch = mapper.readTree("[{\"op\":\"replace\",\"path\":\"\",\"value\":{\"a\":1}}]"); + JsonPatch.apply(doc, patch); + + assertEquals(originalJson, doc.toString()); + } + + // ────────────────────────────────────────────────────────────────────────── + // Additional: move and copy ops + // ────────────────────────────────────────────────────────────────────────── + + @Test + void applyMoveOp() throws Exception { + JsonNode doc = mapper.readTree("{\"a\":1}"); + JsonNode patch = mapper.readTree("[{\"op\":\"move\",\"from\":\"/a\",\"path\":\"/b\"}]"); + + JsonNode result = JsonPatch.apply(doc, patch); + + assertEquals(mapper.readTree("{\"b\":1}"), result); + } + + @Test + void applyCopyOp() throws Exception { + JsonNode doc = mapper.readTree("{\"a\":1}"); + JsonNode patch = mapper.readTree("[{\"op\":\"copy\",\"from\":\"/a\",\"path\":\"/b\"}]"); + + JsonNode result = JsonPatch.apply(doc, patch); + + assertEquals(mapper.readTree("{\"a\":1,\"b\":1}"), result); + } + + // ────────────────────────────────────────────────────────────────────────── + // Additional: prototype pollution guard + // ────────────────────────────────────────────────────────────────────────── + + @Test + void applyRejectsForbiddenTokenInPath() throws Exception { + JsonNode doc = mapper.readTree("{}"); + JsonNode patch = mapper.readTree("[{\"op\":\"add\",\"path\":\"/__proto__/x\",\"value\":1}]"); + + assertThrows(IllegalArgumentException.class, () -> JsonPatch.apply(doc, patch)); + } + + // ────────────────────────────────────────────────────────────────────────── + // (j) diff: arrays are replaced as a single atomic unit + // ────────────────────────────────────────────────────────────────────────── + + @Test + void diffArrayShrinkEmitsSingleReplaceOp() throws Exception { + // diff([1,2,3], [1,2]) must emit exactly ONE replace op with the full new array + JsonNode from = mapper.readTree("{\"items\":[1,2,3]}"); + JsonNode to = mapper.readTree("{\"items\":[1,2]}"); + + JsonNode result = JsonPatch.diff(from, to); + + assertEquals(1, result.size(), "Expected exactly one op, got: " + result); + JsonNode op = result.get(0); + assertEquals("replace", op.get("op").asText()); + assertEquals("/items", op.get("path").asText()); + assertEquals(mapper.readTree("[1,2]"), op.get("value")); + } + + @Test + void diffArrayGrowEmitsSingleReplaceOp() throws Exception { + // diff([1,2], [1,2,3]) must emit exactly ONE replace op with the full new array + JsonNode from = mapper.readTree("{\"items\":[1,2]}"); + JsonNode to = mapper.readTree("{\"items\":[1,2,3]}"); + + JsonNode result = JsonPatch.diff(from, to); + + assertEquals(1, result.size(), "Expected exactly one op, got: " + result); + JsonNode op = result.get(0); + assertEquals("replace", op.get("op").asText()); + assertEquals("/items", op.get("path").asText()); + assertEquals(mapper.readTree("[1,2,3]"), op.get("value")); + } + + @Test + void diffRootLevelArraysEmitsSingleReplaceOpWithEmptyPath() throws Exception { + // When the root documents are both arrays and differ, path must be "" + JsonNode from = mapper.readTree("[1,2,3]"); + JsonNode to = mapper.readTree("[1,2]"); + + JsonNode result = JsonPatch.diff(from, to); + + assertEquals(1, result.size(), "Expected exactly one op, got: " + result); + JsonNode op = result.get(0); + assertEquals("replace", op.get("op").asText()); + assertEquals("", op.get("path").asText()); + assertEquals(mapper.readTree("[1,2]"), op.get("value")); + } + + @Test + void diffEqualArraysEmitsNoOps() throws Exception { + JsonNode from = mapper.readTree("{\"items\":[1,2,3]}"); + JsonNode to = mapper.readTree("{\"items\":[1,2,3]}"); + + JsonNode result = JsonPatch.diff(from, to); + + assertEquals(mapper.readTree("[]"), result); + } + + // ────────────────────────────────────────────────────────────────────────── + // Additional: pointer escaping in diff output + // ────────────────────────────────────────────────────────────────────────── + + @Test + void diffEscapesPointerTokens() throws Exception { + // Keys with "/" and "~" must be escaped in pointer paths + ObjectNode from = mapper.createObjectNode(); + ObjectNode to = mapper.createObjectNode(); + to.put("a/b", 1); + to.put("c~d", 2); + + JsonNode result = JsonPatch.diff(from, to); + + // The paths must use ~1 for "/" and ~0 for "~" + String resultStr = result.toString(); + assertTrue(resultStr.contains("/a~1b"), "Expected escaped /a~1b in: " + resultStr); + assertTrue(resultStr.contains("/c~0d"), "Expected escaped /c~0d in: " + resultStr); + } + + // ────────────────────────────────────────────────────────────────────────── + // Helpers + // ────────────────────────────────────────────────────────────────────────── + + private void assertRoundTrip(JsonNode from, JsonNode to) { + JsonNode patch = JsonPatch.diff(from, to); + JsonNode result = JsonPatch.apply(from, patch); + assertEquals(to, result, "Round-trip failed. patch=" + patch); + } +} diff --git a/docs/astro.config.mjs b/docs/astro.config.mjs index b6497a310..e3ed7269a 100644 --- a/docs/astro.config.mjs +++ b/docs/astro.config.mjs @@ -80,11 +80,25 @@ export default defineConfig({ { label: "Structured Output", slug: "structured-output" }, { label: "Streaming", slug: "streaming" }, { label: "RAG", slug: "rag" }, - { label: "Chat Sessions", slug: "chat-sessions" }, { label: "Evaluations", slug: "evaluations" }, { label: "Middleware", slug: "middleware" }, { label: "Interrupts", slug: "interrupts" }, - { label: "Multi-Agent", slug: "multi-agent" }, + ], + }, + { + label: "Agents (Beta)", + items: [ + { label: "Overview", slug: "agents/overview" }, + { label: "Define Agents", slug: "agents/define-agents" }, + { label: "Run and Stream", slug: "agents/run-and-stream" }, + { label: "Serve over HTTP", slug: "agents/serve-over-http" }, + { label: "Sessions", slug: "agents/sessions" }, + { label: "Session Stores", slug: "agents/session-stores" }, + { label: "Interrupts", slug: "agents/interrupts" }, + { label: "Background Execution", slug: "agents/background-execution" }, + { label: "Multi-agent Delegation", slug: "agents/multi-agent-delegation" }, + { label: "Custom Orchestration", slug: "agents/custom-orchestration" }, + { label: "Error Handling", slug: "agents/error-handling" }, ], }, { diff --git a/docs/src/content/docs/agents/background-execution.md b/docs/src/content/docs/agents/background-execution.md new file mode 100644 index 000000000..9db01fea9 --- /dev/null +++ b/docs/src/content/docs/agents/background-execution.md @@ -0,0 +1,115 @@ +--- +title: Background Execution +description: Run slow agent turns in the background and poll for completion instead of blocking the caller. +--- + +Some agent turns are slow — a long tool chain, a batch job, a multi-minute model call. **Detaching** a turn starts the work in the background and returns to the caller immediately, so you can poll for the result later instead of blocking. This is useful for long-running jobs, HTTP requests that shouldn't hold a connection open, and any turn you'd rather not wait on inline. + +Detach requires a **server-managed** agent — one configured with a `SessionStore` (via `AgentConfig.store(...)` / `CustomAgentConfig.store(...)`) — because there needs to be somewhere to persist the in-progress work and its eventual result. + +## Detaching a turn + +Set `detach(true)` on the input: + +```java +import com.google.genkit.ai.agent.AgentInput; +import com.google.genkit.ai.agent.AgentResponse; +import com.google.genkit.ai.Message; +import java.util.Map; + +AgentResponse> resp = chat.send( + AgentInput.builder() + .message(Message.user("Run the long report")) + .detach(true) + .build()); + +resp.finishReason(); // AgentFinishReason.DETACHED +String snapshotId = resp.snapshotId(); // poll this for the real result +``` + +The call returns almost immediately with `finishReason() == DETACHED` and a `snapshotId`. There's no message or state yet — those only exist once the background work finishes. Over HTTP the same shape comes back as JSON with a `"detached"` finish reason and a `snapshotId` (see [Serve over HTTP](../serve-over-http)). + +A heartbeat keeps a long-running detached turn marked as alive, so a reader can tell a still-running turn apart from one whose process died. + +## Polling for completion + +Poll `getSnapshotData` (in-process) or `POST //getSnapshot` (over HTTP) until the snapshot leaves `PENDING`: + +```java +import com.google.genkit.ai.agent.GetSnapshotRequest; +import com.google.genkit.ai.agent.SessionSnapshot; +import com.google.genkit.ai.agent.SnapshotStatus; + +SessionSnapshot> snap; +do { + Thread.sleep(200); + snap = agent.getSnapshotData( + GetSnapshotRequest.builder().snapshotId(snapshotId).build()); +} while (snap.getStatus() == SnapshotStatus.PENDING); + +if (snap.getStatus() == SnapshotStatus.COMPLETED) { + System.out.println(snap.getState().getMessages()); +} else if (snap.getStatus() == SnapshotStatus.FAILED) { + System.out.println("Failed: " + snap.getError().getMessage()); +} +``` + +Once terminal, the snapshot carries the turn's real outcome — `finishReason` becomes `stop` (or whatever the turn actually produced), never `detached`. The `detached` reason only ever appears on the immediate response to the detach call itself. + +### Snapshot statuses + +| Status | Meaning | +|--------|---------| +| `PENDING` | Reserved and (usually) still running in the background | +| `COMPLETED` | The turn finished normally; `getState()` holds the result | +| `FAILED` | The turn threw; `getError()` holds the error | +| `ABORTED` | The turn was aborted while still pending — see below | +| `EXPIRED` | Reserved for store-specific expiry policies | + +## When a background turn fails + +A background turn that throws doesn't crash anything — the runtime records a `FAILED` snapshot with the error, which you see on your next poll: + +```java +if (snap.getStatus() == SnapshotStatus.FAILED) { + System.err.println("Detached turn failed: " + snap.getError().getMessage()); +} +``` + +So you handle a background failure exactly where you handle a background success: at the poll site, by checking the status. See [Error Handling](../error-handling) for the full error contract. + +## Aborting a detached turn + +`chat.abort()` (or `POST //abort` over HTTP) cancels a still-running detached turn. It marks the snapshot `ABORTED` — so a subsequent poll sees `ABORTED` — and also signals the running turn to stop. A custom `AgentFn` that checks `ctx.isAborted()` in its loop observes this and can stop early: + +```java +import com.google.genkit.ai.agent.AgentFn; +import com.google.genkit.ai.agent.AgentResult; +import com.google.genkit.ai.agent.AgentFinishReason; + +AgentFn> longJob = (sess, ctx) -> { + for (int i = 0; i < 1000 && !ctx.isAborted(); i++) { + // do one unit of work per iteration + } + return AgentResult.builder() + .finishReason(ctx.isAborted() ? AgentFinishReason.ABORTED : AgentFinishReason.STOP) + .build(); +}; +``` + +A turn that never checks `ctx.isAborted()` still runs to the end (an abort can't force-kill work in flight), but the caller reliably sees `ABORTED` either way. Cooperative abort is specific to detached turns — a synchronous `chat.send(...)` turn always runs to completion. + +Abort needs a store that supports live signalling (`FileSessionStore`, `FirestoreSessionStore`); with `InMemorySessionStore` there's nothing to poll or signal, so `abort()` is a no-op. See [Session Stores](../session-stores). + +## Client-managed agents + +Detach has no effect on a client-managed agent (no `SessionStore`): the turn just runs inline and returns its result normally, as if `detach` hadn't been set. Check `finishReason()` if your code needs to confirm a detach actually took effect. + +## See also + +- [Agents Overview](../overview) — Defining agents and the chat API +- [Run and Stream](../run-and-stream) — The normal (non-detached) `send`/`sendStream` turn model +- [Sessions](../sessions) — Snapshot lifecycle and the turn model +- [Session Stores](../session-stores) — Which stores support abort +- [Serve over HTTP](../serve-over-http) — Driving detach / getSnapshot / abort over the wire +- [Error Handling](../error-handling) — How errors surface on a `FAILED` snapshot diff --git a/docs/src/content/docs/agents/custom-orchestration.md b/docs/src/content/docs/agents/custom-orchestration.md new file mode 100644 index 000000000..6b0657d4c --- /dev/null +++ b/docs/src/content/docs/agents/custom-orchestration.md @@ -0,0 +1,153 @@ +--- +title: Custom Orchestration +description: Write your own AgentFn for full control over each turn — session state, streaming, and finish reasons, with no model required. +--- + +Most agents are prompt-backed: you give `defineAgent` a system prompt, tools, and a model, and it runs the model loop for you. A **custom agent** is the opposite — you write the per-turn logic yourself as an `AgentFn`, so you control exactly what happens on each turn. It's the right tool for deterministic, testable orchestration (a state machine, a rules engine, a pipeline of non-LLM steps) that still lives behind the same `AgentChat`, session, and streaming machinery as any other agent. + +## The `AgentFn` contract + +An `AgentFn` receives everything it needs through two parameters — a `SessionRunner` for reading and writing session state, and an `AgentFnContext` for per-turn signals — and returns an `AgentResult`: + +```java +import com.google.genkit.ai.Message; +import com.google.genkit.ai.agent.Agent; +import com.google.genkit.ai.agent.AgentFn; +import com.google.genkit.ai.agent.AgentFinishReason; +import com.google.genkit.ai.agent.AgentResult; +import com.google.genkit.ai.agent.CustomAgentConfig; +import com.google.genkit.ai.agent.InMemorySessionStore; +import java.util.HashMap; +import java.util.Map; + +AgentFn> orchestratorFn = (sess, ctx) -> { + // Read state, do work, optionally emit chunks... + return AgentResult.builder() + .message(Message.model("done")) + .finishReason(AgentFinishReason.STOP) + .build(); +}; + +Agent> orchestrator = genkit.beta().defineCustomAgent( + CustomAgentConfig.>builder() + .name("orchestrator") + .store(new InMemorySessionStore<>()) + .build(), + orchestratorFn); +``` + +`AgentResult` has three fields — `message`, `artifacts`, and `finishReason` (defaults to `STOP`). Whatever you return is what the caller sees as `AgentResponse.message()` / `.artifacts()` / `.finishReason()`. + +## Reading and writing session state + +`sess` (the `SessionRunner`) is your handle onto everything the session knows: + +| Method | Use | +|--------|-----| +| `sess.getMessages()` | The conversation so far, including this turn's user message | +| `sess.addMessages(Message...)` | Append messages (e.g. a tool-request/tool-response pair) | +| `sess.getCustom()` | A copy of the current custom state `S` | +| `sess.updateCustom(UnaryOperator)` | Atomically read-modify-write the custom state | +| `sess.getArtifacts()` / `sess.addArtifacts(Artifact...)` | Artifacts so far / add new ones | +| `sess.turnIndex()` | How many turns have run in this invocation | + +State you write persists across calls. Here a turn counter is carried forward via custom state: + +```java +AgentFn> fn = (sess, ctx) -> { + Map custom = sess.getCustom(); + int previous = custom != null && custom.get("turnCount") != null + ? (Integer) custom.get("turnCount") + : 0; + int turnCount = previous + 1; + + sess.updateCustom(old -> { + Map updated = old != null ? new HashMap<>(old) : new HashMap<>(); + updated.put("turnCount", turnCount); + return updated; + }); + + return AgentResult.builder() + .message(Message.model("turn #" + turnCount)) + .finishReason(AgentFinishReason.STOP) + .build(); +}; +``` + +Two calls to `chat.send(...)` against this agent return `"turn #1"` then `"turn #2"` — the second call sees the state the first one wrote. + +## Emitting chunks during a turn + +`AgentFnContext.sendChunk()` returns a `Consumer`. Call it as often as you like to push incremental output to a caller using `chat.sendStream(...)`: + +```java +AgentFn> fn = (sess, ctx) -> { + for (String step : List.of("fetching", "processing", "done")) { + ctx.sendChunk().accept(AgentStreamChunk.builder().build()); + // Set .modelChunk(...) or .customPatch(...) on the builder to carry a real payload. + } + return AgentResult.builder() + .message(Message.model("finished")) + .finishReason(AgentFinishReason.STOP) + .build(); +}; +``` + +Chunks arrive at the caller's `onChunk` callback. If the caller used `chat.send(...)` (no callback), the chunks simply go nowhere — so emitting them is always safe. + +## Resuming a paused turn + +If your `AgentFn` returns `AgentFinishReason.INTERRUPTED`, the caller can resolve it with `chat.resume(...)` — and the resume payload comes back to you through `ctx.resume()`. Branch on it to distinguish a fresh turn from a resumed one: + +```java +AgentFn> approvalAgent = (sess, ctx) -> { + if (ctx.resume() == null) { + // First turn: ask for confirmation and pause. + return AgentResult.builder() + .message(Message.model("Please confirm: transfer $150 to Alice? (yes/no)")) + .finishReason(AgentFinishReason.INTERRUPTED) + .build(); + } + + // Resumed turn: ctx.resume() carries the parts passed to chat.resume(...). + boolean approved = /* read the response from ctx.resume() */ true; + return AgentResult.builder() + .message(Message.model(approved ? "Transfer completed." : "Transfer cancelled.")) + .finishReason(AgentFinishReason.STOP) + .build(); +}; +``` + +See [Interrupts](../interrupts) for the caller-side flow. + +## Cooperative abort + +For a [detached (background) turn](../background-execution), a long-running body can poll `ctx.isAborted()` and stop early when the caller aborts: + +```java +AgentFn> fn = (sess, ctx) -> { + for (int i = 0; i < 100 && !ctx.isAborted(); i++) { + // do one unit of work + } + return AgentResult.builder() + .finishReason(ctx.isAborted() ? AgentFinishReason.ABORTED : AgentFinishReason.STOP) + .build(); +}; +``` + +Cooperative abort applies only to detached turns — a synchronous turn runs to completion. If you need to stop a foreground turn partway through, wire up your own cancellation flag that your `AgentFn` polls. + +## Detach-awareness + +If you configure a `SessionStore` and the caller sends `detach(true)`, the runtime runs your same `AgentFn` on a background thread and returns `DETACHED` immediately — your function doesn't need to do anything special. Streaming is automatically suppressed for a detached turn (no one is listening), so you can keep emitting chunks unconditionally. See [Background Execution](../background-execution) and [Error Handling](../error-handling) for how a detached failure surfaces. + +## See also + +- [Agents Overview](../overview) — Defining agents and the chat API +- [Define Agents](../define-agents) — `AgentConfig`, `CustomAgentConfig`, and the `AgentFn` contract +- [Run and Stream](../run-and-stream) — `AgentChat.send`/`sendStream`, `AgentResponse`, and chunk shapes +- [Sessions](../sessions) — Session lifecycle, snapshots, and the turn model +- [Session Stores](../session-stores) — Server-managed persistence backends +- [Interrupts](../interrupts) — Reading `chat.resume(...)`'s payload via `ctx.resume()` +- [Multi-agent Delegation](../multi-agent-delegation) — Composing custom agents as sub-agents +- [Error Handling](../error-handling) — Failure modes and how they surface diff --git a/docs/src/content/docs/agents/define-agents.md b/docs/src/content/docs/agents/define-agents.md new file mode 100644 index 000000000..6780051e0 --- /dev/null +++ b/docs/src/content/docs/agents/define-agents.md @@ -0,0 +1,191 @@ +--- +title: Define Agents +description: Build model-backed agents with AgentConfig, or write your own per-turn logic with CustomAgentConfig and AgentFn. +--- + +The [Agents Overview](../overview) shows the quick path with `defineAgent`. This page covers the two ways to define an agent, the full configuration surface, and how to write your own turn logic with `AgentFn` when you need more control. + +## Two ways to define an agent + +| Factory | You provide | Per-turn logic | +|---------|-------------|----------------| +| `defineAgent(AgentConfig)` | A system prompt, model, and tools | Built for you: one model call per turn, with your history and tools | +| `defineCustomAgent(CustomAgentConfig, AgentFn)` | An `AgentFn` | You write it: your function runs once per turn | + +Both register the same kind of `Agent`, so everything in [Run and Stream](../run-and-stream), [Sessions](../sessions), and [Session Stores](../session-stores) applies to agents defined either way. The only difference is what runs during a turn. + +## Model-backed agents: `defineAgent` + +Pass an `AgentConfig`. Only `name` is required. + +```java +import com.google.genkit.agent.AgentConfig; +import com.google.genkit.ai.agent.Agent; +import com.google.genkit.ai.agent.FileSessionStore; + +Agent> weatherAgent = genkit.beta().defineAgent( + AgentConfig.>builder() + .name("weatherAgent") + .description("A helpful weather assistant") + .system("You are a helpful weather assistant. Use the getWeather tool.") + .tools(getWeather) + .model("openai/gpt-4o-mini") + .maxTurns(5) + .store(new FileSessionStore<>("./.snapshots")) + .build()); +``` + +Each turn runs one model call with your system prompt, tools, and the full conversation history. The model may call tools during the turn (up to `maxTurns`), and the request context flows into those tool calls, so a tool can read caller info such as auth. + +### `AgentConfig` fields + +| Field | Type | Description | +|-------|------|-------------| +| `name` | `String` (required) | Registered name; appears as `/agent/` in the Dev UI | +| `description` | `String` | Human-readable description surfaced in agent metadata | +| `model` | `String` | Model ID; falls back to your Genkit default model | +| `system` | `String` | System prompt sent to the model every turn | +| `tools` | `List>` / varargs | Tools the model may call during a turn | +| `config` | `GenerationConfig` | Temperature, max output tokens, and other generation settings | +| `maxTurns` | `Integer` | Cap on tool-calling turns within a single turn | +| `store` | `SessionStore` | Persist the session server-side; omit for client-managed | +| `stateType` | `Class` | Java class for your custom session state | +| `clientTransform` | `ClientTransform` | Adjusts outgoing state before it reaches the caller (client-managed) | +| `promptName` | `String` | Name of a registered prompt for `definePromptAgent` | +| `promptInput` | `Object` | Input interpolated into the prompt template each turn | + +If you set `stateType` to a typed class, the agent's registered metadata includes a generated JSON schema for that state, which makes the shape discoverable to the Dev UI and schema-aware clients. A bare `Map` produces no schema, since an untyped map has no useful shape. + +### Prompt-backed agents: `definePromptAgent` + +`definePromptAgent` is a variant of `defineAgent` that draws its system instructions from a registered prompt instead of the `system` string. The prompt template is rendered **every turn**, interpolating `promptInput` and the current session state, so template variables like `{{topic}}` are filled in with live values for that turn. If no matching prompt is found (or its template is blank), it falls back to the `system` field. + +```java +Agent> supportAgent = genkit.beta().definePromptAgent( + AgentConfig.>builder() + .name("supportAgent") + .promptName("supportSystem") // a registered prompt + .promptInput(Map.of("product", "Genkit")) + .model("openai/gpt-4o-mini") + .store(new FileSessionStore<>("./.snapshots")) + .build()); +``` + +A multi-message prompt is rendered into a single system instruction: variable interpolation works, but a `.prompt` file with several message roles isn't split into separate messages. + +## Custom agents: `defineCustomAgent` and `AgentFn` + +When a model-backed turn isn't the right shape — a deterministic workflow, an agent that only sometimes calls a model, or one that calls a model in a way `defineAgent` doesn't expose — write the turn yourself. + +```java +import com.google.genkit.ai.agent.CustomAgentConfig; +import com.google.genkit.ai.agent.AgentFn; +import com.google.genkit.ai.agent.AgentResult; +import com.google.genkit.ai.agent.AgentFinishReason; +import com.google.genkit.ai.agent.InMemorySessionStore; +import com.google.genkit.ai.Message; +import com.google.genkit.ai.Role; +import java.util.List; +import java.util.Map; + +AgentFn> echoFn = (sess, fnCtx) -> { + List msgs = sess.getMessages(); + String userText = ""; + for (int i = msgs.size() - 1; i >= 0; i--) { + if (msgs.get(i).getRole() == Role.USER) { + userText = msgs.get(i).getText(); + break; + } + } + return AgentResult.builder() + .message(Message.model("echo: " + userText)) + .finishReason(AgentFinishReason.STOP) + .build(); +}; + +Agent> echoAgent = genkit.beta().defineCustomAgent( + CustomAgentConfig.>builder() + .name("docsCustomAgent") + .description("A no-model custom agent") + .store(new InMemorySessionStore<>()) + .build(), + echoFn); +``` + +This agent never calls a model. An `AgentFn` is a plain function that runs once per turn and returns an `AgentResult`. + +### `CustomAgentConfig` fields + +| Field | Type | Description | +|-------|------|-------------| +| `name` | `String` (required) | Registered name | +| `description` | `String` | Human-readable description | +| `stateType` | `Class` | Java class for your custom session state | +| `store` | `SessionStore` | Persist the session server-side; omit for client-managed | +| `clientTransform` | `ClientTransform` | Adjusts outgoing state before it reaches the caller (client-managed) | +| `storeOptions` | `SessionStoreOptions` | Options forwarded to every store operation (e.g. multi-tenant scoping) | + +### Writing an `AgentFn` + +```java +@FunctionalInterface +public interface AgentFn { + AgentResult run(SessionRunner sess, AgentFnContext ctx) throws Exception; +} +``` + +Your function receives two objects. `SessionRunner` lets you read and update the session for this turn; `AgentFnContext` gives you the stream sink, the abort signal, the request context, and any resume payload from the caller. + +#### Reading and updating the session — `SessionRunner` + +| Method | Description | +|--------|-------------| +| `sessionId()` | The session's ID | +| `getMessages()` | The conversation so far, including this turn's user message | +| `addMessages(Message...)` | Append extra messages beyond the reply you return | +| `getCustom()` | The current custom state `S` | +| `updateCustom(UnaryOperator)` | Apply a function to the custom state; also streams the change to streaming callers | +| `getArtifacts()` / `addArtifacts(Artifact...)` | Read or append turn artifacts | +| `turnIndex()` | Number of turns completed so far | +| `lastSnapshotId()` / `lastSnapshot()` | The most recently persisted snapshot (server-managed) | +| `lastTurnFinishReason()` / `lastTurnError()` | Outcome of the previous turn | + +A typical `AgentFn` reads `sess.getMessages()` for context, does its work, optionally calls `sess.updateCustom(...)` or `sess.addArtifacts(...)`, and returns the reply message in an `AgentResult`. The runtime appends that reply to the session for you, so use `addMessages(...)` only when you need to add messages *beyond* the reply. + +#### Streaming, abort, context, and resume — `AgentFnContext` + +| Method | Description | +|--------|-------------| +| `sendChunk()` | The sink for pushing chunks (model text, artifacts, …) to a `sendStream` caller mid-turn | +| `isAborted()` | Whether this turn has been asked to stop (see [Background Execution](../background-execution)) | +| `context()` | The request context, carrying caller info such as auth and request headers | +| `resume()` | The caller's response to a pending interrupt, or `null` when this isn't a resume turn | + +Tools and other code you call synchronously from inside an `AgentFn` can reach the active session through `AgentSessionContext.current()` without you threading it through every call, which is how the [multi-agent delegation](../multi-agent-delegation) tools cooperate. + +### Returning an `AgentResult` + +```java +AgentResult.builder() + .message(assistantReply) // appended to session history for you + .artifacts(extraArtifacts) // optional; merged into the session + .finishReason(AgentFinishReason.STOP) // defaults to STOP + .build(); +``` + +If your `AgentFn` throws instead of returning, the turn is recorded as failed rather than raising to the caller: `send`/`sendStream` return an `AgentResponse` whose `finishReason()` is `FAILED`. See [Error Handling](../error-handling) for the full picture. + +## Where session state lives + +Both `defineAgent` and `defineCustomAgent` choose server- vs. client-managed the same way: pass a `SessionStore` to `.store(...)` to persist server-side, or omit it to round-trip the full state through `AgentChat` each turn. See [Agents Overview](../overview#where-session-state-lives) and [Session Stores](../session-stores) for the available stores. + +## See also + +- [Agents Overview](../overview) — Quick-start and the beta opt-in flag +- [Run and Stream](../run-and-stream) — `send`/`sendStream`, `AgentResponse`, and chunk shapes +- [Sessions](../sessions) — Session lifecycle, snapshots, and resuming +- [Session Stores](../session-stores) — `InMemorySessionStore`, `FileSessionStore`, Firestore +- [Interrupts](../interrupts) — Pausing agents for human-in-the-loop confirmation +- [Multi-agent Delegation](../multi-agent-delegation) — Composing multiple agents together +- [Custom Orchestration](../custom-orchestration) — Driving agents outside of `AgentChat` +- [Error Handling](../error-handling) — Failure modes and how they surface diff --git a/docs/src/content/docs/agents/error-handling.md b/docs/src/content/docs/agents/error-handling.md new file mode 100644 index 000000000..25fd62b88 --- /dev/null +++ b/docs/src/content/docs/agents/error-handling.md @@ -0,0 +1,100 @@ +--- +title: Error Handling +description: When a turn returns a FAILED response versus when chat.send() throws, for both foreground and background turns. +--- + +Agents draw a clear line between two kinds of error, and knowing which is which tells you whether to check a return value or wrap a call in `try`/`catch`: + +- **A turn that fails while running** — your `AgentFn` or a tool it calls throws — never propagates out of `chat.send(...)`. It comes back as a normal `AgentResponse` with `finishReason() == FAILED` and an error attached. +- **Malformed input** — misusing the API, e.g. building an `AgentInput` with the wrong shape — throws a `GenkitException` before any turn runs. + +The rule of thumb: **a running turn reports failure in its response; API misuse throws.** + +## A failing turn returns FAILED + +You don't need a `try`/`catch` around `chat.send(...)` to catch a turn that throws — check `finishReason()` instead: + +```java +import com.google.genkit.ai.agent.AgentFn; +import com.google.genkit.ai.agent.AgentResponse; +import com.google.genkit.ai.agent.AgentFinishReason; +import com.google.genkit.ai.agent.RuntimeError; + +AgentFn> throwingFn = (sess, ctx) -> { + throw new RuntimeException("boom: downstream service unavailable"); +}; + +Agent> agent = genkit.beta().defineCustomAgent( + CustomAgentConfig.>builder() + .name("flaky") + .store(new InMemorySessionStore<>()) + .build(), + throwingFn); + +// No try/catch needed — this returns normally. +AgentResponse> resp = agent.chat().send("do something"); + +if (resp.finishReason() == AgentFinishReason.FAILED) { + RuntimeError err = resp.raw().getError(); + System.out.println("Turn failed: " + err.getMessage()); // "boom: downstream service unavailable" + // decide whether to retry, surface to the user, etc. +} +``` + +`RuntimeError` carries a short `status` code (such as `"INTERNAL"`), a `message`, and optional `details`. This is the same for prompt-backed agents: if the model call or a tool fails, the turn comes back `FAILED` rather than throwing. + +## What does throw: malformed input + +Building input by hand with an invalid shape is treated as a programming error and throws right away, before any turn runs: + +```java +// A tool-response part on a user message (those must go through chat.resume(...)): +Part badPart = new Part(); +badPart.setToolResponse(new ToolResponse()); +Message badMessage = new Message(Role.USER, List.of(badPart)); + +// Throws GenkitException (INVALID_ARGUMENT) — not a FAILED response. +chat.send(AgentInput.builder().message(badMessage).build()); +``` + +The rejected cases are: a message with a role other than `USER`, and a user message that carries `toolRequest`/`toolResponse` parts (those belong in `chat.resume(...)`). You only need to guard against this if you construct `AgentInput` yourself from untrusted shapes — the plain `chat.send(String)` / `chat.sendStream(String, ...)` overloads always build valid input. + +## Resume that references a stale tool call + +Resolving an interrupt that isn't actually pending — resuming a tool call that was never interrupted, or restarting with input that doesn't match the tool's schema — fails cleanly with an `INVALID_ARGUMENT` error rather than crashing the turn. Treat it the same as any other malformed-input case. See [Interrupts](../interrupts) for the normal resume/restart flow. + +## Background (detached) turn failures + +A [detached turn](../background-execution) has already returned to the caller before it runs, so a failure can't come back through the original call — it surfaces on the persisted snapshot instead. The immediate response is `DETACHED`; poll the snapshot and check its status: + +```java +AgentResponse> immediate = chat.send( + AgentInput.builder() + .message(Message.user("start the long job")) + .detach(true) + .build()); + +String snapshotId = immediate.snapshotId(); + +SessionSnapshot> snap; +do { + Thread.sleep(100); + snap = agent.getSnapshotData(GetSnapshotRequest.builder().snapshotId(snapshotId).build()); +} while (snap.getStatus() == SnapshotStatus.PENDING); + +if (snap.getStatus() == SnapshotStatus.FAILED) { + System.out.println("Detached turn failed: " + snap.getError().getMessage()); +} +``` + +So background failures are handled where background successes are — at the poll site, by checking the snapshot status. + +## See also + +- [Agents Overview](../overview) — Defining agents and the chat API +- [Define Agents](../define-agents) — `AgentConfig`, `CustomAgentConfig`, and the `AgentFn` contract +- [Run and Stream](../run-and-stream) — `AgentResponse`, `finishReason()`, and the response surface +- [Sessions](../sessions) — Session lifecycle, snapshots, and abort +- [Interrupts](../interrupts) — How `INTERRUPTED` differs from `FAILED` +- [Background Execution](../background-execution) — Detached turns and how failures surface on a snapshot +- [Custom Orchestration](../custom-orchestration) — Writing the `AgentFn` whose failures this page describes diff --git a/docs/src/content/docs/agents/interrupts.md b/docs/src/content/docs/agents/interrupts.md new file mode 100644 index 000000000..5ad220554 --- /dev/null +++ b/docs/src/content/docs/agents/interrupts.md @@ -0,0 +1,131 @@ +--- +title: Interrupts +description: Pause an agent turn for human-in-the-loop confirmation, then resume it with the caller's real response. +--- + +An **interrupt** pauses a turn so a human (or another system) can approve, reject, or supply data before the agent continues. A turn ends with `AgentFinishReason.INTERRUPTED` when a tool the model called needs input instead of returning a result. You inspect the pending interrupts, gather the answer, and then call `chat.resume(...)` to feed the tool's response back so the model can finish its work. + +This works the same way whether the agent is prompt-backed (`defineAgent`/`definePromptAgent`) or a custom `AgentFn`. + +## A complete example: a confirmation gate + +Define an interrupt tool with `genkit.defineInterrupt(...)`. When the model calls it, the turn pauses instead of executing: + +```java +import com.google.genkit.ai.InterruptConfig; +import com.google.genkit.ai.Tool; +import com.google.genkit.ai.Part; +import com.google.genkit.ai.agent.Agent; +import com.google.genkit.ai.agent.AgentChat; +import com.google.genkit.ai.agent.AgentConfig; +import com.google.genkit.ai.agent.AgentFinishReason; +import com.google.genkit.ai.agent.AgentInterrupt; +import com.google.genkit.ai.agent.AgentResponse; +import java.util.List; +import java.util.Map; + +Tool confirmTransfer = + genkit.defineInterrupt( + InterruptConfig.builder() + .name("confirmTransfer") + .description("Ask for confirmation before transferring money.") + .inputType(TransferRequest.class) + .outputType(ConfirmationOutput.class) + .inputSchema(Map.of("type", "object", "properties", Map.of( + "recipient", Map.of("type", "string"), + "amount", Map.of("type", "number")))) + .build()); + +Agent> bankingAgent = genkit.beta().defineAgent( + AgentConfig.>builder() + .name("bankingAgent") + .system("You are a banking assistant. Use confirmTransfer before any transfer.") + .tools(confirmTransfer) + .model("openai/gpt-4o-mini") + .build()); + +AgentChat> chat = bankingAgent.chat(); + +// 1. The turn pauses when the model calls confirmTransfer. +AgentResponse> res = chat.send("Transfer $150 to Alice for dinner"); + +if (res.finishReason() == AgentFinishReason.INTERRUPTED) { + AgentInterrupt interrupt = res.interrupts().get(0); + System.out.println("Tool: " + interrupt.name()); // "confirmTransfer" + System.out.println("Input: " + interrupt.input()); // {recipient=Alice, amount=150.0} + + // 2. Ask the human, then build the tool's response and resume. + ConfirmationOutput approved = new ConfirmationOutput(true, "User approved"); + Part responsePart = confirmTransfer.respond(interrupt.part(), approved); + + AgentResponse> resumed = chat.resume(List.of(responsePart)); + System.out.println(resumed.text()); // e.g. "Done — I've transferred $150 to Alice." + System.out.println(resumed.finishReason()); // STOP +} +``` + +## Inspecting interrupts + +When `finishReason() == AgentFinishReason.INTERRUPTED`, `res.interrupts()` returns one `AgentInterrupt` per tool the model wanted to run but couldn't. Each one exposes: + +| Method | Returns | +|--------|---------| +| `name()` | The interrupted tool's name | +| `input()` | The input the model wanted to call the tool with | +| `part()` | The originating tool-request part — pass this to `respond(...)` / `restart(...)` | + +A single turn can raise several interrupts at once; collect a response for each and pass them all to `chat.resume(List.of(...))` together. + +## Resuming with a response + +`chat.resume(responseParts)` supplies the tool's output and continues the turn. Build each part with the interrupted tool's `respond(...)` helper: + +```java +Part responsePart = confirmTransfer.respond(interrupt.part(), new ConfirmationOutput(true, "ok")); +AgentResponse> resumed = chat.resume(List.of(responsePart)); +``` + +The model genuinely sees the tool response and reasons about it — it continues the same conversation rather than starting over, and the tool response is recorded in the session history. + +## Restarting a tool call + +Sometimes you don't want to answer the tool — you want to run it again with corrected input or an approval stamp. `chat.restart(restartParts)` re-executes the interrupted tool call. Build each part with the tool's `restart(...)` helper, optionally passing replacement input: + +```java +// Retry the transfer with a corrected amount and an "approved" marker. +Part restartPart = confirmTransfer.restart( + interrupt.part(), + Map.of("approved", true), // resumed metadata + new TransferRequest("Alice", 120.0)); // replacement input + +AgentResponse> resumed = chat.restart(List.of(restartPart)); +``` + +Inside the tool handler, a restart-aware tool can tell it's being re-invoked and read the metadata you supplied via the tool context: + +```java +(ctx, input) -> { + if (ctx.isResumed()) { + // ctx.getResumed() holds the resumed metadata (e.g. {"approved": true}). + // input is the replacement input if you passed one to restart(...). + return doTransfer(input); + } + throw new ToolInterruptException(); // fresh call → pause for confirmation +}; +``` + +Use `respond(...)` when you're answering the tool's question; use `restart(...)` when you're re-running it with new input or an approval decision. + +## Limitations + +- A synchronous `chat.send(...)` / `chat.resume(...)` turn runs to completion — you can't cancel it midway. Cooperative cancellation applies to [detached (background) turns](../background-execution). +- Referencing a tool call that isn't actually pending, or restarting with input that doesn't match the schema, fails cleanly with an `INVALID_ARGUMENT` error rather than crashing the turn — see [Error Handling](../error-handling). + +## See also + +- [Agents Overview](../overview) — Defining agents and the chat API +- [Define Agents](../define-agents) — `AgentConfig`, `CustomAgentConfig`, and the `AgentFn` contract +- [Run and Stream](../run-and-stream) — `AgentResponse`, `finishReason()`, and the response surface +- [Multi-agent Delegation](../multi-agent-delegation) — How an interrupted sub-agent surfaces to a parent +- [Error Handling](../error-handling) — How `FAILED` turns differ from `INTERRUPTED` ones +- [Custom Orchestration](../custom-orchestration) — Reading resume data in your own `AgentFn` diff --git a/docs/src/content/docs/agents/multi-agent-delegation.md b/docs/src/content/docs/agents/multi-agent-delegation.md new file mode 100644 index 000000000..6d85ebf12 --- /dev/null +++ b/docs/src/content/docs/agents/multi-agent-delegation.md @@ -0,0 +1,102 @@ +--- +title: Multi-agent Delegation +description: Let one agent delegate self-contained tasks to specialized sub-agents by exposing them as tools. +--- + +Delegation lets a *parent* agent hand a self-contained task to a specialized *sub-agent*. The middleware plugin's `Agents` helper turns any set of already-registered agents into ordinary tools the parent can call — the parent model decides when to delegate, and the sub-agent's answer flows back into the parent's reasoning like any other tool result. + +## Installation + +```xml + + com.google.genkit + genkit-plugin-middleware + ${genkit.version} + +``` + +## Setting up delegation + +Register your sub-agents first, then build delegation tools from their names and pass them to the parent: + +```java +import com.google.genkit.ai.Tool; +import com.google.genkit.ai.agent.Agent; +import com.google.genkit.ai.agent.AgentConfig; +import com.google.genkit.plugins.middleware.Agents; +import com.google.genkit.plugins.middleware.AgentsOptions; +import java.util.List; +import java.util.Map; + +// researcher and writer are already registered via genkit.beta().defineAgent(...) etc. +AgentsOptions options = AgentsOptions.builder() + .agents("researcher", "writer") + .build(); + +List> delegationTools = Agents.delegationTools(options); +String subAgentsFragment = Agents.systemPromptFragment(options); + +Agent> orchestrator = genkit.beta().defineAgent( + AgentConfig.>builder() + .name("orchestrator") + .system("You coordinate a research task.\n\n" + subAgentsFragment) + .tools(delegationTools.toArray(new Tool[0])) + .model("openai/gpt-4o-mini") + .build()); +``` + +Each sub-agent becomes a tool named `_` — with the default prefix that's `delegate_to_researcher`, `delegate_to_writer`, and so on. The parent calls a delegation tool with a `{"task": "..."}` payload; the sub-agent runs one turn against that task and returns its answer. `Agents.systemPromptFragment(options)` gives you a ready-made description of the available sub-agents to splice into the parent's system prompt. + +## Options + +| Field | Default | Description | +|-------|---------|-------------| +| `agents` | *(required)* | Sub-agent names to expose, at least one | +| `toolPrefix` | `"delegate_to"` | Tool-name prefix; an empty string uses the bare agent name | +| `maxDelegations` | `0` (unlimited) | Cap on total delegation calls across all tools from one `delegationTools(...)` call | +| `historyLength` | `0` (task only) | How many trailing prior messages of this sub-agent's own conversation to forward on each call | +| `artifactStrategy` | `ArtifactStrategy.INLINE` | Whether a delegated artifact's content is inlined in the tool output or referenced by name only | + +## Accumulating sub-agent history + +By default each delegation forwards only the current `task`. Set `historyLength` to let a sub-agent build up its own conversation across repeated delegations *within the same parent conversation*: + +```java +AgentsOptions options = AgentsOptions.builder() + .agents("researcher") + .historyLength(10) // forward up to 10 trailing messages of prior delegation context + .build(); +``` + +With this set, the second time the parent delegates to `researcher` it sees the previous task and answer plus the new task, and so on — trimmed to the last `historyLength` messages. Each (parent conversation, sub-agent) pair keeps its own isolated history, so delegating to `researcher` never leaks into `writer`, and a different parent conversation starts fresh. + +History trimming applies to client-managed sub-agents (no `SessionStore`). A server-managed sub-agent resumes its full stored history on each call; if you need bounded context there, have each `task` restate what the sub-agent needs. + +## Interrupted and failed sub-agents + +A sub-agent isn't a black box — its outcome surfaces to the parent's tool loop as a real event, not swallowed text: + +- **Interrupt.** If a sub-agent pauses for confirmation (see [Interrupts](../interrupts)), the delegation tool raises that interrupt to the parent. The parent's own turn finishes `INTERRUPTED`, and `AgentResponse.interrupts()` reports it — exactly as if the parent had called an interrupting tool directly. +- **Failure.** If a sub-agent's turn fails, the delegation tool raises a real error carrying the sub-agent's message, which propagates through the parent's tool loop like any other tool exception. + +One thing to keep in mind: resolving the *parent's* interrupt with `chat.resume(...)` acknowledges that a sub-agent is paused, but does not automatically re-drive the sub-agent's own paused turn. If a delegated sub-agent might interrupt and you need the human's answer to reach it, have the sub-agent resolve its own confirmation (for example against a default policy or its own tracked state) rather than expecting the answer to thread all the way down. + +## Artifacts across the delegation boundary + +Any artifact a sub-agent produces is namespaced (as `/`) and merged into the parent's artifacts, so it shows up in the parent's own `AgentResponse.artifacts()` and — for a server-managed parent — in its session store. With the default `ArtifactStrategy.INLINE`, the artifact's content is also folded into the delegation tool's result so the parent model can read it directly. Switch to `ArtifactStrategy.SESSION` to return only the namespaced name and read the content back from the parent's session yourself: + +```java +AgentsOptions options = AgentsOptions.builder() + .agents("writer") + .artifactStrategy(ArtifactStrategy.SESSION) // names only, no inline content + .build(); +``` + +## See also + +- [Agents Overview](../overview) — Defining agents and the chat API +- [Define Agents](../define-agents) — `AgentConfig`, `CustomAgentConfig`, and the `AgentFn` contract +- [Interrupts](../interrupts) — What "interrupted" means for a single agent +- [Sessions](../sessions) — Session lifecycle, snapshots, and the turn model +- [Custom Orchestration](../custom-orchestration) — Driving multiple agents yourself instead of via delegation tools +- [Error Handling](../error-handling) — The general `FAILED`/error contract delegation builds on diff --git a/docs/src/content/docs/agents/overview.md b/docs/src/content/docs/agents/overview.md new file mode 100644 index 000000000..e26dde4bb --- /dev/null +++ b/docs/src/content/docs/agents/overview.md @@ -0,0 +1,179 @@ +--- +title: Agents Overview +description: Build stateful, multi-turn AI agents with tool-calling and session persistence using the Genkit Agents API. +--- + +Agents are stateful, multi-turn AI actors. Unlike a single `generate` call, an agent carries conversation history across turns, calls tools, and can persist its session so a conversation survives across requests, processes, or machines. The Agents API is in **beta**, so you opt in explicitly. + +## Enabling the beta API + +Enable experimental features when you build `Genkit`, then reach the agent factories through `genkit.beta()`: + +```java +Genkit genkit = Genkit.builder() + .options(GenkitOptions.builder().experimental(true).build()) + .plugin(OpenAIPlugin.create()) + .build(); +``` + +You can also set `GENKIT_EXPERIMENTAL=true` in the environment before your process starts. Without the flag, calling `genkit.beta()` throws. + +## Defining an agent + +Use `defineAgent` to register a model-backed agent. The only required field is `name`; everything else is optional. + +```java +import com.google.genkit.agent.AgentConfig; +import com.google.genkit.ai.agent.Agent; +import com.google.genkit.ai.agent.FileSessionStore; + +Agent> weatherAgent = genkit.beta().defineAgent( + AgentConfig.>builder() + .name("weatherAgent") + .description("A helpful weather assistant") + .system("You are a helpful weather assistant. Use the getWeather tool.") + .tools(getWeather) // Tool objects from genkit.defineTool + .model("openai/gpt-4o-mini") + .store(new FileSessionStore<>("./.snapshots")) // persist the session; omit to keep it client-side + .build()); +``` + +The options you'll reach for most often: + +| Field | Description | +|-------|-------------| +| `name` | Registered name; also how the agent appears in the Dev UI | +| `system` | System prompt sent at the start of every turn | +| `tools` | Tools the model may call during a turn | +| `model` | Model ID; falls back to your Genkit default model | +| `config` | Generation settings such as temperature and max tokens | +| `maxTurns` | Cap on tool-calling turns within a single `send` | +| `store` | A `SessionStore` to persist the session server-side; omit for client-managed state | +| `stateType` | Java class for your custom session state | + +See [Define agents](../define-agents) for the full configuration surface and for writing your own turn logic. + +## Starting a chat + +`Agent.chat(...)` returns an `AgentChat` that manages turn-by-turn history for you. + +```java +AgentChat> chat = weatherAgent.chat(); +``` + +To continue an earlier conversation, seed the chat with a `snapshotId` or `sessionId`: + +```java +import com.google.genkit.ai.agent.AgentInit; + +AgentChat> chat = weatherAgent.chat( + AgentInit.>builder() + .snapshotId("existing-snapshot-id") + .build()); +``` + +## Sending messages + +Call `send` for a blocking turn, or `sendStream` to observe chunks as the turn runs. Both return the same `AgentResponse`, and both automatically carry state forward to the next turn. + +```java +AgentResponse> res = chat.send("What is the weather in London?"); +System.out.println(res.text()); // "It is sunny and 22°C in London." + +// History carries forward automatically: +AgentResponse> res2 = chat.send("Now say that in French"); +System.out.println(res2.text()); +``` + +```java +chat.sendStream("Summarize in one sentence", chunk -> { + if (chunk.text() != null) { + System.out.print(chunk.text()); // incremental tokens + } +}); +``` + +`AgentResponse` gives you everything about the completed turn: + +| Method | Returns | +|--------|---------| +| `text()` | The assistant message's text | +| `finishReason()` | Why the turn ended (`STOP`, `LENGTH`, `INTERRUPTED`, …) | +| `snapshotId()` | This turn's snapshot ID (server-managed) | +| `sessionId()` | The session ID | +| `custom()` | Your custom session state after the turn | +| `artifacts()` | Any artifacts produced this turn | +| `interrupts()` | Pending interrupts when the turn paused for input | + +See [Run and stream](../run-and-stream) for the full `AgentResponse` surface and every chunk shape. + +## Where session state lives + +You choose where an agent's session is stored: + +| Mode | How to enable | Where state lives | +|------|---------------|-------------------| +| **Server-managed** | Pass a `SessionStore` to `.store(...)` | The store (memory, disk, Firestore, …) | +| **Client-managed** | Omit `.store(...)` | Carried by `AgentChat`, round-tripped each turn | + +Server-managed sessions are the norm when you want conversations to persist or to run across processes. Client-managed sessions are handy for stateless servers or when you'd rather own the conversation history yourself — `AgentChat` carries the messages and custom state along on every turn, so the model always sees the full context. + +```java +// Client-managed (stateless) agent — no .store(...) +Agent> statelessAgent = genkit.beta().defineAgent( + AgentConfig.>builder() + .name("weatherAgentStateless") + .system("You are a helpful weather assistant.") + .tools(getWeather) + .model("openai/gpt-4o-mini") + .build()); + +AgentChat> statelessChat = statelessAgent.chat(ctx); +statelessChat.send("What is the weather in Tokyo?"); +statelessChat.send("Compare that to Paris"); // history round-trips automatically +``` + +## Custom agents (no model required) + +`defineAgent` builds a model-backed turn for you. When you want full control over what happens on each turn — a deterministic workflow, or an agent that only sometimes calls a model — use `defineCustomAgent` and write the turn logic yourself as an `AgentFn`: + +```java +import com.google.genkit.ai.agent.CustomAgentConfig; +import com.google.genkit.ai.agent.AgentFn; +import com.google.genkit.ai.agent.AgentResult; +import com.google.genkit.ai.agent.AgentFinishReason; +import com.google.genkit.ai.agent.InMemorySessionStore; +import com.google.genkit.ai.Message; + +AgentFn> echoFn = (sess, fnCtx) -> { + String userText = sess.getMessages().get(sess.getMessages().size() - 1).getText(); + return AgentResult.builder() + .message(Message.model("echo: " + userText)) + .finishReason(AgentFinishReason.STOP) + .build(); +}; + +Agent> echoAgent = genkit.beta().defineCustomAgent( + CustomAgentConfig.>builder() + .name("echoAgent") + .store(new InMemorySessionStore<>()) + .build(), + echoFn); +``` + +See [Define agents](../define-agents) for the full `AgentFn` and `SessionRunner` API. + +## Full example + +See [`samples/agents-weather`](https://github.com/genkit-ai/genkit-java/tree/main/samples/agents-weather) for a complete, runnable example that combines both modes and demonstrates streaming. + +## See also + +- [Define Agents](../define-agents) — `AgentConfig`, `CustomAgentConfig`, and writing an `AgentFn` +- [Run and Stream](../run-and-stream) — `send`/`sendStream`, `AgentResponse`, and chunk shapes +- [Sessions](../sessions) — Session lifecycle, snapshots, and resuming +- [Session Stores](../session-stores) — `InMemorySessionStore`, `FileSessionStore`, Firestore +- [Interrupts](../interrupts) — Pausing agents for human-in-the-loop confirmation +- [Multi-agent Delegation](../multi-agent-delegation) — Composing multiple agents together +- [Custom Orchestration](../custom-orchestration) — Driving agents outside of `AgentChat` +- [Error Handling](../error-handling) — Failure modes and how they surface diff --git a/docs/src/content/docs/agents/run-and-stream.md b/docs/src/content/docs/agents/run-and-stream.md new file mode 100644 index 000000000..24d6dfddd --- /dev/null +++ b/docs/src/content/docs/agents/run-and-stream.md @@ -0,0 +1,187 @@ +--- +title: Run and Stream +description: Drive an agent turn-by-turn with AgentChat.send and sendStream, and read the AgentResponse and stream chunks. +--- + +The [Agents Overview](../overview) shows the basic `send` and `sendStream` calls. This page covers both methods in depth, the full `AgentResponse` surface, and every kind of chunk `sendStream` can deliver. + +## `send` vs. `sendStream` + +Both methods run exactly **one turn** and return an `AgentResponse`. The only difference is whether you observe chunks as the turn runs — `send` gives you the final result, `sendStream` also fires a callback for each chunk along the way. + +```java +// Blocking: you only see the final response. +AgentResponse> res = chat.send("What is the weather in London?"); + +// Streaming: a callback fires per chunk, and you still get the same final response. +AgentResponse> res2 = chat.sendStream( + "What is the weather in Paris?", + chunk -> { + if (chunk.text() != null) { + System.out.print(chunk.text()); + } + }); +``` + +Both accept either a plain `String` (wrapped into a user message) or a fully-formed `AgentInput`: + +```java +import com.google.genkit.ai.agent.AgentInput; +import com.google.genkit.ai.Message; + +AgentResponse> res = chat.send( + AgentInput.builder().message(Message.user("Custom-built input")).build()); +``` + +## Reading the result — `AgentResponse` + +Both methods return the same `AgentResponse` once the turn completes: + +| Method | Returns | +|--------|---------| +| `text()` | The assistant message's text (`""` if none) | +| `message()` | The raw assistant `Message`, or `null` | +| `toolRequests()` | Tool-request parts on the assistant message | +| `interrupts()` | Pending interrupts — populated when `finishReason()` is `INTERRUPTED` | +| `finishReason()` | Why the turn ended: `STOP`, `LENGTH`, `BLOCKED`, `INTERRUPTED`, `ABORTED`, `DETACHED`, `FAILED`, `OTHER`, `UNKNOWN` | +| `snapshotId()` | This turn's snapshot ID (server-managed); `null` for client-managed | +| `sessionId()` | The session ID | +| `custom()` | Your custom state after the turn (populated for both modes) | +| `state()` | The full `SessionState` — only for client-managed agents; server-managed agents return `null` here (use `snapshotId()` with a store lookup instead) | +| `artifacts()` | Artifacts produced this turn | +| `raw()` | The underlying `AgentOutput` | + +```java +AgentResponse> res = chat.send("What is 2+2?"); +System.out.println(res.text()); // "The answer is 4." +System.out.println(res.finishReason()); // stop +System.out.println(res.snapshotId()); // "1a2b3c…" (server-managed) or null (client-managed) +System.out.println(res.custom()); // whatever your agent tracked, e.g. {"turns": 1} +``` + +## Streaming chunks + +The callback you pass to `sendStream` receives an `AgentChunk` for each chunk, in order: + +| Method | Returns | +|--------|---------| +| `text()` | Streamed model text, or `null` if this chunk carries no model text | +| `modelChunk()` | The raw model chunk, or `null` | +| `artifact()` | An artifact update, or `null` | +| `custom()` | Your custom state with this chunk's change applied, or `null` | +| `raw()` | The underlying chunk (exposes the turn-end signal) | + +A single turn can produce a mix of chunks. Each chunk also carries a role via its underlying model chunk, so you can tell model output apart from other content. + +### Model chunks + +The agent streams model output as it's generated. For model-backed agents this is every token; for custom agents it's whatever you push through `ctx.sendChunk()`. + +```java +chat.sendStream("Summarize this in one sentence", chunk -> { + if (chunk.text() != null) { + System.out.print(chunk.text()); // incremental tokens + } +}); +``` + +### Tool-response chunks + +When the agent calls a tool during a turn, the tool's response streams too, so a client can show tool activity as it happens rather than waiting for the final message. + +### Custom-state chunks + +Each `updateCustom(...)` call inside an `AgentFn` streams the updated custom state. You read the result — the new state — directly through `chunk.custom()`. + +```java +chat.sendStream("Do the thing", chunk -> { + if (chunk.custom() != null) { + System.out.println("Custom state now: " + chunk.custom()); + } +}); +``` + +### Artifact chunks + +When an `AgentFn` pushes an artifact update mid-turn, it arrives as an artifact chunk. (Artifacts added with `sess.addArtifacts(...)` instead attach to the final response rather than streaming.) + +```java +chat.sendStream("Generate a report", chunk -> { + if (chunk.artifact() != null) { + System.out.println("Artifact update: " + chunk.artifact().getName()); + } +}); +``` + +### Turn-end chunk + +Every turn ends with a single chunk carrying a turn-end signal — your cue that no more chunks are coming — with the turn's final `snapshotId` and `finishReason`. This fires even when the turn failed, which is useful for a UI that needs to know exactly when to stop showing a "thinking" indicator. + +```java +import com.google.genkit.ai.agent.TurnEnd; + +chat.sendStream("hi", chunk -> { + TurnEnd end = chunk.raw().getTurnEnd(); + if (end != null) { + System.out.println("Turn ended: " + end.getFinishReason() + " @ " + end.getSnapshotId()); + } +}); +``` + +## A custom agent that emits every chunk kind + +A no-model `AgentFn` that emits all four chunk kinds in one turn: + +```java +AgentFn> fn = (sess, fnCtx) -> { + // Model-style chunks + fnCtx.sendChunk().accept(AgentStreamChunk.builder() + .modelChunk(ModelResponseChunk.text("Hel")).build()); + fnCtx.sendChunk().accept(AgentStreamChunk.builder() + .modelChunk(ModelResponseChunk.text("lo!")).build()); + + // Custom-state change (streams to the caller) + sess.updateCustom(cur -> { + Map next = cur != null ? new HashMap<>(cur) : new HashMap<>(); + next.put("status", "thinking"); + return next; + }); + + // Artifact chunk (streamed mid-turn) + fnCtx.sendChunk().accept(AgentStreamChunk.builder() + .artifact(Artifact.builder().name("draft").parts(List.of()).build()).build()); + + // The turn-end chunk is emitted for you after this returns. + return AgentResult.builder() + .message(Message.model("Hello!")) + .finishReason(AgentFinishReason.STOP) + .build(); +}; +``` + +Driving it with `chat.sendStream("say hello", ...)` delivers, in order: two model-text chunks (`"Hel"`, `"lo!"`), a chunk whose `custom()` shows `status = "thinking"`, an artifact chunk named `"draft"`, and a final turn-end chunk carrying `STOP` and the same `snapshotId()` as the returned response. + +## Resuming after an interrupt + +When a turn pauses for input (`finishReason() == INTERRUPTED`), `AgentChat.resume(List)` sends the caller's response and runs the next turn. The response is genuinely fed back to the model as part of the conversation, so the agent continues where it left off. + +```java +AgentResponse> interrupted = chat.send("please confirm"); +if (interrupted.finishReason() == AgentFinishReason.INTERRUPTED) { + Part respondPart = new Part(); + respondPart.setText("confirmed"); + AgentResponse> resumed = chat.resume(List.of(respondPart)); +} +``` + +See [Interrupts](../interrupts) for the full human-in-the-loop flow. + +## See also + +- [Define Agents](../define-agents) — `AgentConfig`, `CustomAgentConfig`, and writing an `AgentFn` +- [Agents Overview](../overview) — Quick-start and the beta opt-in flag +- [Sessions](../sessions) — Session lifecycle, snapshots, and resuming +- [Session Stores](../session-stores) — `InMemorySessionStore`, `FileSessionStore`, Firestore +- [Interrupts](../interrupts) — Pausing agents for human-in-the-loop confirmation +- [Background Execution](../background-execution) — Detached turns +- [Error Handling](../error-handling) — Failure modes and how they surface diff --git a/docs/src/content/docs/agents/serve-over-http.md b/docs/src/content/docs/agents/serve-over-http.md new file mode 100644 index 000000000..d9ef3bc82 --- /dev/null +++ b/docs/src/content/docs/agents/serve-over-http.md @@ -0,0 +1,204 @@ +--- +title: Serve over HTTP +description: Expose defined agents as HTTP endpoints with the Jetty or Spring plugins, and call them from any client that speaks the wire format. +--- + +Any agent you define can be served over HTTP by the [Jetty](../../plugins/jetty) or [Spring](../../plugins/spring) plugin. Both plugins speak the **same wire format**, so the same Java client works unchanged against either one. + +## Serving your agents + +Agent endpoints are mounted automatically for every registered agent when the plugin starts — there's no separate opt-in beyond adding the plugin and starting it: + +```java +import com.google.genkit.plugins.jetty.JettyPlugin; +import com.google.genkit.plugins.jetty.JettyPluginOptions; + +JettyPlugin jetty = new JettyPlugin(JettyPluginOptions.builder().port(8080).build()); + +Genkit genkit = Genkit.builder() + .options(GenkitOptions.builder().experimental(true).build()) + .plugin(jetty) + .build(); + +// Define your agent(s) before starting. +Agent> echoAgent = genkit.beta().defineCustomAgent( + CustomAgentConfig.>builder() + .name("echoAgent") + .store(new FileSessionStore<>("./.snapshots")) + .build(), + (sess, fnCtx) -> AgentResult.builder() + .message(Message.model("echo: " + sess.getMessages().get(sess.getMessages().size() - 1).getText())) + .finishReason(AgentFinishReason.STOP) + .build()); + +jetty.start(); // mounts /echoAgent (and its companions) alongside any flows +``` + +Swap in `SpringPlugin.create()` for the same result. Agent endpoints are mounted at the root path (`/`), separate from flow endpoints, on the same port. See [Jetty Server](../../plugins/jetty) and [Spring Boot](../../plugins/spring) for plugin-specific setup. + +## Endpoints + +For an agent named ``, the server mounts: + +| Endpoint | Mounted when | Purpose | +|----------|--------------|---------| +| `POST /` | Always | Runs one turn | +| `POST //getSnapshot` | Server-managed agents | Reads back a stored snapshot by `snapshotId` or `sessionId` | +| `POST //abort` | The store supports change notifications | Marks a pending snapshot aborted | + +A client-managed agent (no store) only mounts the turn endpoint, since there's no server-side snapshot to fetch or abort. + +## Running a turn + +`POST /` runs exactly one turn. The JSON body is an envelope with the turn input, optional resume/session `init`, and an optional `context` map: + +```json +{"data": { /* AgentInput */ }, "init": { /* AgentInit */ }, "context": { /* optional */ }} +``` + +```bash +curl -X POST http://localhost:8080/echoAgent \ + -H "Content-Type: application/json" \ + -d '{"data":{"message":{"role":"user","content":[{"text":"Hello, agent!"}]}}}' +``` + +```json +{"result":{"sessionId":"876d78e6-…","snapshotId":"ca350aa8-…","message":{"text":"echo: Hello, agent!","role":"model","content":[{"text":"echo: Hello, agent!"}]},"finishReason":"stop"}} +``` + +### Streaming + +Request Server-Sent Events with `Accept: text/event-stream` or `?stream=true`: + +```bash +curl -N -X POST "http://localhost:8080/echoAgent?stream=true" \ + -H "Content-Type: application/json" \ + -d '{"data":{"message":{"role":"user","content":[{"text":"Hello again"}]}}}' +``` + +The response is `text/event-stream`: zero or more chunk frames, then exactly one terminal frame carrying the final result: + +``` +data: {"message": } + +data: {"result": } + +``` + +A failed turn sends `data: {"error": {...}}` as the terminal frame instead. How much each frame carries depends on the agent — a model-backed agent streams model text as it's generated, while a custom agent streams only what it pushes through `ctx.sendChunk()`. + +## Reading a snapshot + +`POST //getSnapshot` (server-managed agents only) reads back a stored snapshot. This is the main way to poll a [detached turn](../background-execution) over HTTP. + +```bash +curl -X POST http://localhost:8080/echoAgent/getSnapshot \ + -H "Content-Type: application/json" \ + -d '{"data":{"snapshotId":"ca350aa8-…"}}' +``` + +```json +{"result":{"snapshotId":"ca350aa8-…","sessionId":"876d78e6-…","status":"completed","finishReason":"stop","state":{"messages":[/* … */],"artifacts":[]}}} +``` + +## Aborting a turn + +`POST //abort` marks a pending snapshot aborted: + +```bash +curl -X POST http://localhost:8080/echoAgent/abort \ + -H "Content-Type: application/json" \ + -d '{"data":{"snapshotId":"29b02672-…"}}' +``` + +```json +{"result":{"snapshotId":"29b02672-…","status":"aborted"}} +``` + +This is most useful for [detached turns](../background-execution): a background turn that checks `ctx.isAborted()` can stop early, and once aborted the snapshot stays aborted even if the work finishes. The endpoint needs a store that supports change notifications (`InMemorySessionStore` does not; see [Session Stores](../session-stores)). See [Sessions](../sessions#aborting-a-turn) for the full abort behavior. + +## Error responses + +Errors use a structured envelope: + +```json +{"error":{"status":"INVALID_ARGUMENT","message":"…","details":{"stack":"…"}}} +``` + +`status` maps to an HTTP code the same way on both plugins: + +| `status` | HTTP code | +|----------|-----------| +| `INVALID_ARGUMENT`, `FAILED_PRECONDITION`, `OUT_OF_RANGE` | 400 | +| `UNAUTHENTICATED` | 401 | +| `PERMISSION_DENIED` | 403 | +| `NOT_FOUND` | 404 | +| `ALREADY_EXISTS`, `ABORTED` | 409 | +| `RESOURCE_EXHAUSTED` | 429 | +| `UNIMPLEMENTED` | 501 | +| `UNAVAILABLE` | 503 | +| `DEADLINE_EXCEEDED` | 504 | +| anything else | 500 | + +See [Error Handling](../error-handling) for where these statuses come from. + +## Calling a served agent from Java + +`RemoteAgent.chat(...)` gives you an `AgentChat` backed by HTTP — the same client works against a Jetty- or Spring-backed server: + +```java +import com.google.genkit.client.RemoteAgent; +import com.google.genkit.client.RemoteAgentOptions; +import com.google.genkit.ai.agent.AgentChat; +import com.google.genkit.ai.agent.AgentResponse; + +AgentChat> chat = RemoteAgent.chat( + RemoteAgentOptions.builder() + .url("http://localhost:8080/echoAgent") + .build()); + +AgentResponse> resp = chat.send("hello"); +System.out.println(resp.text()); // "echo: hello" +System.out.println(chat.snapshotId()); // the next send() resumes from here automatically + +chat.abort(); // POSTs to /echoAgent/abort with the tracked snapshotId +``` + +For a client-managed agent, set `serverManaged(false)` so state round-trips in the request instead of being stored server-side: + +```java +RemoteAgentOptions opts = RemoteAgentOptions.builder() + .url("http://localhost:8080/counterAgent") + .serverManaged(false) + .build(); +``` + +### Passing headers + +Headers you set with `RemoteAgentOptions.headers(...)` are sent as real HTTP request headers, and the server makes them available to your `AgentFn` and tools through the request context under a `"headers"` key (standard framing headers like `content-type` and `host` are excluded). This is a natural place to carry an auth token: + +```java +// Client +RemoteAgentOptions opts = RemoteAgentOptions.builder() + .url("http://localhost:8080/echoAgent") + .headers(Map.of("Authorization", "Bearer sk-…")) + .build(); +``` + +```java +// Server — inside an AgentFn or tool handler +Map headers = + (Map) fnCtx.context().getContext().get("headers"); +String token = headers != null ? headers.get("Authorization") : null; +``` + +## See also + +- [Agents Overview](../overview) — Defining agents and the chat API +- [Run and Stream](../run-and-stream) — The `send`/`sendStream` semantics `AgentChat` wraps over HTTP +- [Sessions](../sessions) — Snapshot lifecycle and aborting +- [Session Stores](../session-stores) — Which stores support the `/abort` endpoint +- [Background Execution](../background-execution) — Detaching turns and polling over HTTP +- [Error Handling](../error-handling) — The status codes behind the error envelope +- [Jetty Server](../../plugins/jetty) — Jetty-specific setup and flow endpoints +- [Spring Boot](../../plugins/spring) — Spring-specific setup and flow endpoints diff --git a/docs/src/content/docs/agents/session-stores.md b/docs/src/content/docs/agents/session-stores.md new file mode 100644 index 000000000..a8b721e3f --- /dev/null +++ b/docs/src/content/docs/agents/session-stores.md @@ -0,0 +1,236 @@ +--- +title: Session Stores +description: Persist agent session state with in-memory, file-based, Firestore, DynamoDB, or Cosmos DB session stores. +--- + +A `SessionStore` is where a server-managed agent keeps its sessions. Pass one to `.store(...)` when you define an agent to persist snapshots; omit it to run the agent client-managed (stateless). Several implementations ship out of the box, and you can write your own. + +| Store | Persistence | Good for | +|-------|-------------|----------| +| `InMemorySessionStore` | In-process; lost on restart | Tests and short-lived processes | +| `FileSessionStore` | JSON files on local disk | Local development and single-process deployments | +| `FirestoreSessionStore` | Google Cloud Firestore | Production on Google Cloud (Cloud Run, Firebase Functions) | +| `DynamoDbSessionStore` | Amazon DynamoDB | Production on AWS | +| `CosmosSessionStore` | Azure Cosmos DB | Production on Azure | + +If you plan to use `chat.abort()` or the `/abort` HTTP endpoint, pick a store that supports change notifications: `FileSessionStore`, `FirestoreSessionStore`, `DynamoDbSessionStore`, and `CosmosSessionStore` do; `InMemorySessionStore` does not, so `abort()` is a no-op there. See [Sessions](../sessions#aborting-a-turn). + +## InMemorySessionStore + +Holds snapshots in memory. State is not shared across processes and is discarded on exit. + +```java +import com.google.genkit.ai.agent.InMemorySessionStore; + +Agent> agent = genkit.beta().defineAgent( + AgentConfig.>builder() + .name("myAgent") + .system("You are a helpful assistant.") + .store(new InMemorySessionStore<>()) + .build()); +``` + +## FileSessionStore + +Persists each snapshot as a JSON file under a directory you choose. + +```java +import com.google.genkit.ai.agent.FileSessionStore; + +Agent> agent = genkit.beta().defineAgent( + AgentConfig.>builder() + .name("myAgent") + .system("You are a helpful assistant.") + .store(new FileSessionStore<>("./.snapshots")) + .build()); +``` + +Writes are atomic — a snapshot is written to a temporary file and then renamed — so a process interrupted mid-write never leaves a partial snapshot behind. + +### Options + +Use the builder for finer control: + +```java +FileSessionStore> store = + FileSessionStore.>builder("./.snapshots") + .prefix("prod") // subdirectory to group sessions (default "global") + .maxPersistedChainLength(10) // prune older snapshots in a chain; 0 = keep all + .rejectBranchingSessions(true) // fail if a session ends up with more than one branch + .snapshotWatchPollIntervalMs(500) // how often change subscribers are polled (default 2000) + .build(); +``` + +Snapshots for a session are stored under the prefix subdirectory, with a small pointer file tracking the latest snapshot per session. + +### Watching for changes + +`FileSessionStore` can notify you when a snapshot's status changes — useful for tracking a background turn: + +```java +try (AutoCloseable sub = store.onSnapshotStateChange( + snapshotId, + snap -> System.out.println("Status changed: " + snap.getStatus()), + null)) { + // ... long-running turn ... +} +``` + +## FirestoreSessionStore + +Stores snapshots in Google Cloud Firestore. It lives in the Firebase plugin, so add the dependency: + +```xml + + com.google.genkit + genkit-plugin-firebase + ${genkit.version} + +``` + +```java +import com.google.genkit.plugins.firebase.session.FirestoreSessionStore; +import com.google.cloud.firestore.Firestore; +import com.google.firebase.cloud.FirestoreClient; + +Firestore db = FirestoreClient.getFirestore(); +FirestoreSessionStore> store = + new FirestoreSessionStore<>(db); // uses the default "genkit-sessions" collection + +Agent> agent = genkit.beta().defineAgent( + AgentConfig.>builder() + .name("myAgent") + .system("You are a helpful assistant.") + .store(store) + .build()); +``` + +To customize the collection name or checkpointing behavior, pass `FirestoreSessionStoreOptions` as the second constructor argument. See the [Firebase plugin](../../plugins/firebase#firestore-session-store) for details. + +## DynamoDbSessionStore + +Stores snapshots in Amazon DynamoDB. It lives in the AWS Bedrock plugin, so add the dependency: + +```xml + + com.google.genkit + genkit-plugin-aws-bedrock + ${genkit.version} + +``` + +```java +import com.google.genkit.plugins.awsbedrock.session.DynamoDbSessionStore; +import software.amazon.awssdk.services.dynamodb.DynamoDbClient; + +DynamoDbClient db = DynamoDbClient.create(); // uses the default AWS credential chain and region +DynamoDbSessionStore> store = + new DynamoDbSessionStore<>(db); // uses the default "genkit-sessions" table + +Agent> agent = genkit.beta().defineAgent( + AgentConfig.>builder() + .name("myAgent") + .system("You are a helpful assistant.") + .store(store) + .build()); +``` + +All records live in a single table with a partition key `pk` and sort key `sk` (both strings). Create the table ahead of time, or let the store create it on first use: + +```java +DynamoDbSessionStore> store = + new DynamoDbSessionStore<>( + db, + DynamoDbSessionStoreOptions.builder() + .tableName("genkit-sessions") // table name (default "genkit-sessions") + .createTableIfNotExists(true) // auto-create the table on first use (default false) + .pollIntervalMs(500) // how often change subscribers are polled (default 2000) + .build()); +``` + +The default shard size is 350 KiB, kept under DynamoDB's 400 KB item-size limit. `DynamoDbSessionStore` supports `onSnapshotStateChange` (via polling), so `chat.abort()` works. + +## CosmosSessionStore + +Stores snapshots in Azure Cosmos DB. It lives in the Azure AI Foundry plugin, so add the dependency: + +```xml + + com.google.genkit + genkit-plugin-azure-foundry + ${genkit.version} + +``` + +```java +import com.google.genkit.plugins.azurefoundry.session.CosmosSessionStore; +import com.azure.cosmos.CosmosClient; +import com.azure.cosmos.CosmosClientBuilder; + +CosmosClient client = new CosmosClientBuilder() + .endpoint(System.getenv("COSMOS_ENDPOINT")) + .key(System.getenv("COSMOS_KEY")) + .buildClient(); + +CosmosSessionStore> store = + new CosmosSessionStore<>(client); // uses database "genkit", container "genkit-sessions" + +Agent> agent = genkit.beta().defineAgent( + AgentConfig.>builder() + .name("myAgent") + .system("You are a helpful assistant.") + .store(store) + .build()); +``` + +All records live in a single container partitioned by `/pk`. Create the database and container ahead of time, or let the store create them on first use: + +```java +CosmosSessionStore> store = + new CosmosSessionStore<>( + client, + CosmosSessionStoreOptions.builder() + .databaseName("genkit") // database name (default "genkit") + .containerName("genkit-sessions") // container name (default "genkit-sessions") + .createIfNotExists(true) // auto-create database + container (default false) + .pollIntervalMs(500) // how often change subscribers are polled (default 2000) + .build()); +``` + +The default shard size is 1 MiB, kept under Cosmos DB's 2 MB document-size limit. `CosmosSessionStore` supports `onSnapshotStateChange` (via polling), so `chat.abort()` works. + +## Writing your own store + +Implement `SessionStore` to use any backend — a database, a cache, a cloud object store: + +```java +import com.google.genkit.ai.agent.SessionStore; +import com.google.genkit.ai.agent.SessionSnapshot; +import com.google.genkit.ai.agent.GetSnapshotOptions; +import com.google.genkit.ai.agent.SessionStoreOptions; +import com.google.genkit.ai.agent.SnapshotMutator; + +public class RedisSessionStore implements SessionStore { + + @Override + public SessionSnapshot getSnapshot(GetSnapshotOptions opts) { + // Read by snapshotId or sessionId. + ... + } + + @Override + public String saveSnapshot(String snapshotId, SnapshotMutator mutator, + SessionStoreOptions options) { + // Apply the mutator to the existing snapshot, persist the result, + // and return the final snapshotId. + ... + } +} +``` + +## See also + +- [Agents Overview](../overview) — Defining agents and the chat API +- [Sessions](../sessions) — Session lifecycle, snapshots, and resuming +- [Serve over HTTP](../serve-over-http) — Exposing a store-backed agent through Jetty or Spring +- [Background Execution](../background-execution) — Detached turns and the store's role in polling diff --git a/docs/src/content/docs/agents/sessions.md b/docs/src/content/docs/agents/sessions.md new file mode 100644 index 000000000..0044e3e22 --- /dev/null +++ b/docs/src/content/docs/agents/sessions.md @@ -0,0 +1,114 @@ +--- +title: Sessions +description: Understand the agent session model, snapshots, and how to resume conversations. +--- + +Every agent conversation is backed by a **session** — a named container for conversation history, custom application state, and (optionally) artifacts. Each completed turn produces a **snapshot**: an immutable record of the session at that point. Snapshots are what let you resume a conversation later, or branch it in a new direction. + +## Session and snapshot IDs + +| Identifier | Scope | Purpose | +|------------|-------|---------| +| `sessionId` | The whole conversation | Groups all snapshots in one thread | +| `snapshotId` | A single turn | Identifies one turn's state, so you can resume from an exact point | + +A new session is created on the first `send` when you don't supply either ID. The server assigns both, and you can read them back from `AgentResponse.sessionId()` and `AgentResponse.snapshotId()`. + +## What happens on a turn + +For each `chat.send(text)` call: + +1. `AgentChat` adds the user message to the session history. +2. The agent's turn logic runs with that history. For a model-backed agent this is one model call with your system prompt, tools, and history. +3. The model may call tools (up to `maxTurns`). +4. The reply is written back as the assistant message. +5. If the agent has a `SessionStore`, the updated state is saved as a new snapshot and the `snapshotId` advances. +6. `AgentChat` updates its tracked state and returns an `AgentResponse`. + +## Resuming a session + +### By snapshot ID + +Resume from an exact turn: + +```java +import com.google.genkit.ai.agent.AgentInit; +import com.google.genkit.ai.agent.AgentChat; + +AgentChat> chat = weatherAgent.chat(ctx, + AgentInit.>builder() + .snapshotId("f3a7-…") + .build()); + +chat.send("Continue from where we left off"); +``` + +### By session ID + +Resume from a session's latest snapshot: + +```java +AgentChat> chat = weatherAgent.chat(ctx, + AgentInit.>builder() + .sessionId("session-abc-123") + .build()); +``` + +### With `loadChat` + +`loadChat` reads the snapshot from the store and hydrates the chat in one call: + +```java +import com.google.genkit.ai.agent.GetSnapshotRequest; + +AgentChat> chat = weatherAgent.loadChat(ctx, + GetSnapshotRequest.builder() + .sessionId("session-abc-123") + .build()); + +chat.send("Next turn"); +``` + +## Client-managed sessions + +When an agent has no `SessionStore`, `AgentChat` carries the full session state itself and passes the accumulated messages and custom state along on every `send`. This avoids any server-side persistence, at the cost of larger requests. + +```java +AgentChat> chat = statelessAgent.chat(ctx); +chat.send("Hello"); // fresh session +chat.send("And in French?"); // history round-tripped automatically +``` + +## Snapshots chain and branch + +Each turn's snapshot descends from the one before it, forming a chain that is your conversation history. Because a snapshot is immutable, you can **branch** by resuming from an earlier one — the new turn starts a fresh chain from that point while the original chain stays intact. + +```java +// Resume from an earlier snapshot to take a different path +AgentChat> branch = weatherAgent.chat(ctx, + AgentInit.>builder() + .snapshotId("snap-turn-2") + .build()); + +branch.send("Take a different path"); +``` + +## Aborting a turn + +`AgentChat.abort()` marks the chat's latest snapshot as `ABORTED` in the session store: + +```java +SnapshotStatus status = chat.abort(); +System.out.println("Status after abort: " + status); +``` + +Abort applies to **background (detached) turns** — a detached turn that cooperatively checks `ctx.isAborted()` can stop early, and once a turn is aborted it stays aborted even if the background work later finishes. A synchronous `chat.send(...)` turn runs to completion regardless, so `abort()` on a foreground turn only records that the caller gave up rather than stopping work in progress. Abort needs a store that supports change notifications — `FileSessionStore` and the Firestore store do; `InMemorySessionStore` does not. See [Background Execution](../background-execution) for detached-turn mechanics. + +## See also + +- [Agents Overview](../overview) — Defining agents and the chat API +- [Define Agents](../define-agents) — `AgentFn` and `AgentFnContext` +- [Run and Stream](../run-and-stream) — `send`/`sendStream` and resuming after an interrupt +- [Session Stores](../session-stores) — Persistence backends +- [Background Execution](../background-execution) — Detached turns and their lifecycle +- [Error Handling](../error-handling) — Failure modes and how they surface diff --git a/docs/src/content/docs/chat-sessions.md b/docs/src/content/docs/chat-sessions.md deleted file mode 100644 index c4942e614..000000000 --- a/docs/src/content/docs/chat-sessions.md +++ /dev/null @@ -1,244 +0,0 @@ ---- -title: Chat Sessions -description: Build multi-turn chat experiences with session persistence. ---- - -Genkit supports multi-turn chat sessions with automatic history management, typed session state, multiple conversation threads, and pluggable persistence via `SessionStore`. - -## Creating a session - -```java -Session session = genkit.createSession(); -``` - -### With session options - -```java -import com.google.genkit.ai.session.*; - -Session session = genkit.createSession( - SessionOptions.builder() - .store(sessionStore) - .initialState(new ConversationState("John")) - .sessionId("custom-session-id") // optional, auto-generated if omitted - .build()); -``` - -## Starting a chat - -Create a `Chat` from a session. The chat manages message history, tool execution, and streaming automatically: - -```java -Chat chat = session.chat( - ChatOptions.builder() - .model("openai/gpt-4o-mini") - .system("You are a helpful assistant.") - .tools(List.of(noteTool, searchTool)) - .config(GenerationConfig.builder().temperature(0.7).build()) - .build()); -``` - -## Sending messages - -Each `send()` call automatically includes the full conversation history: - -```java -ModelResponse r1 = chat.send("What is the capital of France?"); -System.out.println(r1.getText()); // "Paris" - -ModelResponse r2 = chat.send("And what about Germany?"); -System.out.println(r2.getText()); // "Berlin" — knows we're asking about capitals -``` - -### With per-message options - -Override model or tools for a specific message: - -```java -ModelResponse response = chat.send("Summarize everything", - Chat.SendOptions.builder() - .model("openai/gpt-4o") // use a different model for this turn - .maxTurns(5) - .build()); -``` - -### Streaming - -```java -ModelResponse response = chat.sendStream("Tell me a long story", - (chunk) -> System.out.print(chunk.getText())); -``` - -## Session state - -Sessions can hold typed state that persists across messages and is saved to the store: - -```java -public class ConversationState { - private String userName; - private int messageCount; - private List topics; - // constructors, getters, setters... -} - -// Read state -ConversationState state = session.getState(); -System.out.println("User: " + state.getUserName()); - -// Update state -state.incrementMessageCount(); -state.getTopics().add("geography"); -session.updateState(state).join(); -``` - -## Conversation threads - -A single session can have multiple independent conversation threads, each with its own history: - -```java -Chat generalChat = session.chat("general", - ChatOptions.builder() - .model("openai/gpt-4o-mini") - .system("You are a general assistant.") - .build()); - -Chat supportChat = session.chat("support", - ChatOptions.builder() - .model("openai/gpt-4o-mini") - .system("You are a support agent.") - .build()); - -// Each thread maintains separate history -generalChat.send("What's the weather like?"); -supportChat.send("I have a billing issue"); -``` - -### Accessing history - -```java -List history = chat.getHistory(); -List supportHistory = session.getMessages("support"); -List mainHistory = session.getMessages(); // default "main" thread -``` - -## Loading existing sessions - -Resume a previous session by loading it from the store: - -```java -Session loaded = genkit.loadSession( - "session-123", - SessionOptions.builder() - .store(sessionStore) - .build() -).get(); // returns CompletableFuture - -// Continue the conversation with full history -Chat chat = loaded.chat(); -chat.send("Where were we?"); -``` - -## SessionStore - -The `SessionStore` interface defines how sessions are persisted. Genkit includes an in-memory implementation, and you can create custom stores for any backend. - -### Interface - -```java -public interface SessionStore { - CompletableFuture> get(String sessionId); - CompletableFuture save(String sessionId, SessionData sessionData); - CompletableFuture delete(String sessionId); -} -``` - -### InMemorySessionStore - -The default implementation stores sessions in a `ConcurrentHashMap`. Useful for development and testing — data is lost when the process stops: - -```java -InMemorySessionStore store = new InMemorySessionStore<>(); - -Session session = genkit.createSession( - SessionOptions.builder() - .store(store) - .initialState(new ConversationState("John")) - .build()); -``` - -### Creating a custom SessionStore - -Implement the `SessionStore` interface to persist sessions to any backend — Redis, a database, Firebase, etc. - -```java -import com.google.genkit.ai.session.SessionStore; -import com.google.genkit.ai.session.SessionData; - -public class RedisSessionStore implements SessionStore { - private final RedisClient redis; - private final ObjectMapper mapper; - - public RedisSessionStore(RedisClient redis, ObjectMapper mapper) { - this.redis = redis; - this.mapper = mapper; - } - - @Override - public CompletableFuture> get(String sessionId) { - return CompletableFuture.supplyAsync(() -> { - String json = redis.get("session:" + sessionId); - if (json == null) return null; - return mapper.readValue(json, SessionData.class); - }); - } - - @Override - public CompletableFuture save(String sessionId, SessionData data) { - return CompletableFuture.runAsync(() -> { - String json = mapper.writeValueAsString(data); - redis.set("session:" + sessionId, json); - }); - } - - @Override - public CompletableFuture delete(String sessionId) { - return CompletableFuture.runAsync(() -> { - redis.del("session:" + sessionId); - }); - } -} -``` - -Then use it like any other store: - -```java -RedisSessionStore store = - new RedisSessionStore<>(redisClient, objectMapper); - -Session session = genkit.createSession( - SessionOptions.builder() - .store(store) - .initialState(new ConversationState("John")) - .build()); -``` - -### SessionData structure - -The `SessionData` object is what gets persisted. It contains: - -| Field | Type | Description | -|-------|------|-------------| -| `id` | `String` | Session ID | -| `state` | `S` | Your typed session state | -| `threads` | `Map>` | All conversation threads and their message histories | - -```java -// Access thread data directly -SessionData data = store.get("session-123").get(); -List mainThread = data.getThread("main"); -Map> allThreads = data.getThreads(); -``` - -## Sample - -See the [chat-session sample](https://github.com/genkit-ai/genkit-java/tree/main/samples/chat-session) for a complete multi-turn chat implementation with state, tools, and session persistence. diff --git a/docs/src/content/docs/multi-agent.md b/docs/src/content/docs/multi-agent.md deleted file mode 100644 index e137d6dc7..000000000 --- a/docs/src/content/docs/multi-agent.md +++ /dev/null @@ -1,155 +0,0 @@ ---- -title: Multi-Agent -description: Build multi-agent AI orchestration patterns. ---- - -Multi-agent systems coordinate multiple AI agents, each with specialized capabilities, to solve complex tasks. In Genkit, agents are defined with their own system prompt, model, and tools. A triage agent routes user requests to the right specialist by calling that agent as a tool, triggering an automatic handoff. - -## Defining agents - -Use `genkit.defineAgent()` to create specialized agents. Each agent has its own system prompt, model, tools, and optional generation config: - -```java -import com.google.genkit.ai.Agent; -import com.google.genkit.ai.AgentConfig; - -// Define tools for the reservation agent -Tool makeReservationTool = genkit.defineTool( - "makeReservation", "Makes a restaurant reservation", - ReservationInput.class, String.class, - (ctx, input) -> "Reserved table for " + input.getPartySize() - + " at " + input.getTime()); - -Tool cancelReservationTool = genkit.defineTool( - "cancelReservation", "Cancels a reservation", - String.class, String.class, - (ctx, id) -> "Reservation " + id + " cancelled"); - -// Create a specialized agent -Agent reservationAgent = genkit.defineAgent( - AgentConfig.builder() - .name("reservationAgent") - .description("Handles restaurant reservations. " - + "Transfer here when the customer wants to book, " - + "modify, or cancel a reservation.") - .system("You are a reservation specialist. Help customers " - + "book and manage restaurant reservations.") - .model("openai/gpt-4o-mini") - .tools(List.of(makeReservationTool, cancelReservationTool)) - .config(GenerationConfig.builder().temperature(0.3).build()) - .build()); -``` - -## Creating the triage agent - -The triage (orchestrator) agent has sub-agents listed in its config. Genkit automatically registers each sub-agent as a tool that the triage agent can call: - -```java -Agent menuAgent = genkit.defineAgent( - AgentConfig.builder() - .name("menuAgent") - .description("Provides menu information and recommendations.") - .system("You are a menu expert. Help customers explore " - + "the menu and make recommendations.") - .model("openai/gpt-4o-mini") - .tools(List.of(getMenuTool, searchMenuTool)) - .build()); - -// Triage agent routes to specialists -Agent triageAgent = genkit.defineAgent( - AgentConfig.builder() - .name("triageAgent") - .description("Routes customer requests to the right specialist") - .system("You are the main customer service agent. Route requests:\n" - + "- reservationAgent: for booking/modifying/cancelling reservations\n" - + "- menuAgent: for menu questions and recommendations\n" - + "Transfer to the appropriate specialist using their tools.") - .model("openai/gpt-4o-mini") - .agents(List.of( - reservationAgent.getConfig(), - menuAgent.getConfig())) - .build()); -``` - -## Running a multi-agent chat - -Use `genkit.getAllToolsForAgent()` to get the triage agent's own tools plus all sub-agent tools, then start a chat session: - -```java -List> allTools = genkit.getAllToolsForAgent(triageAgent); - -Session session = genkit.createSession(); -Chat chat = session.chat( - ChatOptions.builder() - .model(triageAgent.getModel()) - .system(triageAgent.getSystem()) - .tools(allTools) - .build()); - -// User asks about the menu → routes to menuAgent -ModelResponse r1 = chat.send("What appetizers do you have?"); - -// User wants to book → routes to reservationAgent -ModelResponse r2 = chat.send("Great, book a table for 4 at 7pm"); -``` - -## How agent handoff works - -When the triage agent decides to delegate, it calls a sub-agent's tool. This triggers an `AgentHandoffException` internally, which Genkit handles automatically: - -1. The triage agent calls the sub-agent tool (e.g., `reservationAgent`) -2. Genkit catches the handoff and switches the chat's active system prompt, model, and tools to the target agent -3. Subsequent messages are handled by the new agent until another handoff occurs - -You can check which agent is currently active: - -```java -String currentAgent = chat.getCurrentAgentName(); -System.out.println("Now talking to: " + currentAgent); -``` - -## Stateful multi-agent conversations - -Combine agents with session state to track context across handoffs: - -```java -Session session = genkit.createSession( - SessionOptions.builder() - .store(sessionStore) - .initialState(new CustomerState("John")) - .build()); - -Chat chat = session.chat( - ChatOptions.builder() - .model(triageAgent.getModel()) - .system(triageAgent.getSystem()) - .tools(allTools) - .build()); - -// State persists across agent handoffs -CustomerState state = session.getState(); -``` - -## Agent registry - -All defined agents are registered in a global agent registry. You can look up agents by name: - -```java -Agent agent = genkit.getAgent("reservationAgent"); -Map allAgents = genkit.getAgentRegistry(); -``` - -## Patterns - -### Orchestrator pattern -A central triage agent delegates to specialized agents based on the user's intent. Best for customer service, help desks, and general-purpose assistants. - -### Pipeline pattern -Agents process data sequentially, each contributing to the result. Define each stage as an agent and chain them in your flow logic. - -### Debate pattern -Multiple agents propose different solutions, then a judge agent selects the best one. Useful for creative tasks or complex decision-making. - -## Sample - -See the [multi-agent sample](https://github.com/genkit-ai/genkit-java/tree/main/samples/multi-agent) for a complete restaurant assistant with triage, reservation, and menu agents. diff --git a/docs/src/content/docs/plugins/aws-bedrock.md b/docs/src/content/docs/plugins/aws-bedrock.md index a3414a19b..b13d93b61 100644 --- a/docs/src/content/docs/plugins/aws-bedrock.md +++ b/docs/src/content/docs/plugins/aws-bedrock.md @@ -57,6 +57,20 @@ ModelResponse response = genkit.generate( - INFERENCE_PROFILE support for advanced models - Text generation, streaming, tool calling +## Session store + +The plugin also ships `DynamoDbSessionStore`, a DynamoDB-backed agent session store. Construct it and pass it to an agent's `.store(...)` to persist server-managed sessions in DynamoDB: + +```java +import com.google.genkit.plugins.awsbedrock.session.DynamoDbSessionStore; +import software.amazon.awssdk.services.dynamodb.DynamoDbClient; + +DynamoDbSessionStore> store = + new DynamoDbSessionStore<>(DynamoDbClient.create()); +``` + +See [Session Stores](../../agents/session-stores#dynamodbsessionstore) for options and the agents-dynamodb-session sample. + ## Sample See the [aws-bedrock sample](https://github.com/genkit-ai/genkit-java/tree/main/samples/aws-bedrock). diff --git a/docs/src/content/docs/plugins/azure-foundry.md b/docs/src/content/docs/plugins/azure-foundry.md index 87d1a7471..f10fcd35f 100644 --- a/docs/src/content/docs/plugins/azure-foundry.md +++ b/docs/src/content/docs/plugins/azure-foundry.md @@ -47,6 +47,23 @@ ModelResponse response = genkit.generate( - Claude 4.x via Azure marketplace - And many more +## Session store + +The plugin also ships `CosmosSessionStore`, an Azure Cosmos DB-backed agent session store. Construct it and pass it to an agent's `.store(...)` to persist server-managed sessions in Cosmos DB: + +```java +import com.google.genkit.plugins.azurefoundry.session.CosmosSessionStore; +import com.azure.cosmos.CosmosClientBuilder; + +CosmosSessionStore> store = new CosmosSessionStore<>( + new CosmosClientBuilder() + .endpoint(System.getenv("COSMOS_ENDPOINT")) + .key(System.getenv("COSMOS_KEY")) + .buildClient()); +``` + +See [Session Stores](../../agents/session-stores#cosmossessionstore) for options and the agents-cosmos-session sample. + ## Sample See the [azure-foundry sample](https://github.com/genkit-ai/genkit-java/tree/main/samples/azure-foundry). diff --git a/docs/src/content/docs/plugins/firebase.md b/docs/src/content/docs/plugins/firebase.md index 55c8df8a3..6d004122d 100644 --- a/docs/src/content/docs/plugins/firebase.md +++ b/docs/src/content/docs/plugins/firebase.md @@ -3,10 +3,11 @@ title: Firebase description: Firebase integration with Firestore vector search, Cloud Functions deployment, and Google Cloud telemetry. --- -The Firebase plugin provides three key capabilities: +The Firebase plugin provides these key capabilities: - **[Firestore Vector Store](/genkit-java/plugins/firebase-vector-store/)** — Native vector similarity search for RAG applications. - **[Cloud Functions Deployment](/genkit-java/plugins/firebase-functions/)** — Deploy Genkit flows as scalable Cloud Functions with auth and streaming. +- **Firestore Session Store** — Persist server-managed agent sessions in Google Cloud Firestore. - **Telemetry** — Export traces and metrics to Google Cloud observability. ## Installation @@ -38,6 +39,47 @@ Enable these APIs in your Google Cloud project: - Cloud Trace API - Cloud Monitoring API +## Firestore Session Store + +`FirestoreSessionStore` persists server-managed agent sessions in Google Cloud Firestore, making it a good fit for production deployments on Cloud Run or Firebase Functions. Pass it to `.store(...)` when defining an agent: + +```java +import com.google.genkit.plugins.firebase.session.FirestoreSessionStore; +import com.google.cloud.firestore.Firestore; +import com.google.firebase.cloud.FirestoreClient; + +Firestore db = FirestoreClient.getFirestore(); + +FirestoreSessionStore> store = new FirestoreSessionStore<>(db); + +Agent> agent = genkit.beta().defineAgent( + AgentConfig.>builder() + .name("myAgent") + .system("You are a helpful assistant.") + .store(store) + .build()); +``` + +### Options + +Use `FirestoreSessionStoreOptions` to customize the collection name and checkpointing behavior: + +```java +import com.google.genkit.plugins.firebase.session.FirestoreSessionStoreOptions; + +FirestoreSessionStore> store = new FirestoreSessionStore<>( + db, + FirestoreSessionStoreOptions.builder() + .collection("agent-sessions") // top-level collection (default "genkit-sessions") + .checkpointInterval(25) // turns between full checkpoints (default 25) + .shardSize(512 * 1024) // checkpoint shard size in bytes (default 512 KiB) + .build()); +``` + +The store derives three collections from the configured name: `` for snapshot metadata, `-shards` for checkpoint state, and `-pointers` for the current leaf snapshot per session. + +See [Session Stores](/genkit-java/agents/session-stores/) for the full comparison of available stores. + ## Requirements - Firebase project on the Blaze (pay-as-you-go) plan diff --git a/docs/src/content/docs/plugins/jetty.md b/docs/src/content/docs/plugins/jetty.md index 681e42247..9f71629b9 100644 --- a/docs/src/content/docs/plugins/jetty.md +++ b/docs/src/content/docs/plugins/jetty.md @@ -46,3 +46,17 @@ curl -X POST http://localhost:8080/api/flows/tellJoke \ -H "Content-Type: application/json" \ -d '{"data": "pirates"}' ``` + +## Serving agents + +Any agent defined with `genkit.beta().defineAgent(...)` / `defineCustomAgent(...)` is mounted automatically — just define it before calling `jetty.start()`. Agents are served at the **root path** (`/`), separate from flows' `/api/flows/...` base: + +```bash +curl -X POST http://localhost:8080/echoAgent \ + -H "Content-Type: application/json" \ + -d '{"data":{"message":{"role":"user","content":[{"text":"Hello, agent!"}]}}}' +``` + +For server-managed agents (those configured with a `SessionStore`), the companion endpoints `POST //getSnapshot` and `POST //abort` are mounted too. Streaming turns are available via `Accept: text/event-stream` or `?stream=true`. + +The same Java client (`RemoteAgent` / `AgentChat`) works against a Jetty or Spring server unchanged. See [Serve over HTTP](../../agents/serve-over-http) for the full request/response contract, calling a served agent from Java, and how request headers reach your agent — and [Background Execution](../../agents/background-execution) for detached (background) turns. diff --git a/docs/src/content/docs/plugins/spring.md b/docs/src/content/docs/plugins/spring.md index 740ebccb6..8c1e5413f 100644 --- a/docs/src/content/docs/plugins/spring.md +++ b/docs/src/content/docs/plugins/spring.md @@ -42,6 +42,20 @@ You must call `spring.start()` after building the `Genkit` instance and defining | `GET /api/flows` | List all registered flows | | `POST /api/flows/{flowName}` | Execute a flow | +## Serving agents + +Any agent defined with `genkit.beta().defineAgent(...)` / `defineCustomAgent(...)` is mounted automatically once the server starts — just define it before calling `spring.start()`. Agents are served at the **root path** (`/`), separate from flows' `/api/flows/...` base: + +```bash +curl -X POST http://localhost:8080/echoAgent \ + -H "Content-Type: application/json" \ + -d '{"data":{"message":{"role":"user","content":[{"text":"Hello, agent!"}]}}}' +``` + +For server-managed agents (those configured with a `SessionStore`), the companion endpoints `POST //getSnapshot` and `POST //abort` are mounted too. Streaming turns are available via `Accept: text/event-stream` or `?stream=true`. + +The same Java client (`RemoteAgent` / `AgentChat`) works against a Spring or Jetty server unchanged. See [Serve over HTTP](../../agents/serve-over-http) for the full request/response contract, calling a served agent from Java, and how request headers reach your agent — and [Background Execution](../../agents/background-execution) for detached (background) turns. + ## Configuration | Option | Default | Description | diff --git a/docs/src/content/docs/samples.md b/docs/src/content/docs/samples.md index e17c86dce..1a7e6be9d 100644 --- a/docs/src/content/docs/samples.md +++ b/docs/src/content/docs/samples.md @@ -42,15 +42,14 @@ genkit start -- ./run.sh | Sample | Description | |--------|-------------| +| **agents-weather** | Weather assistant agent demonstrating the beta Agents API (server-managed and client-managed sessions) | | **dotprompt** | DotPrompt files with complex inputs/outputs, variants, and partials | | **structured-output** | Type-safe structured output generation | | **rag** | RAG application with local vector store | -| **chat-session** | Multi-turn chat with session persistence | | **evaluations** | Custom evaluators and evaluation workflows | | **evaluators-plugin** | Pre-built RAGAS-style evaluators plugin demo | | **complex-io** | Complex nested types, arrays, maps in flow inputs/outputs | | **middleware** | Middleware patterns for logging, caching, rate limiting | -| **multi-agent** | Multi-agent orchestration patterns | | **interrupts** | Flow interrupts and human-in-the-loop patterns | | **mcp** | Model Context Protocol (MCP) integration | diff --git a/genkit/pom.xml b/genkit/pom.xml index 897f17f43..9a3849bad 100644 --- a/genkit/pom.xml +++ b/genkit/pom.xml @@ -98,5 +98,11 @@ mockito-core test + + + com.fasterxml.jackson.dataformat + jackson-dataformat-yaml + test + diff --git a/genkit/src/main/java/com/google/genkit/Genkit.java b/genkit/src/main/java/com/google/genkit/Genkit.java index 67314e07e..b1eec3c6c 100644 --- a/genkit/src/main/java/com/google/genkit/Genkit.java +++ b/genkit/src/main/java/com/google/genkit/Genkit.java @@ -21,7 +21,6 @@ import com.google.genkit.ai.*; import com.google.genkit.ai.evaluation.*; import com.google.genkit.ai.middleware.*; -import com.google.genkit.ai.session.*; import com.google.genkit.ai.telemetry.ModelTelemetryHelper; import com.google.genkit.core.*; import com.google.genkit.core.middleware.Middleware; @@ -32,7 +31,6 @@ import java.util.ArrayList; import java.util.List; import java.util.Map; -import java.util.concurrent.CompletableFuture; import java.util.concurrent.ConcurrentHashMap; import java.util.function.BiFunction; import java.util.function.Function; @@ -53,7 +51,6 @@ public class Genkit { private final List plugins; private final GenkitOptions options; private final Map> promptCache; - private final Map agentRegistry; private ReflectionServer reflectionServer; private ReflectionServerV2 reflectionServerV2; private EvaluationManager evaluationManager; @@ -73,7 +70,6 @@ public Genkit(GenkitOptions options) { this.registry = new DefaultRegistry(); this.plugins = new ArrayList<>(); this.promptCache = new ConcurrentHashMap<>(); - this.agentRegistry = new ConcurrentHashMap<>(); } /** @@ -674,7 +670,12 @@ private ModelResponse generateInternal(GenerateOptions options) throws Genkit private ModelResponse generateInternal( GenerateOptions options, java.util.function.Consumer streamCallback) throws GenkitException { - ActionContext ctx = new ActionContext(registry); + // Thread the caller-supplied user context (e.g. {"auth": {...}}) into the ActionContext used + // to execute tools during this generate. Tools therefore observe ctx.getContext() == + // options.getContext(). This is additional to (and independent of) the model-grounding use of + // options.getContext() at generateObject. + ActionContext ctx = + ActionContext.builder().registry(registry).context(options.getContext()).build(); int maxTurns = options.getMaxTurns() != null ? options.getMaxTurns() : 5; @@ -713,6 +714,13 @@ private ModelResponse generateInternal( // Use an array to hold the reference for recursive WrapGenerate wrapping final GenerateNext[] generateRef = new GenerateNext[1]; + // Streaming role/message-index tracking (mirrors Go's wrappedCb in generate.go:337-357 and JS + // makeChunk in action.ts:303-334). Model chunks default to role=model; the message index bumps + // whenever the streamed role transitions (e.g. model → tool). These are shared across the whole + // tool loop so indices stay monotonic across turns. + final Role[] streamCurrentRole = {Role.MODEL}; + final int[] streamCurrentIndex = {0}; + // Core generate iteration: resolve options → model call → tool handling → recurse GenerateNext rawGenerate = (actx, params) -> { @@ -736,11 +744,19 @@ private ModelResponse generateInternal( List restarts = new java.util.ArrayList<>(pendingRestarts); pendingRestarts.clear(); - // Convert restart requests to tool request parts for middleware execution - List restartParts = - restarts.stream() - .map(Part::toolRequest) - .collect(java.util.stream.Collectors.toList()); + // Convert restart requests to tool request parts for middleware execution, PRESERVING + // the request's Part-level metadata (resumed / replacedInput) so the restarted tool + // can observe its resumed status via ActionContext.isResumed()/getResumed() (mirrors + // Go handleResumedToolRequest → ToolContext.Resumed). The restart directive's metadata + // travels on the ToolRequest itself (set by GenkitBeta.toResumeOptions). + List restartParts = new java.util.ArrayList<>(); + for (ToolRequest restart : restarts) { + Part restartPart = Part.toolRequest(restart); + if (restart.getMetadata() != null && !restart.getMetadata().isEmpty()) { + restartPart.setMetadata(new java.util.HashMap<>(restart.getMetadata())); + } + restartParts.add(restartPart); + } // Execute through WrapTool chain (fires wrapTool hooks) ToolExecutionResult toolResult = @@ -795,6 +811,16 @@ private ModelResponse generateInternal( // Recurse through WrapGenerate hooks for the next turn (propagate onChunk) GenerateActionOptions nextOpts = opts.withMessages(updatedMessages); int nextMsgIdx = params.getMessageIndex() + 1; + + // Stream the resumed tool-response message as a chunk before recursing (mirrors Go + // generate.go:232-241 and JS action.ts:324-327). + emitToolResponseChunk( + params.getOnChunk(), + toolResult.getResponses(), + nextMsgIdx, + streamCurrentRole, + streamCurrentIndex); + return generateRef[0].apply( actx, new GenerateParams(nextOpts, turn + 1, nextMsgIdx, params.getOnChunk())); } @@ -804,8 +830,29 @@ private ModelResponse generateInternal( ModelNext wrappedModelCall = buildWrappedModelCall(model, opts.getModel(), actx, middlewares); - // Call model through WrapModel chain (propagate streaming callback from GenerateParams) - ModelParams mparams = new ModelParams(req, params.getOnChunk()); + // Call model through WrapModel chain (propagate streaming callback from GenerateParams). + // Wrap the callback so model chunks carry role=model and a monotonic message index, + // bumping the index on any role transition (mirrors Go generate.go:337-357). + final java.util.function.Consumer rawOnChunk = params.getOnChunk(); + java.util.function.Consumer wrappedOnChunk = null; + if (rawOnChunk != null) { + streamCurrentIndex[0] = params.getMessageIndex(); + streamCurrentRole[0] = Role.MODEL; + wrappedOnChunk = + chunk -> { + Role chunkRole = chunk.getRole(); + if (chunkRole != null && chunkRole != streamCurrentRole[0]) { + streamCurrentIndex[0]++; + streamCurrentRole[0] = chunkRole; + } + chunk.setIndex(streamCurrentIndex[0]); + if (chunk.getRole() == null) { + chunk.setRole(Role.MODEL); + } + rawOnChunk.accept(chunk); + }; + } + ModelParams mparams = new ModelParams(req, wrappedOnChunk); ModelResponse response = wrappedModelCall.apply(actx, mparams); // Check if the model requested tool calls @@ -840,6 +887,17 @@ private ModelResponse generateInternal( // Recurse through the wrapped generate function (goes through WrapGenerate hooks) int nextMsgIdx = params.getMessageIndex() + 1; + + // Stream the tool-response message as a chunk before recursing (mirrors Go + // generate.go:865-874 and JS action.ts:420-425). The chunk carries role=tool, the tool + // response parts, and the incremented message index. + emitToolResponseChunk( + params.getOnChunk(), + toolResult.getResponses(), + nextMsgIdx, + streamCurrentRole, + streamCurrentIndex); + return generateRef[0].apply( actx, new GenerateParams(nextOpts, turn + 1, nextMsgIdx, params.getOnChunk())); }; @@ -851,6 +909,30 @@ private ModelResponse generateInternal( return generateRef[0].apply(ctx, new GenerateParams(actionOpts, 0, 0, streamCallback)); } + /** + * Emits the tool-response message as a streaming chunk before the generate loop recurses into the + * next model turn (mirrors Go {@code generate.go} and JS {@code action.ts}: the tool message is + * streamed as a {@code role: tool} chunk at the incremented message index). No-op when {@code + * onChunk} is null or there are no responses. Advances the shared streaming role/index trackers + * so a subsequent model chunk bumps the index correctly. + */ + private static void emitToolResponseChunk( + java.util.function.Consumer onChunk, + List toolResponseParts, + int messageIndex, + Role[] streamCurrentRole, + int[] streamCurrentIndex) { + if (onChunk == null || toolResponseParts == null || toolResponseParts.isEmpty()) { + return; + } + ModelResponseChunk chunk = new ModelResponseChunk(toolResponseParts); + chunk.setRole(Role.TOOL); + chunk.setIndex(messageIndex); + streamCurrentRole[0] = Role.TOOL; + streamCurrentIndex[0] = messageIndex; + onChunk.accept(chunk); + } + /** Creates fresh middleware instances for a single generate invocation. */ private List createMiddlewareInstances(List use) { if (use == null || use.isEmpty()) { @@ -1117,9 +1199,24 @@ private ToolExecutionResult executeToolsWithMiddleware( } try { - // Execute through WrapTool chain + // Execute through WrapTool chain. When this tool request is a RESTART (its Part carries + // `resumed` metadata attached by Tool.restart(...) / GenkitBeta.toResumeOptions), thread + // the + // resumed value and original input into the ActionContext so a restart-aware tool handler + // can observe ctx.isResumed()/getResumed()/getOriginalInput() (mirrors Go's + // handleResumedToolRequest → ToolContext.Resumed and JS ToolRunOptions.resumed). + ActionContext toolCtx = ctx; + Map partMeta = toolRequestPart.getMetadata(); + if (partMeta != null && partMeta.containsKey("resumed")) { + Object resumedValue = partMeta.get("resumed"); + Object originalInput = + partMeta.containsKey("replacedInput") + ? partMeta.get("replacedInput") + : toolRequest.getInput(); + toolCtx = ctx.withResumed(resumedValue, originalInput); + } ToolParams tparams = new ToolParams(toolRequestPart, tool); - Part responsePart = wrappedToolCall.apply(ctx, tparams); + Part responsePart = wrappedToolCall.apply(toolCtx, tparams); responseParts.add(responsePart); @@ -1906,6 +2003,40 @@ public GenkitOptions getOptions() { return options; } + /** + * Returns the beta (experimental) API surface for this Genkit instance. + * + *

The beta API exposes experimental features such as the {@code defineAgent}, {@code + * definePromptAgent}, and {@code defineCustomAgent} bidi agent actions. These methods are gated + * behind the {@code experimental} flag (see {@link GenkitOptions.Builder#experimental} or the + * {@code GENKIT_EXPERIMENTAL} environment variable) and throw a {@link GenkitException} if it is + * not enabled. + * + * 

{@code
+   * Genkit ai = new Genkit(GenkitOptions.builder().experimental(true).build());
+   * Agent agent = ai.beta().defineAgent(
+   *     AgentConfig.builder()
+   *         .name("helper")
+   *         .system("You are helpful.")
+   *         .model("googleai/gemini-2.0-flash")
+   *         .build());
+   * }
+ * + * @return the beta API surface + */ + public GenkitBeta beta() { + return new GenkitBeta(this); + } + + /** + * Returns whether experimental features are enabled for this instance. + * + * @return true if experimental features are enabled + */ + public boolean isExperimental() { + return options.isExperimental(); + } + /** * Gets the registered plugins. * @@ -1972,201 +2103,9 @@ public void stop() { } // ========================================================================= - // Session Methods - // ========================================================================= - - /** - * Creates a new session with default options. - * - *

Sessions provide stateful multi-turn conversations with automatic history persistence. Each - * session can have multiple named conversation threads. - * - *

Example usage: - * - * 

{@code
-   * Session session = genkit.createSession();
-   * Chat chat = session.chat(
-   *     ChatOptions.builder()
-   *         .model("openai/gpt-4o")
-   *         .system("You are a helpful assistant.")
-   *         .build());
-   * chat.send("Hello!");
-   * }
- * - * @param the session state type - * @return a new session - */ - public Session createSession() { - return Session.create(registry, SessionOptions.builder().build(), agentRegistry); - } - - /** - * Creates a new session with the given options. - * - *

Example usage: - * - * 

{@code
-   * // With custom state
-   * Session session = genkit.createSession(
-   *     SessionOptions.builder().initialState(new MyState("John")).build());
-   *
-   * // With custom store and session ID
-   * Session session = genkit.createSession(
-   *     SessionOptions.builder()
-   *         .store(new RedisSessionStore<>())
-   *         .sessionId("my-session-123")
-   *         .initialState(new MyState())
-   *         .build());
-   * }
- * - * @param the session state type - * @param options the session options - * @return a new session - */ - public Session createSession(SessionOptions options) { - return Session.create(registry, options, agentRegistry); - } - - /** - * Loads an existing session from a store. - * - *

Example usage: - * - * 

{@code
-   * CompletableFuture> sessionFuture = genkit.loadSession(
-   *     "session-123",
-   *     SessionOptions.builder().store(mySessionStore).build());
-   * Session session = sessionFuture.get();
-   * if (session != null) {
-   *   Chat chat = session.chat();
-   *   // Continue conversation...
-   * }
-   * }
- * - * @param the session state type - * @param sessionId the session ID to load - * @param options the session options (must include store) - * @return a CompletableFuture containing the session, or null if not found - */ - public CompletableFuture> loadSession( - String sessionId, SessionOptions options) { - return Session.load(registry, sessionId, options, agentRegistry); - } - - /** - * Creates a simple chat without session persistence. - * - *

This is a convenience method for quick interactions without full session management. Use - * {@link #createSession()} for persistent multi-turn conversations. - * - *

Example usage: - * - * 

{@code
-   * Chat chat = genkit.chat(
-   *     ChatOptions.builder()
-   *         .model("openai/gpt-4o")
-   *         .system("You are a helpful assistant.")
-   *         .build());
-   * ModelResponse response = chat.send("Hello!");
-   * }
- * - * @param the state type (usually Void for simple chats) - * @param options the chat options - * @return a new chat instance - */ - public Chat chat(ChatOptions options) { - Session session = createSession(); - return session.chat(options); - } - - // ========================================================================= - // Agent and Interrupt Methods + // Interrupt Methods // ========================================================================= - /** - * Defines an agent that can be used as a tool in multi-agent systems. - * - *

Agents are specialized conversational components that can be delegated to by other agents. - * When an agent is called as a tool, it takes over the conversation with its own system prompt, - * model, and tools. - * - *

Example usage: - * - * 

{@code
-   * // Define a specialized agent
-   * Agent reservationAgent = genkit.defineAgent(
-   *     AgentConfig.builder()
-   *         .name("reservationAgent")
-   *         .description("Handles restaurant reservations")
-   *         .system("You are a reservation specialist...")
-   *         .model("openai/gpt-4o")
-   *         .tools(List.of(reservationTool, lookupTool))
-   *         .build());
-   *
-   * // Use in a parent agent
-   * Agent triageAgent = genkit.defineAgent(
-   *     AgentConfig.builder()
-   *         .name("triageAgent")
-   *         .description("Routes customer requests to specialists")
-   *         .system("You route customer requests to the appropriate specialist")
-   *         .agents(List.of(reservationAgent.getConfig()))
-   *         .build());
-   *
-   * // Start chat with triage agent
-   * Chat chat = genkit.chat(
-   *     ChatOptions.builder()
-   *         .model("openai/gpt-4o")
-   *         .system(triageAgent.getSystem())
-   *         .tools(triageAgent.getAllTools(agentRegistry))
-   *         .build());
-   * }
- * - * @param config the agent configuration - * @return the created agent - */ - public Agent defineAgent(AgentConfig config) { - Agent agent = new Agent(config); - // Register the agent as a tool - registry.registerAction(ActionType.TOOL, agent.asTool()); - // Register in agent registry for getAllTools lookup - agentRegistry.put(config.getName(), agent); - return agent; - } - - /** - * Gets an agent by name. - * - * @param name the agent name - * @return the agent, or null if not found - */ - public Agent getAgent(String name) { - return agentRegistry.get(name); - } - - /** - * Gets the agent registry. - * - *

This returns an unmodifiable view of all registered agents. - * - * @return the agent registry - */ - public Map getAgentRegistry() { - return java.util.Collections.unmodifiableMap(agentRegistry); - } - - /** - * Gets all tools for an agent, including sub-agent tools. - * - *

This is a convenience method that collects all tools from an agent, including tools from any - * sub-agents defined in its configuration. - * - * @param agent the agent - * @return the list of all tools - */ - public List> getAllToolsForAgent(Agent agent) { - return agent.getAllTools(agentRegistry); - } - /** * Defines an interrupt tool for human-in-the-loop interactions. * @@ -2260,34 +2199,6 @@ public Tool defineInterrupt(InterruptConfig config) { return interruptTool; } - /** - * Gets the current session from the context. - * - *

This method can be called from within tool execution to access the current session state. It - * uses a thread-local context that is set during chat execution. - * - *

Example usage: - * - * 

{@code
-   * Tool myTool = genkit
-   *     .defineTool("myTool", Input.class, Output.class, (ctx, input) -> {
-   *       Session session = genkit.currentSession();
-   *       if (session != null) {
-   *         Object state = session.getState();
-   *         // Use session state...
-   *       }
-   *       return new Output();
-   *     });
-   * }
- * - * @param the session state type - * @return the current session, or null if not in a session context - */ - @SuppressWarnings("unchecked") - public Session currentSession() { - return (Session) SessionContext.currentSession(); - } - // ========================================================================= // Evaluation Methods // ========================================================================= diff --git a/genkit/src/main/java/com/google/genkit/GenkitBeta.java b/genkit/src/main/java/com/google/genkit/GenkitBeta.java new file mode 100644 index 000000000..bac53f8c2 --- /dev/null +++ b/genkit/src/main/java/com/google/genkit/GenkitBeta.java @@ -0,0 +1,555 @@ +/* + * 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. + * + * SPDX-License-Identifier: Apache-2.0 + */ + +package com.google.genkit; + +import com.google.genkit.ai.FinishReason; +import com.google.genkit.ai.GenerateOptions; +import com.google.genkit.ai.Message; +import com.google.genkit.ai.ModelResponse; +import com.google.genkit.ai.Part; +import com.google.genkit.ai.ResumeOptions; +import com.google.genkit.ai.Role; +import com.google.genkit.ai.ToolRequest; +import com.google.genkit.ai.ToolResponse; +import com.google.genkit.ai.agent.Agent; +import com.google.genkit.ai.agent.AgentFinishReason; +import com.google.genkit.ai.agent.AgentFn; +import com.google.genkit.ai.agent.AgentResult; +import com.google.genkit.ai.agent.AgentStreamChunk; +import com.google.genkit.ai.agent.CustomAgentConfig; +import com.google.genkit.ai.agent.SessionRunner; +import com.google.genkit.ai.agent.ToolResume; +import com.google.genkit.core.GenkitException; +import com.google.genkit.prompt.ExecutablePrompt; +import java.util.ArrayList; +import java.util.HashMap; +import java.util.List; +import java.util.Map; + +/** + * Beta (experimental) API surface for Genkit. + * + *

This is the ergonomic, user-facing API for defining agents that run as bidi {@code agent} + * actions (so they appear and run in the Dev UI). It mirrors the JavaScript {@code GenkitBeta} + * (defineAgent / definePromptAgent / defineCustomAgent) and the Go {@code WithExperimental} gating. + * + *

All methods are gated behind the {@code experimental} flag (see {@link + * GenkitOptions.Builder#experimental} or the {@code GENKIT_EXPERIMENTAL} environment variable). + * When the flag is not enabled they throw a {@link GenkitException}. + * + *

Obtain an instance via {@link Genkit#beta()}. + */ +public final class GenkitBeta { + + private final Genkit genkit; + + /** + * Constructs a GenkitBeta for the given Genkit instance. Package-private; use {@link + * Genkit#beta()}. + * + * @param genkit the owning Genkit instance + */ + GenkitBeta(Genkit genkit) { + this.genkit = genkit; + } + + /** + * Throws if experimental features are not enabled on the owning Genkit instance. + * + * @throws GenkitException if the {@code experimental} flag is not set + */ + private void requireExperimental() { + if (!genkit.isExperimental()) { + throw new GenkitException( + "This is an experimental API. Enable experimental features via " + + "GenkitOptions.builder().experimental(true) (or set the GENKIT_EXPERIMENTAL " + + "environment variable to 'true') to use it."); + } + } + + /** + * Defines a custom agent from an explicit {@link CustomAgentConfig} and {@link AgentFn}. + * + *

This is the lowest-level beta agent factory: the caller supplies the per-turn logic + * directly. The agent is registered as a bidi {@code /agent/} action (plus companion + * snapshot/abort actions when a store is configured). + * + * @param the custom session state type + * @param config the custom agent configuration + * @param fn the per-turn agent function + * @return the registered agent + * @throws GenkitException if experimental features are not enabled + */ + public Agent defineCustomAgent(CustomAgentConfig config, AgentFn fn) { + requireExperimental(); + return com.google.genkit.ai.agent.internal.AgentActions.defineCustomAgent( + genkit.getRegistry(), config, fn); + } + + /** + * Defines an ergonomic, prompt-backed agent. + * + *

This is the {@code ai.defineAgent({name, system, tools, model, store})} entry point. It + * builds a {@link CustomAgentConfig} from the facade config and synthesizes an {@link AgentFn} + * that, for each turn, calls {@code generateStream} with the configured system prompt + tools and + * the session history (which already includes the just-added user message), streams model chunks + * back to the caller, and returns the model's response message as the turn result. + * + * @param the custom session state type + * @param config the agent configuration ({@code name} required) + * @return the registered agent + * @throws GenkitException if experimental features are not enabled + */ + public Agent defineAgent(com.google.genkit.agent.AgentConfig config) { + requireExperimental(); + // Plain-system path: the system instructions are fixed for the life of the agent, so the + // per-turn resolver simply returns the configured (static) system text on every turn. + String system = config.getSystem(); + return defineGenerateBackedAgent(config, runner -> system); + } + + /** + * Defines an agent backed by a registered prompt. + * + *

The agent's system instructions are produced by rendering the named prompt's Handlebars + * template on every turn: the prompt named by {@code config.getPromptName()} (falling + * back to {@code config.getName()}) is loaded once via {@link Genkit#prompt(String)}, and each + * turn its template is rendered with {@code config.getPromptInput()} merged with the current + * session state (see {@link #renderPromptSystem}). This means template variables such as {@code + * {{topic}}} interpolate the live {@code promptInput}/state values for that turn rather than + * leaking through as raw {@code {{...}}} placeholders. + * + *

If no matching prompt can be loaded (or its template is blank), this falls back to {@code + * config.getSystem()}, behaving like {@link #defineAgent(com.google.genkit.agent.AgentConfig)}. + * + * @param the custom session state type + * @param config the agent configuration ({@code name} required) + * @return the registered agent + * @throws GenkitException if experimental features are not enabled + */ + public Agent definePromptAgent(com.google.genkit.agent.AgentConfig config) { + requireExperimental(); + // Resolve the backing prompt (its static template) once, at definition time; the actual + // Handlebars render with promptInput/session-state happens per turn inside the SystemResolver. + ExecutablePrompt> prompt = resolvePrompt(config); + SystemResolver resolver = + runner -> { + if (prompt != null) { + String rendered = renderPromptSystem(prompt, config, runner); + if (rendered != null && !rendered.isBlank()) { + return rendered; + } + } + return config.getSystem(); + }; + return defineGenerateBackedAgent(config, resolver); + } + + /** + * Resolves the system instructions for a single turn of a generate-backed agent. + * + *

For {@link #defineAgent} this returns a fixed string; for {@link #definePromptAgent} it + * renders the backing prompt's Handlebars template against {@code promptInput}/session state. + * + * @param the custom session state type + */ + @FunctionalInterface + private interface SystemResolver { + /** + * Produces the system prompt for the current turn. + * + * @param runner the session runner for this turn (carries current session state) + * @return the system text to feed this turn's generate call, or {@code null} for none + */ + String resolve(SessionRunner runner); + } + + /** + * Builds a {@link CustomAgentConfig} from the facade config and registers a generate-backed agent + * whose {@link AgentFn} runs one generate (streaming) call per turn, resolving the system prompt + * for that turn via {@code systemResolver}. + */ + private Agent defineGenerateBackedAgent( + com.google.genkit.agent.AgentConfig config, SystemResolver systemResolver) { + + CustomAgentConfig customConfig = + CustomAgentConfig.builder() + .name(config.getName()) + .description(config.getDescription()) + .stateType(config.getStateType()) + .store(config.getStore()) + .clientTransform(config.getClientTransform()) + .build(); + + AgentFn fn = buildAgentFn(config, systemResolver); + + return com.google.genkit.ai.agent.internal.AgentActions.defineCustomAgent( + genkit.getRegistry(), customConfig, fn); + } + + /** + * Builds the per-turn agent function: it reads the session history (which already includes the + * user message added by {@code SessionRunner.runTurn}), generates a streaming response, forwards + * model chunks to the caller, and returns the model's reply message. + * + *

When the turn is a resume turn ({@code ctx.resume() != null} — i.e. {@code + * AgentChat.resume(...)} was called after a tool interrupt), this resumes the SAME generate call + * that was interrupted rather than starting a fresh one: it appends the resume's tool-response + * message to the session directly (so it is durably threaded into history — see {@code + * genkit.agent.yaml}'s "interrupt resume state accumulation" conformance case) and passes {@link + * ResumeOptions} built from {@link ToolResume} through to {@code generate}, using the session's + * last message (the previously-interrupted model message, already in {@code runner.getMessages()} + * from the interrupted turn) as the tail of the message history {@code generate} resumes from. + */ + private AgentFn buildAgentFn( + com.google.genkit.agent.AgentConfig config, SystemResolver systemResolver) { + return (SessionRunner runner, com.google.genkit.ai.agent.AgentFnContext ctx) -> { + ToolResume resume = ctx.resume(); + + // Resolve the system instructions for THIS turn. For defineAgent this is the fixed + // config.getSystem(); for definePromptAgent this renders the backing prompt's Handlebars + // template against the current promptInput/session state, so template variables interpolate + // per turn rather than reaching the model as raw {{...}} placeholders. + String system = systemResolver.resolve(runner); + + // Validate a resume directive against the resumed session's pending interrupts BEFORE running + // generate. A respond/restart whose (name, ref) does not match a tool request in the last + // model message — or a restart whose input differs from the original — is rejected with + // INVALID_ARGUMENT, which SessionRunner turns into a graceful FAILED output (not a thrown + // error). Mirrors JS resolve-tool-requests.ts:244-266 / Go generate.go:1142-1187. + if (resume != null) { + validateResumeDirectives(resume, runner.getMessages()); + } + + // Forward the run's request-scoped user context (e.g. {"auth": {...}}) into the generate + // call so that tools executed during this turn observe it via ctx.getContext(). + Map userContext = ctx.context() != null ? ctx.context().getContext() : null; + + GenerateOptions.Builder optsBuilder = + GenerateOptions.builder() + .model(config.getModel()) + .system(system) + .tools(config.getTools()) + .config(config.getConfig()) + .context(userContext) + .maxTurns(config.getMaxTurns()); + + // Snapshot the messages we hand to generate (the session history, including the just-added + // user message, plus — on a resume — the previously-interrupted model message that is still + // the session tail). generate resumes from this list; on the resume path, + // Genkit.handleResumeOption requires the LAST message here to still be the interrupted MODEL + // message (it builds its own tool-response message from ResumeOptions and appends it before + // calling the model). So we do NOT pre-append any tool-response message to the runner. + List sentMessages = new ArrayList<>(runner.getMessages()); + optsBuilder.messages(sentMessages); + if (resume != null) { + optsBuilder.resume(toResumeOptions(resume)); + } + + GenerateOptions opts = optsBuilder.build(); + + ModelResponse resp = + genkit.generateStream( + opts, + chunk -> + ctx.sendChunk().accept(AgentStreamChunk.builder().modelChunk(chunk).build())); + + // Thread the FULL post-turn message history into session state (matches JS action.ts:427-431 + // and Go generate.go:876-879, which carry the intermediate model-tool-request and + // tool-response messages alongside the final model message). ModelResponse.getMessages() + // returns the request messages (which the tool loop grows in place) followed by the final + // model message; everything past the messages we sent is newly produced this turn. The + // generate call may prepend a synthesized system message, so account for that offset. + int systemOffset = system != null ? 1 : 0; + List allMessages = resp.getMessages(); + int newStart = systemOffset + sentMessages.size(); + Message finalMessage = resp.getMessage(); + if (allMessages != null && newStart >= 0 && newStart < allMessages.size()) { + // Fold every newly-produced message except the terminal one directly into history; the + // terminal message is returned as the AgentResult so AgentActions appends it (and surfaces + // it as the output message). + for (int i = newStart; i < allMessages.size() - 1; i++) { + runner.addMessages(allMessages.get(i)); + } + finalMessage = allMessages.get(allMessages.size() - 1); + } + + return AgentResult.builder() + .message(finalMessage) + .finishReason(mapFinishReason(resp)) + .build(); + }; + } + + /** + * Converts an agent-level {@link ToolResume} (parts) into generate-level {@link ResumeOptions} + * (typed tool responses/requests), matching the shapes {@code Genkit.generate(...).resume(...)} + * expects (see {@code samples/interrupts}). + */ + private static ResumeOptions toResumeOptions(ToolResume resume) { + ResumeOptions.Builder builder = ResumeOptions.builder(); + if (resume.getRespond() != null) { + List responses = new ArrayList<>(); + for (Part part : resume.getRespond()) { + if (part != null && part.getToolResponse() != null) { + responses.add(part.getToolResponse()); + } + } + if (!responses.isEmpty()) { + builder.respond(responses); + } + } + if (resume.getRestart() != null) { + List requests = new ArrayList<>(); + for (Part part : resume.getRestart()) { + if (part != null && part.getToolRequest() != null) { + ToolRequest req = part.getToolRequest(); + // Carry the restart directive's Part-level metadata (resumed / replacedInput) onto the + // ToolRequest so it survives into ResumeOptions.restart (a List that drops + // the Part wrapper). The generate loop re-attaches it to the restart Part so the tool + // observes its resumed status. Without this, restart metadata would be silently lost. + if (part.getMetadata() != null && !part.getMetadata().isEmpty()) { + Map merged = + req.getMetadata() != null ? new HashMap<>(req.getMetadata()) : new HashMap<>(); + merged.putAll(part.getMetadata()); + req.setMetadata(merged); + } + requests.add(req); + } + } + if (!requests.isEmpty()) { + builder.restart(requests); + } + } + return builder.build(); + } + + /** + * Validates a resume directive against the pending interrupted tool requests in the resumed + * session history. Throws {@link GenkitException} with {@code INVALID_ARGUMENT} — which {@code + * SessionRunner} turns into a graceful {@code finishReason: failed} output — when: + * + *
    + *
  • a {@code respond} directive's (name, ref) matches no tool request in the last model + * message ("not found in session history"), or + *
  • a {@code restart} directive's (name, ref) matches none ("not found in session history"), + * or its input differs from the original tool request's input ("modified inputs"). + *
+ * + *

Matching key is (name, ref), against the tool requests in the last model message of the + * session (the interrupted turn). Mirrors JS {@code resolve-tool-requests.ts:244-266} and Go + * {@code generate.go:1142-1187}. + */ + private static void validateResumeDirectives(ToolResume resume, List history) { + List pending = lastModelToolRequests(history); + + if (resume.getRespond() != null) { + for (Part part : resume.getRespond()) { + if (part == null || part.getToolResponse() == null) { + continue; + } + ToolResponse tr = part.getToolResponse(); + if (findToolRequest(pending, tr.getName(), tr.getRef()) == null) { + throw GenkitException.builder() + .message( + "tool response for '" + + tr.getName() + + "' (ref=" + + tr.getRef() + + ") not found in session history") + .errorCode("INVALID_ARGUMENT") + .build(); + } + } + } + + if (resume.getRestart() != null) { + for (Part part : resume.getRestart()) { + if (part == null || part.getToolRequest() == null) { + continue; + } + ToolRequest req = part.getToolRequest(); + ToolRequest original = findToolRequest(pending, req.getName(), req.getRef()); + if (original == null) { + throw GenkitException.builder() + .message( + "restart for tool '" + + req.getName() + + "' (ref=" + + req.getRef() + + ") not found in session history") + .errorCode("INVALID_ARGUMENT") + .build(); + } + if (!inputsEqual(original.getInput(), req.getInput())) { + throw GenkitException.builder() + .message( + "restart for tool '" + + req.getName() + + "' has modified inputs (does not match the original tool request)") + .errorCode("INVALID_ARGUMENT") + .build(); + } + } + } + } + + /** + * Returns the tool requests in the last {@code model}-role message of {@code history} (the + * interrupted turn), or an empty list if there is no such message. + */ + private static List lastModelToolRequests(List history) { + List out = new ArrayList<>(); + if (history == null) { + return out; + } + for (int i = history.size() - 1; i >= 0; i--) { + Message m = history.get(i); + if (m == null || m.getRole() != Role.MODEL) { + continue; + } + if (m.getContent() != null) { + for (Part p : m.getContent()) { + if (p.getToolRequest() != null) { + out.add(p.getToolRequest()); + } + } + } + break; // only the most recent model message + } + return out; + } + + /** Finds a tool request in {@code pending} matching {@code name} and {@code ref}, or null. */ + private static ToolRequest findToolRequest(List pending, String name, String ref) { + for (ToolRequest tr : pending) { + if (java.util.Objects.equals(tr.getName(), name) + && java.util.Objects.equals(tr.getRef(), ref)) { + return tr; + } + } + return null; + } + + /** + * Compares two tool inputs for equality by canonical JSON (tolerates map/POJO representations). + */ + private static boolean inputsEqual(Object a, Object b) { + try { + var mapper = com.google.genkit.core.JsonUtils.getObjectMapper(); + return mapper.valueToTree(a).equals(mapper.valueToTree(b)); + } catch (Exception e) { + return java.util.Objects.equals(a, b); + } + } + + /** + * Loads the backing {@link ExecutablePrompt} for a prompt-backed agent, once, at definition time. + * Attempts the prompt named by {@code config.getPromptName()} (falling back to {@code + * config.getName()}). Returns {@code null} when no matching prompt can be loaded (in which case + * the agent falls back to {@code config.getSystem()} per turn). The prompt's template is static; + * only the per-turn {@linkplain #renderPromptSystem render} varies with {@code + * promptInput}/state. + */ + private ExecutablePrompt> resolvePrompt( + com.google.genkit.agent.AgentConfig config) { + String promptName = config.getPromptName() != null ? config.getPromptName() : config.getName(); + if (promptName == null) { + return null; + } + try { + return genkit.prompt(promptName); + } catch (Exception e) { + // No matching prompt available — the caller falls back to the explicit system text. + return null; + } + } + + /** + * Renders a prompt-backed agent's system instructions for the current turn by applying the + * prompt's Handlebars template to a rendering context built from {@code config.getPromptInput()} + * merged with the current session state. This reuses the existing dotprompt/Handlebars engine + * ({@link ExecutablePrompt#render}) rather than introducing a separate templating layer. + * + *

Context precedence (later wins): the session's custom state (when it is a {@code Map}), then + * the configured {@code promptInput}. When {@code promptInput} is a non-{@code Map} POJO it is + * used directly as the render context (its bean properties are visible to Handlebars) and session + * state is not merged. On any render failure this returns {@code null} so the caller falls back + * to {@code config.getSystem()}. + */ + private String renderPromptSystem( + ExecutablePrompt> prompt, + com.google.genkit.agent.AgentConfig config, + SessionRunner runner) { + try { + Object promptInput = config.getPromptInput(); + if (promptInput != null && !(promptInput instanceof Map)) { + // POJO promptInput: render directly against the bean so its fields interpolate. The backing + // DotPrompt is typed to Map, but Handlebars renders any bean, so we render + // through a raw reference to feed the POJO context. + @SuppressWarnings({"unchecked", "rawtypes"}) + String rendered = + ((com.google.genkit.prompt.DotPrompt) prompt.getDotPrompt()).render(promptInput); + return rendered; + } + Map context = new HashMap<>(); + Object custom = runner != null ? runner.getCustom() : null; + if (custom instanceof Map) { + @SuppressWarnings("unchecked") + Map customMap = (Map) custom; + context.putAll(customMap); + } + if (promptInput instanceof Map) { + @SuppressWarnings("unchecked") + Map inputMap = (Map) promptInput; + context.putAll(inputMap); + } + return prompt.render(context); + } catch (Exception e) { + // Rendering failed (e.g. missing helper) — fall back to the explicit system text. + return null; + } + } + + /** + * Maps a model {@link FinishReason} to the agent-level {@link AgentFinishReason}. Defaults to + * {@link AgentFinishReason#STOP} when the response has no finish reason. + */ + private static AgentFinishReason mapFinishReason(ModelResponse resp) { + FinishReason fr = resp != null ? resp.getFinishReason() : null; + if (fr == null) { + return AgentFinishReason.STOP; + } + switch (fr) { + case LENGTH: + return AgentFinishReason.LENGTH; + case BLOCKED: + return AgentFinishReason.BLOCKED; + case INTERRUPTED: + return AgentFinishReason.INTERRUPTED; + case OTHER: + return AgentFinishReason.OTHER; + case UNKNOWN: + return AgentFinishReason.UNKNOWN; + case STOP: + default: + return AgentFinishReason.STOP; + } + } +} diff --git a/genkit/src/main/java/com/google/genkit/GenkitOptions.java b/genkit/src/main/java/com/google/genkit/GenkitOptions.java index a6d3ae748..67670e8fa 100644 --- a/genkit/src/main/java/com/google/genkit/GenkitOptions.java +++ b/genkit/src/main/java/com/google/genkit/GenkitOptions.java @@ -22,6 +22,7 @@ public class GenkitOptions { private final boolean devMode; + private final boolean experimental; private final int reflectionPort; private final String projectRoot; private final String promptDir; @@ -29,6 +30,7 @@ public class GenkitOptions { private GenkitOptions(Builder builder) { this.devMode = builder.devMode; + this.experimental = builder.experimental; this.reflectionPort = builder.reflectionPort; this.projectRoot = builder.projectRoot; this.promptDir = builder.promptDir; @@ -53,6 +55,15 @@ public boolean isDevMode() { return devMode; } + /** + * Returns whether experimental features (such as the beta agent APIs) are enabled. + * + * @return true if experimental features are enabled + */ + public boolean isExperimental() { + return experimental; + } + /** * Returns the reflection server port. * @@ -93,6 +104,7 @@ public String getName() { /** Builder for GenkitOptions. */ public static class Builder { private boolean devMode = isDevModeFromEnv(); + private boolean experimental = isExperimentalFromEnv(); private int reflectionPort = getReflectionPortFromEnv(); private String projectRoot = System.getProperty("user.dir"); private String promptDir = "/prompts"; @@ -102,6 +114,11 @@ private static boolean isDevModeFromEnv() { return "dev".equalsIgnoreCase(System.getenv("GENKIT_ENV")); } + private static boolean isExperimentalFromEnv() { + String value = System.getenv("GENKIT_EXPERIMENTAL"); + return "true".equalsIgnoreCase(value) || "1".equals(value); + } + private static int getReflectionPortFromEnv() { String port = System.getenv("GENKIT_REFLECTION_PORT"); if (port != null) { @@ -119,6 +136,11 @@ public Builder devMode(boolean devMode) { return this; } + public Builder experimental(boolean experimental) { + this.experimental = experimental; + return this; + } + public Builder reflectionPort(int reflectionPort) { this.reflectionPort = reflectionPort; return this; diff --git a/genkit/src/main/java/com/google/genkit/ReflectionServer.java b/genkit/src/main/java/com/google/genkit/ReflectionServer.java index d03d7be9e..6cc99bcb8 100644 --- a/genkit/src/main/java/com/google/genkit/ReflectionServer.java +++ b/genkit/src/main/java/com/google/genkit/ReflectionServer.java @@ -24,6 +24,8 @@ import com.google.genkit.core.ActionContext; import com.google.genkit.core.ActionDesc; import com.google.genkit.core.ActionRunResult; +import com.google.genkit.core.BidiAction; +import com.google.genkit.core.BufferedInputSource; import com.google.genkit.core.GenkitException; import com.google.genkit.core.JsonUtils; import com.google.genkit.core.Registry; @@ -479,9 +481,11 @@ private String handleRunAction(String body) throws GenkitException { throw new GenkitException("Action not found: " + key); } - ActionContext context = new ActionContext(registry); + ActionContext context = + ActionContext.builder().registry(registry).context(parseContext(requestNode)).build(); + JsonNode init = requestNode.has("init") ? requestNode.get("init") : null; - ActionRunResult result = action.runJsonWithTelemetry(context, input, null); + ActionRunResult result = runActionWithInit(action, context, input, init, null); // Build response according to Genkit reflection API spec: // { result: ..., telemetry: { traceId: "..." } } @@ -540,6 +544,56 @@ private String handleStreamTrace(String body) { } } + /** + * Runs an action, threading the optional {@code init} (session source) for agent (bidi) + * actions. The Dev UI sends {@code init} on each agent turn — client-managed {@code state} or + * server {@code snapshotId} — which a plain unary run would drop, breaking multi-turn chat. + * Bidi actions are therefore driven one-turn-per-request with the init; non-bidi actions use + * the existing unary path unchanged. + */ + /** + * Parses the optional {@code context} object from a runAction request body into a {@code + * Map} using the shared ObjectMapper. The Dev UI "Execution context" panel sends + * this (e.g. {@code {"auth": {"user": "alice"}}}); it is threaded into the run's ActionContext. + * + * @param requestNode the parsed request body + * @return the parsed context map, or null if absent/blank + */ + private Map parseContext(JsonNode requestNode) { + if (requestNode == null + || !requestNode.has("context") + || requestNode.get("context").isNull()) { + return null; + } + JsonNode contextNode = requestNode.get("context"); + if (!contextNode.isObject()) { + return null; + } + return JsonUtils.getObjectMapper() + .convertValue( + contextNode, + new com.fasterxml.jackson.core.type.TypeReference>() {}); + } + + private ActionRunResult runActionWithInit( + Action action, + ActionContext context, + JsonNode input, + JsonNode init, + java.util.function.Consumer streamCallback) + throws GenkitException { + if (action instanceof BidiAction) { + BufferedInputSource inputs = new BufferedInputSource<>(); + if (input != null) { + inputs.offer(input); + } + inputs.end(); + return ((BidiAction) action) + .runBidiJsonWithTelemetry(context, init, inputs, streamCallback); + } + return action.runJsonWithTelemetry(context, input, streamCallback); + } + /** * Handle runAction with streaming format (when ?stream=true is set). The Dev UI expects: 1. * Content-Type: text/plain with Content-Length 2. X-Genkit-Trace-Id and X-Genkit-Version @@ -560,8 +614,10 @@ private void handleStreamingRunAction(String body, Response response, Callback c throw new GenkitException("Action not found: " + actionKey); } - ActionContext context = new ActionContext(registry); - ActionRunResult result = action.runJsonWithTelemetry(context, input, null); + ActionContext context = + ActionContext.builder().registry(registry).context(parseContext(requestNode)).build(); + JsonNode init = requestNode.has("init") ? requestNode.get("init") : null; + ActionRunResult result = runActionWithInit(action, context, input, init, null); // Build the final response with result and telemetry Map responseData = new HashMap<>(); diff --git a/genkit/src/main/java/com/google/genkit/ReflectionServerV2.java b/genkit/src/main/java/com/google/genkit/ReflectionServerV2.java index 7f8a20058..d6e145910 100644 --- a/genkit/src/main/java/com/google/genkit/ReflectionServerV2.java +++ b/genkit/src/main/java/com/google/genkit/ReflectionServerV2.java @@ -23,6 +23,9 @@ import com.google.genkit.core.Action; import com.google.genkit.core.ActionContext; import com.google.genkit.core.ActionDesc; +import com.google.genkit.core.BidiAction; +import com.google.genkit.core.BufferedInputSource; +import com.google.genkit.core.InputSource; import com.google.genkit.core.JsonUtils; import com.google.genkit.core.Registry; import com.google.genkit.core.tracing.Tracer; @@ -79,6 +82,21 @@ public class ReflectionServerV2 { /** Maps traceId → Thread for targeted cancellation of running actions. */ private final ConcurrentHashMap activeActions = new ConcurrentHashMap<>(); + /** + * Maps a bidi {@code runAction} requestId → its input source. Populated by {@code runAction} with + * {@code streamInput:true} and by {@code sendInputStreamChunk}/{@code endInputStream} (via {@code + * computeIfAbsent}) so that input chunks arriving before the action handler starts are buffered + * rather than dropped. + */ + private final ConcurrentHashMap> bidiSessions = + new ConcurrentHashMap<>(); + + /** + * Optional outbound-message sink, used by tests to capture JSON-RPC messages without a live + * WebSocket. When null (production), messages go to {@link #webSocket}. + */ + private volatile java.util.function.Consumer outboundSink; + private final ScheduledExecutorService scheduler = Executors.newSingleThreadScheduledExecutor(); private final ExecutorService actionExecutor = Executors.newCachedThreadPool(); @@ -181,12 +199,30 @@ private void scheduleReconnect() { // ========================================================================= private void send(String message) { + java.util.function.Consumer sink = this.outboundSink; + if (sink != null) { + sink.accept(message); + return; + } WebSocket ws = this.webSocket; if (ws != null) { ws.sendText(message, true); } } + /** + * Test seam: routes outbound JSON-RPC messages to {@code sink} instead of the WebSocket, and lets + * a test feed inbound messages via {@link #handleMessageForTesting(String)}. + */ + void setOutboundSinkForTesting(java.util.function.Consumer sink) { + this.outboundSink = sink; + } + + /** Test seam: dispatches a raw inbound JSON-RPC message as if received over the WebSocket. */ + void handleMessageForTesting(String text) { + handleMessage(text); + } + private void sendResponse(String id, Object result) { Map response = new HashMap<>(); response.put("jsonrpc", "2.0"); @@ -320,6 +356,12 @@ private void handleRequest(JsonNode request) { case "runAction": handleRunAction(id, params); break; + case "sendInputStreamChunk": + handleSendInputStreamChunk(params); + break; + case "endInputStream": + handleEndInputStream(params); + break; case "configure": handleConfigure(params); break; @@ -395,6 +437,26 @@ private void handleListValues(String requestId, JsonNode params) { private void handleRunAction(String requestId, JsonNode params) { if (requestId == null) return; + // For bidi actions driven with streamInput, pre-register the input source synchronously + // (before dispatching to the executor) so that sendInputStreamChunk/endInputStream + // notifications arriving early are buffered rather than dropped. + final boolean streamInput = + params != null && params.has("streamInput") && params.get("streamInput").asBoolean(); + if (streamInput) { + bidiSessions.computeIfAbsent(requestId, k -> new BufferedInputSource<>()); + } + + String _key = params != null && params.has("key") ? params.get("key").asText() : null; + boolean _stream = params != null && params.has("stream") && params.get("stream").asBoolean(); + boolean _initPresent = params != null && params.has("init") && !params.get("init").isNull(); + logger.debug( + "→ runAction id={} key={} stream={} streamInput={} initPresent={}", + requestId, + _key, + _stream, + streamInput, + _initPresent); + // Run action in a separate thread so we don't block the WebSocket message loop actionExecutor.submit( () -> { @@ -402,6 +464,7 @@ private void handleRunAction(String requestId, JsonNode params) { try { String key = params.has("key") ? params.get("key").asText() : null; JsonNode input = params.has("input") ? params.get("input") : null; + JsonNode init = params.has("init") ? params.get("init") : null; boolean stream = params.has("stream") && params.get("stream").asBoolean(); if (key == null) { @@ -420,6 +483,8 @@ private void handleRunAction(String requestId, JsonNode params) { if (stream) { streamCallback = (chunk) -> { + logger.debug( + "← streamChunk requestId={} kind={}", requestId, streamChunkKind(chunk)); Map chunkParams = new HashMap<>(); chunkParams.put("requestId", requestId); chunkParams.put("chunk", chunk); @@ -427,7 +492,8 @@ private void handleRunAction(String requestId, JsonNode params) { }; } - ActionContext context = new ActionContext(registry); + ActionContext context = + ActionContext.builder().registry(registry).context(parseContext(params)).build(); // Create the telemetry span manually so we can get the traceId // BEFORE the action starts producing stream chunks. @@ -459,7 +525,21 @@ private void handleRunAction(String requestId, JsonNode params) { sendNotification("runActionState", stateParams); } - // Now run the action - stream chunks will be sent AFTER traceId + // Now run the action - stream chunks will be sent AFTER traceId. + // Bidi actions driven with streamInput consume their inputs from the + // pre-registered input source rather than the single `input` param. + if (streamInput && action instanceof BidiAction) { + InputSource inputs = bidiSessions.get(requestId); + if (inputs == null) { + inputs = new BufferedInputSource<>(); + } + return ((BidiAction) action) + .runBidiJson( + context.withSpanContext(spanCtx), + init, + inputs, + finalStreamCallback); + } return action.runJson( context.withSpanContext(spanCtx), in, finalStreamCallback); }); @@ -480,6 +560,19 @@ private void handleRunAction(String requestId, JsonNode params) { telemetry.put("traceId", traceId); responseResult.put("telemetry", telemetry); } + if (logger.isDebugEnabled()) { + String resultKeys = + jsonResult != null + ? String.join( + ",", + java.util.stream.StreamSupport.stream( + java.util.Spliterators.spliteratorUnknownSize( + jsonResult.fieldNames(), 0), + false) + .collect(java.util.stream.Collectors.toList())) + : "null"; + logger.debug("← result id={} (keys={})", requestId, resultKeys); + } sendResponse(requestId, responseResult); } catch (Exception e) { @@ -506,10 +599,70 @@ private void handleRunAction(String requestId, JsonNode params) { if (traceId != null) { activeActions.remove(traceId); } + BufferedInputSource src = bidiSessions.remove(requestId); + if (src != null) { + src.close(); + } } }); } + /** + * Handles a {@code sendInputStreamChunk} notification: enqueues one input chunk onto the bidi + * session's input source. Uses {@code computeIfAbsent} so a chunk that arrives before the + * matching {@code runAction} handler started is still buffered. + */ + private void handleSendInputStreamChunk(JsonNode params) { + if (params == null || !params.has("requestId")) { + return; + } + String requestId = params.get("requestId").asText(); + JsonNode chunk = params.has("chunk") ? params.get("chunk") : null; + if (chunk == null) { + return; + } + logger.debug( + "→ sendInputStreamChunk requestId={} chunkKeys={}", requestId, chunkFieldNames(chunk)); + bidiSessions.computeIfAbsent(requestId, k -> new BufferedInputSource<>()).offer(chunk); + } + + /** Handles an {@code endInputStream} notification: signals end-of-input for the bidi session. */ + private void handleEndInputStream(JsonNode params) { + if (params == null || !params.has("requestId")) { + return; + } + String requestId = params.get("requestId").asText(); + logger.debug("→ endInputStream requestId={}", requestId); + bidiSessions.computeIfAbsent(requestId, k -> new BufferedInputSource<>()).end(); + } + + /** Returns a comma-separated list of field names present on a chunk node (null-safe). */ + private static String chunkFieldNames(JsonNode chunk) { + if (chunk == null || !chunk.isObject()) return "null"; + StringBuilder sb = new StringBuilder(); + chunk + .fieldNames() + .forEachRemaining( + f -> { + if (sb.length() > 0) sb.append(','); + sb.append(f); + }); + return sb.length() > 0 ? sb.toString() : "(empty)"; + } + + /** + * Returns the kind of an {@code AgentStreamChunk} for debug logging: whichever of {@code + * modelChunk}, {@code customPatch}, {@code artifact}, or {@code turnEnd} is present, else the raw + * field names. + */ + private static String streamChunkKind(JsonNode chunk) { + if (chunk == null || !chunk.isObject()) return "null"; + for (String known : new String[] {"turnEnd", "modelChunk", "customPatch", "artifact"}) { + if (chunk.has(known)) return known; + } + return chunkFieldNames(chunk); + } + private void handleConfigure(JsonNode params) { if (params == null) return; @@ -554,6 +707,26 @@ private static String getStackTraceString(Throwable e) { return sw.toString(); } + /** + * Parses the optional {@code context} object from runAction params into a {@code + * Map}. The Dev UI "Execution context" panel sends this (e.g. {@code {"auth": + * {"user": "alice"}}}); it is threaded into the run's ActionContext so tools/flows can read it. + * + * @param params the JSON-RPC params for runAction + * @return the parsed context map, or null if absent/blank + */ + private static Map parseContext(JsonNode params) { + if (params == null || !params.has("context") || params.get("context").isNull()) { + return null; + } + JsonNode contextNode = params.get("context"); + if (!contextNode.isObject()) { + return null; + } + return objectMapper.convertValue( + contextNode, new com.fasterxml.jackson.core.type.TypeReference>() {}); + } + // ========================================================================= // WebSocket Listener // ========================================================================= diff --git a/genkit/src/main/java/com/google/genkit/agent/AgentConfig.java b/genkit/src/main/java/com/google/genkit/agent/AgentConfig.java new file mode 100644 index 000000000..8f6a445c8 --- /dev/null +++ b/genkit/src/main/java/com/google/genkit/agent/AgentConfig.java @@ -0,0 +1,374 @@ +/* + * 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. + * + * SPDX-License-Identifier: Apache-2.0 + */ + +package com.google.genkit.agent; + +import com.google.genkit.ai.GenerationConfig; +import com.google.genkit.ai.Tool; +import com.google.genkit.ai.agent.ClientTransform; +import com.google.genkit.ai.agent.SessionStore; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.List; + +/** + * Facade-level configuration for the ergonomic, prompt-backed agents created via {@code + * GenkitBeta.defineAgent} and {@code GenkitBeta.definePromptAgent}. + * + *

Unlike {@link com.google.genkit.ai.agent.CustomAgentConfig} (which requires the caller to + * supply the per-turn {@code AgentFn} themselves), this config captures the high-level pieces of a + * generate-backed agent — model, system prompt, tools, generation config — and the {@code + * defineAgent} implementation synthesizes the {@code AgentFn} that calls {@code generate}. + * + *

The only required field is {@link #getName()}. + * + * @param the type of custom session state + */ +public final class AgentConfig { + + private final String name; + private final String description; + private final String model; + private final String system; + private final List> tools; + private final GenerationConfig config; + private final Integer maxTurns; + private final SessionStore store; + private final Class stateType; + private final ClientTransform clientTransform; + private final String promptName; + private final Object promptInput; + + private AgentConfig(Builder builder) { + this.name = builder.name; + this.description = builder.description; + this.model = builder.model; + this.system = builder.system; + this.tools = builder.tools; + this.config = builder.config; + this.maxTurns = builder.maxTurns; + this.store = builder.store; + this.stateType = builder.stateType; + this.clientTransform = builder.clientTransform; + this.promptName = builder.promptName; + this.promptInput = builder.promptInput; + } + + /** + * Creates a builder for AgentConfig. + * + * @param the type of custom session state + * @return a new builder + */ + public static Builder builder() { + return new Builder<>(); + } + + /** + * Returns the agent's registered name. + * + * @return the agent name (never null on a built config) + */ + public String getName() { + return name; + } + + /** + * Returns the agent's human-readable description. + * + * @return the description, or {@code null} if not set + */ + public String getDescription() { + return description; + } + + /** + * Returns the model name to use for generation. + * + * @return the model name, or {@code null} to rely on the Genkit default model + */ + public String getModel() { + return model; + } + + /** + * Returns the system prompt for the agent. + * + * @return the system prompt, or {@code null} if not set + */ + public String getSystem() { + return system; + } + + /** + * Returns the tools available to the agent. + * + * @return the tools, or {@code null} if not set + */ + public List> getTools() { + return tools; + } + + /** + * Returns the generation configuration. + * + * @return the config, or {@code null} if not set + */ + public GenerationConfig getConfig() { + return config; + } + + /** + * Returns the maximum number of tool-execution turns for a single generate call. + * + * @return the max turns, or {@code null} to use the generate default + */ + public Integer getMaxTurns() { + return maxTurns; + } + + /** + * Returns the session store for server-managed agents. + * + *

When {@code null}, the agent operates in client-managed mode. + * + * @return the session store, or {@code null} for client-managed mode + */ + public SessionStore getStore() { + return store; + } + + /** + * Returns the Java class for the agent's custom state type. + * + * @return the state type class, or {@code null} if not specified + */ + public Class getStateType() { + return stateType; + } + + /** + * Returns the client-transform applied to session state before returning it to the caller in + * client-managed mode. + * + * @return the client transform, or {@code null} if not set + */ + public ClientTransform getClientTransform() { + return clientTransform; + } + + /** + * Returns the name of a registered prompt to drive {@code definePromptAgent}. + * + * @return the prompt name, or {@code null} if not set + */ + public String getPromptName() { + return promptName; + } + + /** + * Returns the input passed to the prompt for {@code definePromptAgent}. + * + * @return the prompt input, or {@code null} if not set + */ + public Object getPromptInput() { + return promptInput; + } + + /** + * Builder for {@link AgentConfig}. + * + * @param the type of custom session state + */ + public static final class Builder { + + private String name; + private String description; + private String model; + private String system; + private List> tools; + private GenerationConfig config; + private Integer maxTurns; + private SessionStore store; + private Class stateType; + private ClientTransform clientTransform; + private String promptName; + private Object promptInput; + + private Builder() {} + + /** + * Sets the agent's registered name (required). + * + * @param name the agent name + * @return this builder + */ + public Builder name(String name) { + this.name = name; + return this; + } + + /** + * Sets the agent's human-readable description. + * + * @param description the description + * @return this builder + */ + public Builder description(String description) { + this.description = description; + return this; + } + + /** + * Sets the model name to use for generation. When {@code null}, the Genkit default model is + * used. + * + * @param model the model name + * @return this builder + */ + public Builder model(String model) { + this.model = model; + return this; + } + + /** + * Sets the system prompt for the agent. + * + * @param system the system prompt + * @return this builder + */ + public Builder system(String system) { + this.system = system; + return this; + } + + /** + * Sets the tools available to the agent. + * + * @param tools the tools + * @return this builder + */ + public Builder tools(List> tools) { + this.tools = tools != null ? new ArrayList<>(tools) : null; + return this; + } + + /** + * Sets the tools available to the agent. + * + * @param tools the tools + * @return this builder + */ + public Builder tools(Tool... tools) { + this.tools = tools != null ? new ArrayList<>(Arrays.asList(tools)) : null; + return this; + } + + /** + * Sets the generation configuration. + * + * @param config the config + * @return this builder + */ + public Builder config(GenerationConfig config) { + this.config = config; + return this; + } + + /** + * Sets the maximum number of tool-execution turns for a single generate call. + * + * @param maxTurns the max turns + * @return this builder + */ + public Builder maxTurns(Integer maxTurns) { + this.maxTurns = maxTurns; + return this; + } + + /** + * Sets the session store for server-managed mode. Pass {@code null} (or omit) for + * client-managed mode. + * + * @param store the session store + * @return this builder + */ + public Builder store(SessionStore store) { + this.store = store; + return this; + } + + /** + * Sets the Java class for the agent's custom state type. + * + * @param stateType the state type class + * @return this builder + */ + public Builder stateType(Class stateType) { + this.stateType = stateType; + return this; + } + + /** + * Sets the client-transform applied to session state before returning it to the caller in + * client-managed mode. + * + * @param clientTransform the transform + * @return this builder + */ + public Builder clientTransform(ClientTransform clientTransform) { + this.clientTransform = clientTransform; + return this; + } + + /** + * Sets the name of a registered prompt to drive {@code definePromptAgent}. + * + * @param promptName the prompt name + * @return this builder + */ + public Builder promptName(String promptName) { + this.promptName = promptName; + return this; + } + + /** + * Sets the input passed to the prompt for {@code definePromptAgent}. + * + * @param promptInput the prompt input + * @return this builder + */ + public Builder promptInput(Object promptInput) { + this.promptInput = promptInput; + return this; + } + + /** + * Builds the {@link AgentConfig}. + * + * @return a new {@link AgentConfig} + * @throws IllegalStateException if {@code name} is null or blank + */ + public AgentConfig build() { + if (name == null || name.isBlank()) { + throw new IllegalStateException("name is required"); + } + return new AgentConfig<>(this); + } + } +} diff --git a/genkit/src/main/java/com/google/genkit/client/HttpAgentTransport.java b/genkit/src/main/java/com/google/genkit/client/HttpAgentTransport.java new file mode 100644 index 000000000..5a95a5e3b --- /dev/null +++ b/genkit/src/main/java/com/google/genkit/client/HttpAgentTransport.java @@ -0,0 +1,267 @@ +/* + * 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. + * + * SPDX-License-Identifier: Apache-2.0 + */ + +package com.google.genkit.client; + +import com.fasterxml.jackson.databind.JsonNode; +import com.fasterxml.jackson.databind.ObjectMapper; +import com.fasterxml.jackson.databind.node.ObjectNode; +import com.google.genkit.ai.agent.AgentInit; +import com.google.genkit.ai.agent.AgentInput; +import com.google.genkit.ai.agent.AgentOutput; +import com.google.genkit.ai.agent.AgentStreamChunk; +import com.google.genkit.ai.agent.AgentTransport; +import com.google.genkit.ai.agent.GetSnapshotRequest; +import com.google.genkit.ai.agent.SessionSnapshot; +import com.google.genkit.ai.agent.SnapshotStatus; +import com.google.genkit.core.GenkitException; +import com.google.genkit.core.JsonUtils; +import java.io.BufferedReader; +import java.io.InputStreamReader; +import java.net.URI; +import java.net.http.HttpClient; +import java.net.http.HttpRequest; +import java.net.http.HttpResponse; +import java.nio.charset.StandardCharsets; +import java.util.Map; +import java.util.function.Consumer; + +/** + * HTTP implementation of {@link AgentTransport} that speaks the Jetty {@code AgentHandler} wire + * format. + * + *

Wire format: + * + *

    + *
  • Turn: {@code POST url} with body {@code {"data":,"init":}} and + * {@code Accept: text/event-stream}. SSE frames are {@code data: {"message":}} then + * {@code data: {"result":}}. Error frames are {@code data: {"error":{...}}}. + *
  • getSnapshot: {@code POST url/getSnapshot} with body {@code {"data":}} → {@code + * {"result":}}. + *
  • abort: {@code POST url/abort} with body {@code {"data":{"snapshotId":...}}} → {@code + * {"result":{"status":...}}}. + *
+ * + * @param the type of custom session state + */ +public final class HttpAgentTransport implements AgentTransport { + + private static final String SSE_DATA_PREFIX = "data: "; + + private final RemoteAgentOptions opts; + private final ObjectMapper mapper; + private final HttpClient httpClient; + + /** + * Constructs an {@link HttpAgentTransport}. + * + * @param opts the options specifying the endpoint URLs, headers, and serverManaged flag + */ + public HttpAgentTransport(RemoteAgentOptions opts) { + this.opts = opts; + this.mapper = JsonUtils.getObjectMapper(); + this.httpClient = HttpClient.newHttpClient(); + } + + @Override + public AgentOutput runTurn( + AgentInput input, AgentInit init, Consumer onChunk) { + try { + // Build request envelope: {"data": , "init": } + ObjectNode envelope = mapper.createObjectNode(); + envelope.set("data", mapper.valueToTree(input != null ? input : new AgentInput())); + if (init != null) { + envelope.set("init", mapper.valueToTree(init)); + } + String body = mapper.writeValueAsString(envelope); + + HttpRequest.Builder reqBuilder = + HttpRequest.newBuilder() + .uri(URI.create(opts.url())) + .header("Content-Type", "application/json") + .header("Accept", "text/event-stream") + .POST(HttpRequest.BodyPublishers.ofString(body, StandardCharsets.UTF_8)); + + for (Map.Entry entry : opts.headers().entrySet()) { + reqBuilder.header(entry.getKey(), entry.getValue()); + } + + HttpResponse response = + httpClient.send(reqBuilder.build(), HttpResponse.BodyHandlers.ofInputStream()); + + if (response.statusCode() >= 400) { + String errorBody = new String(response.body().readAllBytes(), StandardCharsets.UTF_8); + throw new GenkitException("Agent HTTP error " + response.statusCode() + ": " + errorBody); + } + + // Parse SSE stream + String contentType = response.headers().firstValue("Content-Type").orElse(""); + if (contentType.contains("text/event-stream")) { + return parseSseStream(response.body(), onChunk); + } else { + // Non-streaming response: {"result": } + String responseBody = new String(response.body().readAllBytes(), StandardCharsets.UTF_8); + JsonNode root = mapper.readTree(responseBody); + return deserializeOutput(root.get("result")); + } + } catch (GenkitException e) { + throw e; + } catch (Exception e) { + throw new GenkitException("HttpAgentTransport.runTurn failed", e); + } + } + + /** + * Parses an SSE stream, dispatching {@code data: {"message":...}} frames to {@code onChunk} and + * returning the {@link AgentOutput} from the terminal {@code data: {"result":...}} frame. + */ + private AgentOutput parseSseStream( + java.io.InputStream inputStream, Consumer onChunk) throws Exception { + try (BufferedReader reader = + new BufferedReader(new InputStreamReader(inputStream, StandardCharsets.UTF_8))) { + String line; + while ((line = reader.readLine()) != null) { + if (!line.startsWith(SSE_DATA_PREFIX)) { + continue; + } + String json = line.substring(SSE_DATA_PREFIX.length()).trim(); + if (json.isEmpty()) { + continue; + } + JsonNode frame = mapper.readTree(json); + + if (frame.has("error")) { + JsonNode err = frame.get("error"); + String msg = err.path("message").asText("agent error"); + throw new GenkitException(msg); + } + + if (frame.has("result")) { + return deserializeOutput(frame.get("result")); + } + + if (frame.has("message") && onChunk != null) { + JsonNode chunkNode = frame.get("message"); + AgentStreamChunk chunk = mapper.treeToValue(chunkNode, AgentStreamChunk.class); + onChunk.accept(chunk); + } + } + } + throw new GenkitException("SSE stream ended without a result frame"); + } + + @Override + @SuppressWarnings("unchecked") + public SessionSnapshot getSnapshot(GetSnapshotRequest req) { + try { + ObjectNode envelope = mapper.createObjectNode(); + envelope.set("data", mapper.valueToTree(req)); + String body = mapper.writeValueAsString(envelope); + + HttpRequest.Builder reqBuilder = + HttpRequest.newBuilder() + .uri(URI.create(opts.getSnapshotUrl())) + .header("Content-Type", "application/json") + .POST(HttpRequest.BodyPublishers.ofString(body, StandardCharsets.UTF_8)); + + for (Map.Entry entry : opts.headers().entrySet()) { + reqBuilder.header(entry.getKey(), entry.getValue()); + } + + HttpResponse response = + httpClient.send(reqBuilder.build(), HttpResponse.BodyHandlers.ofString()); + + if (response.statusCode() >= 400) { + throw new GenkitException( + "getSnapshot HTTP error " + response.statusCode() + ": " + response.body()); + } + + JsonNode root = mapper.readTree(response.body()); + JsonNode result = root.get("result"); + if (result == null || result.isNull()) { + return null; + } + return (SessionSnapshot) mapper.treeToValue(result, SessionSnapshot.class); + } catch (GenkitException e) { + throw e; + } catch (Exception e) { + throw new GenkitException("HttpAgentTransport.getSnapshot failed", e); + } + } + + @Override + public SnapshotStatus abort(String snapshotId) { + try { + ObjectNode data = mapper.createObjectNode(); + data.put("snapshotId", snapshotId); + ObjectNode envelope = mapper.createObjectNode(); + envelope.set("data", data); + String body = mapper.writeValueAsString(envelope); + + HttpRequest.Builder reqBuilder = + HttpRequest.newBuilder() + .uri(URI.create(opts.abortUrl())) + .header("Content-Type", "application/json") + .POST(HttpRequest.BodyPublishers.ofString(body, StandardCharsets.UTF_8)); + + for (Map.Entry entry : opts.headers().entrySet()) { + reqBuilder.header(entry.getKey(), entry.getValue()); + } + + HttpResponse response = + httpClient.send(reqBuilder.build(), HttpResponse.BodyHandlers.ofString()); + + if (response.statusCode() >= 400) { + throw new GenkitException( + "abort HTTP error " + response.statusCode() + ": " + response.body()); + } + + JsonNode root = mapper.readTree(response.body()); + JsonNode result = root.get("result"); + if (result == null || result.isNull()) { + return null; + } + String statusStr = result.path("status").asText(null); + if (statusStr == null || statusStr.isEmpty()) { + return null; + } + try { + return SnapshotStatus.fromValue(statusStr); + } catch (IllegalArgumentException e) { + return null; + } + } catch (GenkitException e) { + throw e; + } catch (Exception e) { + throw new GenkitException("HttpAgentTransport.abort failed", e); + } + } + + @Override + public boolean serverManaged() { + return opts.serverManaged(); + } + + @SuppressWarnings("unchecked") + private AgentOutput deserializeOutput(JsonNode node) throws Exception { + if (node == null || node.isNull()) { + return new AgentOutput<>(); + } + return (AgentOutput) mapper.treeToValue(node, AgentOutput.class); + } +} diff --git a/genkit/src/main/java/com/google/genkit/client/RemoteAgent.java b/genkit/src/main/java/com/google/genkit/client/RemoteAgent.java new file mode 100644 index 000000000..df12105ca --- /dev/null +++ b/genkit/src/main/java/com/google/genkit/client/RemoteAgent.java @@ -0,0 +1,52 @@ +/* + * 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. + * + * SPDX-License-Identifier: Apache-2.0 + */ + +package com.google.genkit.client; + +import com.google.genkit.ai.agent.AgentChat; + +/** + * Factory for creating a remote agent chat client. + * + *

Use {@link #chat(RemoteAgentOptions)} to connect to an agent served by the Jetty plugin (or + * any compatible server that speaks the Genkit agent wire format): + * + * 

{@code
+ * AgentChat> chat = RemoteAgent.>chat(
+ *     RemoteAgentOptions.builder()
+ *         .url("http://localhost:8080/myAgent")
+ *         .build());
+ * AgentResponse> resp = chat.send("hello");
+ * }
+ */ +public final class RemoteAgent { + + private RemoteAgent() {} + + /** + * Creates a new {@link AgentChat} backed by an {@link HttpAgentTransport} that speaks to the + * agent at {@code opts.url()}. + * + * @param the type of custom session state + * @param opts the remote agent options (URL, headers, serverManaged flag) + * @return a fresh {@link AgentChat} ready to send turns + */ + public static AgentChat chat(RemoteAgentOptions opts) { + return AgentChat.over(new HttpAgentTransport(opts), null); + } +} diff --git a/genkit/src/main/java/com/google/genkit/client/RemoteAgentOptions.java b/genkit/src/main/java/com/google/genkit/client/RemoteAgentOptions.java new file mode 100644 index 000000000..d2ed3e651 --- /dev/null +++ b/genkit/src/main/java/com/google/genkit/client/RemoteAgentOptions.java @@ -0,0 +1,192 @@ +/* + * 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. + * + * SPDX-License-Identifier: Apache-2.0 + */ + +package com.google.genkit.client; + +import java.util.Collections; +import java.util.HashMap; +import java.util.Map; + +/** + * Options for configuring a {@link RemoteAgent} HTTP client. + * + *

Build with {@link #builder()}. + * + *

    + *
  • {@link #url()} — the agent turn endpoint (e.g. {@code http://host:8080/myAgent}). + *
  • {@link #getSnapshotUrl()} — companion snapshot endpoint; defaults to {@code url + + * "/getSnapshot"}. + *
  • {@link #abortUrl()} — companion abort endpoint; defaults to {@code url + "/abort"}. + *
  • {@link #headers()} — optional extra request headers. + *
  • {@link #serverManaged()} — whether state is server-managed; default {@code true}. + *
+ */ +public final class RemoteAgentOptions { + + private final String url; + private final String getSnapshotUrl; + private final String abortUrl; + private final Map headers; + private final boolean serverManaged; + + private RemoteAgentOptions(Builder builder) { + this.url = builder.url; + this.getSnapshotUrl = + builder.getSnapshotUrl != null ? builder.getSnapshotUrl : builder.url + "/getSnapshot"; + this.abortUrl = builder.abortUrl != null ? builder.abortUrl : builder.url + "/abort"; + this.headers = + builder.headers != null + ? Collections.unmodifiableMap(new HashMap<>(builder.headers)) + : Collections.emptyMap(); + this.serverManaged = builder.serverManaged; + } + + /** + * Creates a builder for {@link RemoteAgentOptions}. + * + * @return a new builder + */ + public static Builder builder() { + return new Builder(); + } + + /** + * Returns the agent turn endpoint URL. + * + * @return the URL + */ + public String url() { + return url; + } + + /** + * Returns the companion getSnapshot URL (defaults to {@code url + "/getSnapshot"}). + * + * @return the getSnapshot URL + */ + public String getSnapshotUrl() { + return getSnapshotUrl; + } + + /** + * Returns the companion abort URL (defaults to {@code url + "/abort"}). + * + * @return the abort URL + */ + public String abortUrl() { + return abortUrl; + } + + /** + * Returns extra HTTP request headers sent on every request. + * + * @return an unmodifiable map of headers (never null) + */ + public Map headers() { + return headers; + } + + /** + * Returns whether the agent is server-managed (default {@code true}). + * + * @return {@code true} if server-managed + */ + public boolean serverManaged() { + return serverManaged; + } + + /** Builder for {@link RemoteAgentOptions}. */ + public static final class Builder { + + private String url; + private String getSnapshotUrl; + private String abortUrl; + private Map headers; + private boolean serverManaged = true; + + private Builder() {} + + /** + * Sets the agent turn endpoint URL (required). + * + * @param url the URL + * @return this builder + */ + public Builder url(String url) { + this.url = url; + return this; + } + + /** + * Overrides the getSnapshot URL. Defaults to {@code url + "/getSnapshot"}. + * + * @param getSnapshotUrl the URL + * @return this builder + */ + public Builder getSnapshotUrl(String getSnapshotUrl) { + this.getSnapshotUrl = getSnapshotUrl; + return this; + } + + /** + * Overrides the abort URL. Defaults to {@code url + "/abort"}. + * + * @param abortUrl the URL + * @return this builder + */ + public Builder abortUrl(String abortUrl) { + this.abortUrl = abortUrl; + return this; + } + + /** + * Sets extra HTTP request headers. + * + * @param headers headers to include on every request + * @return this builder + */ + public Builder headers(Map headers) { + this.headers = headers; + return this; + } + + /** + * Sets whether the agent is server-managed. Default {@code true}. + * + * @param serverManaged true for server-managed + * @return this builder + */ + public Builder serverManaged(boolean serverManaged) { + this.serverManaged = serverManaged; + return this; + } + + /** + * Builds the {@link RemoteAgentOptions}. + * + * @return a new options instance + * @throws IllegalStateException if {@code url} is null + */ + public RemoteAgentOptions build() { + if (url == null || url.isBlank()) { + throw new IllegalStateException("url is required"); + } + return new RemoteAgentOptions(this); + } + } +} diff --git a/genkit/src/test/java/com/google/genkit/ExecutionContextTest.java b/genkit/src/test/java/com/google/genkit/ExecutionContextTest.java new file mode 100644 index 000000000..ed0c61765 --- /dev/null +++ b/genkit/src/test/java/com/google/genkit/ExecutionContextTest.java @@ -0,0 +1,247 @@ +/* + * 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. + * + * SPDX-License-Identifier: Apache-2.0 + */ + +package com.google.genkit; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertNotNull; + +import com.fasterxml.jackson.databind.JsonNode; +import com.google.genkit.agent.AgentConfig; +import com.google.genkit.ai.Candidate; +import com.google.genkit.ai.FinishReason; +import com.google.genkit.ai.Message; +import com.google.genkit.ai.Model; +import com.google.genkit.ai.ModelInfo; +import com.google.genkit.ai.ModelRequest; +import com.google.genkit.ai.ModelResponse; +import com.google.genkit.ai.ModelResponseChunk; +import com.google.genkit.ai.Part; +import com.google.genkit.ai.Role; +import com.google.genkit.ai.Tool; +import com.google.genkit.ai.ToolRequest; +import com.google.genkit.ai.ToolResponse; +import com.google.genkit.ai.agent.Agent; +import com.google.genkit.ai.agent.AgentInput; +import com.google.genkit.ai.agent.AgentResult; +import com.google.genkit.ai.agent.CustomAgentConfig; +import com.google.genkit.core.ActionContext; +import com.google.genkit.core.ActionRunResult; +import com.google.genkit.core.BufferedInputSource; +import com.google.genkit.core.JsonUtils; +import java.util.List; +import java.util.Map; +import java.util.function.Consumer; +import org.junit.jupiter.api.Test; + +/** + * End-to-end test proving the full execution-context propagation chain that powers the Dev UI + * "Execution context" panel: + * + * 
+ * reflection request body context
+ *   → ActionContext.context
+ *   → (run path) withSpanContext keeps it
+ *   → AgentFnContext.context()
+ *   → GenkitBeta.buildAgentFn → GenerateOptions.context
+ *   → Genkit.generate's tool-executing ActionContext
+ *   → Tool handler reads ctx.getContext()
+ *
+ */ +class ExecutionContextTest { + + private static final Map AUTH_CONTEXT = Map.of("auth", Map.of("user", "alice")); + + // ── helpers ────────────────────────────────────────────────────────────────── + + private static Genkit experimentalGenkit() { + return new Genkit(GenkitOptions.builder().experimental(true).build()); + } + + private static JsonNode initJson() { + return JsonUtils.toJsonNode(new com.google.genkit.ai.agent.AgentInit>()); + } + + private static BufferedInputSource inputSourceWith(String userText) { + BufferedInputSource src = new BufferedInputSource<>(); + src.offer(JsonUtils.toJsonNode(AgentInput.builder().message(Message.user(userText)).build())); + src.end(); + return src; + } + + /** + * A fake model that, on the first turn, requests the {@code whoami} tool; on the next turn (once + * a TOOL message carrying the tool's output is present in the history) it echoes the tool output + * as its final text reply. This drives the real generate tool-execution loop so we can prove the + * tool observes the injected execution context. + */ + private static Model toolCallingModel(String name, String toolName) { + return new Model() { + @Override + public String getName() { + return name; + } + + @Override + public ModelInfo getInfo() { + return new ModelInfo(); + } + + @Override + public boolean supportsStreaming() { + return true; + } + + @Override + public ModelResponse run(ActionContext ctx, ModelRequest request) { + return run(ctx, request, null); + } + + @Override + public ModelResponse run( + ActionContext ctx, ModelRequest request, Consumer streamCallback) { + // Look for a TOOL message carrying the tool output from a previous turn. + String toolOutput = findToolOutput(request); + if (toolOutput != null) { + String reply = "tool said: " + toolOutput; + if (streamCallback != null) { + ModelResponseChunk chunk = new ModelResponseChunk(); + chunk.setContent(List.of(Part.text(reply))); + streamCallback.accept(chunk); + } + Candidate candidate = new Candidate(Message.model(reply), FinishReason.STOP); + ModelResponse response = new ModelResponse(List.of(candidate)); + response.setFinishReason(FinishReason.STOP); + response.setRequest(request); + return response; + } + + // First turn: ask to call the tool. + Part toolRequestPart = Part.toolRequest(new ToolRequest(toolName, Map.of())); + Message assistant = new Message(Role.MODEL, List.of(toolRequestPart)); + Candidate candidate = new Candidate(assistant, FinishReason.STOP); + ModelResponse response = new ModelResponse(List.of(candidate)); + response.setFinishReason(FinishReason.STOP); + response.setRequest(request); + return response; + } + + private String findToolOutput(ModelRequest request) { + if (request.getMessages() == null) { + return null; + } + for (Message m : request.getMessages()) { + if (m.getRole() == Role.TOOL && m.getContent() != null) { + for (Part p : m.getContent()) { + ToolResponse tr = p.getToolResponse(); + if (tr != null && tr.getOutput() != null) { + return String.valueOf(tr.getOutput()); + } + } + } + } + return null; + } + }; + } + + // ── Layer 1: reflection ActionContext → AgentFnContext (custom agent) ───────── + + @Test + void customAgentReadsContextFromAgentFnContext() throws Exception { + Genkit genkit = experimentalGenkit(); + + CustomAgentConfig> config = + CustomAgentConfig.>builder().name("ctxAgent").build(); + + // The AgentFn reads the injected execution context off its AgentFnContext and returns the + // nested auth.user value as the turn's message. + Agent> agent = + genkit + .beta() + .defineCustomAgent( + config, + (runner, ctx) -> { + assertNotNull(ctx.context(), "AgentFnContext must carry the run ActionContext"); + Map userContext = ctx.context().getContext(); + assertNotNull(userContext, "user context must propagate to the agent"); + @SuppressWarnings("unchecked") + Map auth = (Map) userContext.get("auth"); + return AgentResult.builder() + .message(Message.model(String.valueOf(auth.get("user")))) + .build(); + }); + + // Simulate the reflection server: build an ActionContext WITH the user context, then drive the + // bidi action through runBidiJsonWithTelemetry (which internally calls withSpanContext). + ActionContext ctx = + ActionContext.builder().registry(genkit.getRegistry()).context(AUTH_CONTEXT).build(); + + ActionRunResult result = + agent.runBidiJsonWithTelemetry(ctx, initJson(), inputSourceWith("hi"), null); + + JsonNode out = result.getResult(); + assertEquals("alice", out.get("message").get("content").get(0).get("text").asText()); + } + + // ── Layer 2: full chain → generate → tool reads ctx.getContext() ────────────── + + @Test + void toolObservesExecutionContextDuringAgentGenerate() throws Exception { + Genkit genkit = experimentalGenkit(); + genkit.registerModel(toolCallingModel("toolModel", "whoami")); + + // Tool handler reads the execution context off its ActionContext. + Tool, String> whoami = + genkit.defineTool( + "whoami", + "returns the authenticated user from the execution context", + (ActionContext ctx, Map input) -> { + Map userContext = ctx.getContext(); + if (userContext == null) { + return "NO_CONTEXT"; + } + @SuppressWarnings("unchecked") + Map auth = (Map) userContext.get("auth"); + return auth == null ? "NO_AUTH" : String.valueOf(auth.get("user")); + }, + (Class>) (Class) Map.class, + String.class); + + AgentConfig> config = + AgentConfig.>builder() + .name("toolAgent") + .model("toolModel") + .system("You are helpful.") + .tools(List.of(whoami)) + .build(); + + Agent> agent = genkit.beta().defineAgent(config); + + ActionContext ctx = + ActionContext.builder().registry(genkit.getRegistry()).context(AUTH_CONTEXT).build(); + + ActionRunResult result = + agent.runBidiJsonWithTelemetry(ctx, initJson(), inputSourceWith("who am I?"), null); + + JsonNode out = result.getResult(); + // The model echoes the tool output; the tool output is the auth.user from the execution + // context. Proves the full chain delivered "alice" all the way into the tool handler. + assertEquals("tool said: alice", out.get("message").get("content").get(0).get("text").asText()); + } +} diff --git a/genkit/src/test/java/com/google/genkit/GenkitBetaTest.java b/genkit/src/test/java/com/google/genkit/GenkitBetaTest.java new file mode 100644 index 000000000..08ea0cdd2 --- /dev/null +++ b/genkit/src/test/java/com/google/genkit/GenkitBetaTest.java @@ -0,0 +1,576 @@ +/* + * 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. + * + * SPDX-License-Identifier: Apache-2.0 + */ + +package com.google.genkit; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import com.fasterxml.jackson.databind.JsonNode; +import com.google.genkit.agent.AgentConfig; +import com.google.genkit.ai.Candidate; +import com.google.genkit.ai.FinishReason; +import com.google.genkit.ai.Message; +import com.google.genkit.ai.Model; +import com.google.genkit.ai.ModelInfo; +import com.google.genkit.ai.ModelRequest; +import com.google.genkit.ai.ModelResponse; +import com.google.genkit.ai.ModelResponseChunk; +import com.google.genkit.ai.Part; +import com.google.genkit.ai.Role; +import com.google.genkit.ai.agent.Agent; +import com.google.genkit.ai.agent.AgentFinishReason; +import com.google.genkit.ai.agent.AgentInput; +import com.google.genkit.ai.agent.AgentRef; +import com.google.genkit.ai.agent.AgentResult; +import com.google.genkit.ai.agent.InMemorySessionStore; +import com.google.genkit.core.ActionContext; +import com.google.genkit.core.BufferedInputSource; +import com.google.genkit.core.GenkitException; +import com.google.genkit.core.JsonUtils; +import java.util.ArrayList; +import java.util.List; +import java.util.Map; +import java.util.function.Consumer; +import org.junit.jupiter.api.Test; + +/** TDD tests for {@link GenkitBeta} (Task 5.3). */ +class GenkitBetaTest { + + // ── helpers ────────────────────────────────────────────────────────────────── + + /** A fake echo model that streams one chunk and returns "echo: ". */ + private static Model echoModel(String name) { + return new Model() { + @Override + public String getName() { + return name; + } + + @Override + public ModelInfo getInfo() { + return new ModelInfo(); + } + + @Override + public boolean supportsStreaming() { + return true; + } + + @Override + public ModelResponse run(ActionContext ctx, ModelRequest request) { + return run(ctx, request, null); + } + + @Override + public ModelResponse run( + ActionContext ctx, ModelRequest request, Consumer streamCallback) { + String userText = ""; + if (request.getMessages() != null) { + for (int i = request.getMessages().size() - 1; i >= 0; i--) { + Message m = request.getMessages().get(i); + if (m.getRole() == Role.USER) { + userText = m.getText(); + break; + } + } + } + String reply = "echo: " + userText; + if (streamCallback != null) { + ModelResponseChunk chunk = new ModelResponseChunk(); + chunk.setContent(List.of(Part.text(reply))); + streamCallback.accept(chunk); + } + Candidate candidate = new Candidate(Message.model(reply), FinishReason.STOP); + ModelResponse response = new ModelResponse(List.of(candidate)); + response.setFinishReason(FinishReason.STOP); + response.setRequest(request); + return response; + } + }; + } + + private static Genkit experimentalGenkit() { + Genkit genkit = new Genkit(GenkitOptions.builder().experimental(true).build()); + genkit.registerModel(echoModel("echoModel")); + return genkit; + } + + /** + * A model that records the SYSTEM-role text of the last request it received (so tests can assert + * exactly what system prompt reached the model), then replies "ok". + */ + private static final class RecordingModel implements Model { + private final String name; + private volatile String lastSystem; + + RecordingModel(String name) { + this.name = name; + } + + @Override + public String getName() { + return name; + } + + @Override + public ModelInfo getInfo() { + ModelInfo info = new ModelInfo(); + ModelInfo.ModelCapabilities caps = new ModelInfo.ModelCapabilities(); + caps.setSystemRole(true); + caps.setMultiturn(true); + info.setSupports(caps); + return info; + } + + @Override + public boolean supportsStreaming() { + return true; + } + + @Override + public ModelResponse run(ActionContext ctx, ModelRequest request) { + return run(ctx, request, null); + } + + @Override + public ModelResponse run( + ActionContext ctx, ModelRequest request, Consumer streamCallback) { + if (request.getMessages() != null) { + for (Message m : request.getMessages()) { + if (m.getRole() == Role.SYSTEM) { + lastSystem = m.getText(); + break; + } + } + } + if (streamCallback != null) { + ModelResponseChunk chunk = new ModelResponseChunk(); + chunk.setContent(List.of(Part.text("ok"))); + streamCallback.accept(chunk); + } + Candidate candidate = new Candidate(Message.model("ok"), FinishReason.STOP); + ModelResponse response = new ModelResponse(List.of(candidate)); + response.setFinishReason(FinishReason.STOP); + response.setRequest(request); + return response; + } + } + + private static JsonNode initJson() { + return JsonUtils.toJsonNode(new com.google.genkit.ai.agent.AgentInit>()); + } + + private static BufferedInputSource inputSourceWith(String userText) { + BufferedInputSource src = new BufferedInputSource<>(); + src.offer(JsonUtils.toJsonNode(AgentInput.builder().message(Message.user(userText)).build())); + src.end(); + return src; + } + + // ── gating: defineAgent ────────────────────────────────────────────────────── + + @Test + void defineAgentWithoutExperimentalThrows() { + Genkit genkit = new Genkit(GenkitOptions.builder().build()); + AgentConfig> config = + AgentConfig.>builder().name("a1").build(); + + GenkitException ex = + assertThrows(GenkitException.class, () -> genkit.beta().defineAgent(config)); + assertTrue( + ex.getMessage().toLowerCase().contains("experimental"), + "message should mention experimental, got: " + ex.getMessage()); + } + + // ── gating: defineCustomAgent ──────────────────────────────────────────────── + + @Test + void defineCustomAgentWithoutExperimentalThrows() { + Genkit genkit = new Genkit(GenkitOptions.builder().build()); + com.google.genkit.ai.agent.CustomAgentConfig> config = + com.google.genkit.ai.agent.CustomAgentConfig.>builder() + .name("custom1") + .build(); + + GenkitException ex = + assertThrows( + GenkitException.class, + () -> + genkit + .beta() + .defineCustomAgent( + config, + (runner, ctx) -> + AgentResult.builder().message(Message.model("x")).build())); + assertTrue(ex.getMessage().toLowerCase().contains("experimental")); + } + + @Test + void defineCustomAgentWithExperimentalRegisters() { + Genkit genkit = experimentalGenkit(); + com.google.genkit.ai.agent.CustomAgentConfig> config = + com.google.genkit.ai.agent.CustomAgentConfig.>builder() + .name("custom2") + .build(); + + Agent> agent = + genkit + .beta() + .defineCustomAgent( + config, (runner, ctx) -> AgentResult.builder().message(Message.model("x")).build()); + + assertNotNull(agent); + assertNotNull(genkit.getRegistry().lookupAction("/agent/custom2")); + } + + // ── defineAgent registration ───────────────────────────────────────────────── + + @Test + void defineAgentWithExperimentalRegisters() { + Genkit genkit = experimentalGenkit(); + AgentConfig> config = + AgentConfig.>builder() + .name("helper") + .description("a helper agent") + .model("echoModel") + .system("You are helpful.") + .build(); + + Agent> agent = genkit.beta().defineAgent(config); + + assertNotNull(agent); + assertNotNull(genkit.getRegistry().lookupAction("/agent/helper")); + } + + // ── one turn through the agent bidi action ─────────────────────────────────── + + @Test + void defineAgentRunsOneTurn() throws Exception { + Genkit genkit = experimentalGenkit(); + AgentConfig> config = + AgentConfig.>builder() + .name("turnAgent") + .model("echoModel") + .system("You are helpful.") + .build(); + + Agent> agent = genkit.beta().defineAgent(config); + + ActionContext ctx = new ActionContext(genkit.getRegistry()); + List chunks = new ArrayList<>(); + JsonNode out = agent.runBidiJson(ctx, initJson(), inputSourceWith("hi"), chunks::add); + + assertNotNull(out); + // The final message is the model's reply. + assertEquals("echo: hi", out.get("message").get("content").get(0).get("text").asText()); + + // A turnEnd chunk was observed, and a model chunk was streamed. + boolean sawTurnEnd = chunks.stream().anyMatch(c -> c.has("turnEnd")); + assertTrue(sawTurnEnd, "expected a turnEnd chunk"); + boolean sawModelChunk = chunks.stream().anyMatch(c -> c.has("modelChunk")); + assertTrue(sawModelChunk, "expected a modelChunk"); + } + + // ── definePromptAgent (first-cut: uses config.system like defineAgent) ──────── + + @Test + void definePromptAgentWithoutExperimentalThrows() { + Genkit genkit = new Genkit(GenkitOptions.builder().build()); + AgentConfig> config = + AgentConfig.>builder().name("p1").build(); + + GenkitException ex = + assertThrows(GenkitException.class, () -> genkit.beta().definePromptAgent(config)); + assertTrue(ex.getMessage().toLowerCase().contains("experimental")); + } + + @Test + void definePromptAgentWithExperimentalRegistersAndRuns() throws Exception { + Genkit genkit = experimentalGenkit(); + AgentConfig> config = + AgentConfig.>builder() + .name("promptAgent") + .model("echoModel") + .system("You are helpful.") + .build(); + + Agent> agent = genkit.beta().definePromptAgent(config); + assertNotNull(genkit.getRegistry().lookupAction("/agent/promptAgent")); + + ActionContext ctx = new ActionContext(genkit.getRegistry()); + JsonNode out = agent.runBidiJson(ctx, initJson(), inputSourceWith("yo"), c -> {}); + assertEquals("echo: yo", out.get("message").get("content").get(0).get("text").asText()); + } + + // ── definePromptAgent renders its template with promptInput PER TURN ────────── + + /** + * Proves the prompt-backed agent renders its Handlebars template with {@code promptInput} on each + * turn: the {@code topicAgent.prompt} template is {@code "You are an expert on {{topic}}. Answer + * questions about {{topic}} for {{userName}}."}. With {@code promptInput = {topic: "wombats", + * userName: "Ada"}}, the SYSTEM text that actually reaches the model must contain the + * interpolated values, and must NOT contain the raw {@code {{topic}}} placeholder. + */ + @Test + void definePromptAgentRendersPromptInputPerTurn() throws Exception { + Genkit genkit = new Genkit(GenkitOptions.builder().experimental(true).build()); + RecordingModel model = new RecordingModel("echoModel"); + genkit.registerModel(model); + + Map promptInput = new java.util.HashMap<>(); + promptInput.put("topic", "wombats"); + promptInput.put("userName", "Ada"); + + AgentConfig> config = + AgentConfig.>builder() + .name("topicAgent") + .model("echoModel") + .promptName("topicAgent") + .promptInput(promptInput) + .build(); + + Agent> agent = genkit.beta().definePromptAgent(config); + + ActionContext ctx = new ActionContext(genkit.getRegistry()); + agent.runBidiJson(ctx, initJson(), inputSourceWith("hi"), c -> {}); + + assertNotNull(model.lastSystem, "the model should have received a SYSTEM message"); + assertTrue( + model.lastSystem.contains("You are an expert on wombats"), + "system should interpolate promptInput.topic, got: " + model.lastSystem); + assertTrue( + model.lastSystem.contains("for Ada"), + "system should interpolate promptInput.userName, got: " + model.lastSystem); + assertFalse( + model.lastSystem.contains("{{"), + "system must be RENDERED, not the raw template, got: " + model.lastSystem); + } + + /** + * Proves the render context also folds in the current session state: with an empty {@code + * promptInput} but a session custom-state map carrying {@code {topic: "quokkas", userName: + * "Grace"}}, those state values interpolate into the rendered system prompt for the turn. + */ + @Test + void definePromptAgentRendersSessionStatePerTurn() throws Exception { + Genkit genkit = new Genkit(GenkitOptions.builder().experimental(true).build()); + RecordingModel model = new RecordingModel("echoModel"); + genkit.registerModel(model); + + AgentConfig> config = + AgentConfig.>builder() + .name("topicAgent") + .model("echoModel") + .promptName("topicAgent") + .build(); + + Agent> agent = genkit.beta().definePromptAgent(config); + + Map customState = new java.util.HashMap<>(); + customState.put("topic", "quokkas"); + customState.put("userName", "Grace"); + com.google.genkit.ai.agent.SessionState> state = + com.google.genkit.ai.agent.SessionState.>builder() + .custom(customState) + .build(); + JsonNode initWithState = + JsonUtils.toJsonNode( + com.google.genkit.ai.agent.AgentInit.>builder() + .state(state) + .build()); + + ActionContext ctx = new ActionContext(genkit.getRegistry()); + agent.runBidiJson(ctx, initWithState, inputSourceWith("hi"), c -> {}); + + assertNotNull(model.lastSystem, "the model should have received a SYSTEM message"); + assertTrue( + model.lastSystem.contains("You are an expert on quokkas"), + "system should interpolate session state topic, got: " + model.lastSystem); + assertTrue( + model.lastSystem.contains("for Grace"), + "system should interpolate session state userName, got: " + model.lastSystem); + } + + // ── finish reason mapping ──────────────────────────────────────────────────── + + @Test + void agentTurnUsesStopFinishReason() throws Exception { + Genkit genkit = experimentalGenkit(); + AgentConfig> config = + AgentConfig.>builder().name("frAgent").model("echoModel").build(); + Agent> agent = genkit.beta().defineAgent(config); + + ActionContext ctx = new ActionContext(genkit.getRegistry()); + JsonNode out = agent.runBidiJson(ctx, initJson(), inputSourceWith("hi"), c -> {}); + assertEquals(AgentFinishReason.STOP.getValue(), out.get("finishReason").asText()); + assertFalse(out.get("message").isNull()); + } + + // ── isExperimental flag exposed on Genkit ──────────────────────────────────── + + @Test + void isExperimentalReflectsOptions() { + assertFalse(new Genkit(GenkitOptions.builder().build()).isExperimental()); + assertTrue(new Genkit(GenkitOptions.builder().experimental(true).build()).isExperimental()); + } + + // ── description exposed on the agent's ref() ───────────────────────────────── + + @Test + void defineAgentExposesDescriptionOnRef() { + Genkit genkit = experimentalGenkit(); + AgentConfig> config = + AgentConfig.>builder() + .name("describedAgent") + .description("some text") + .model("echoModel") + .build(); + + Agent> agent = genkit.beta().defineAgent(config); + + AgentRef ref = agent.ref(); + assertEquals("describedAgent", ref.getName()); + assertEquals("some text", ref.getDescription()); + } + + // ── metadata: description + stateManagement/abortable ──────────────────────── + + /** + * Client-managed (no store): {@code defineAgent} should surface the description and a {@code + * stateManagement=client} / {@code abortable=false} sub-map under the {@code "agent"} metadata + * key. This config does not set a {@code stateType}, so no {@code stateSchema} is expected either + * (see {@link #defineAgentMetadataIncludesStateSchemaForTypedState} for the populated case). + */ + @Test + @SuppressWarnings("unchecked") + void defineAgentMetadataReflectsDescriptionAndClientManagedState() { + Genkit genkit = experimentalGenkit(); + AgentConfig> config = + AgentConfig.>builder() + .name("metaAgentClient") + .description("client managed agent") + .model("echoModel") + .build(); + + Agent> agent = genkit.beta().defineAgent(config); + + Map metadata = agent.getMetadata(); + assertNotNull(metadata); + assertEquals("client managed agent", metadata.get("description")); + + Map agentMeta = (Map) metadata.get("agent"); + assertNotNull(agentMeta, "expected an \"agent\" sub-map in metadata"); + assertEquals("client", agentMeta.get("stateManagement")); + assertEquals(false, agentMeta.get("abortable")); + assertFalse( + agentMeta.containsKey("stateSchema"), + "no stateType was configured, so no stateSchema should be generated"); + + assertFalse(agent.serverManaged()); + } + + /** + * Server-managed (store configured with an {@link InMemorySessionStore}, which also implements + * {@code SnapshotSubscriber}): {@code stateManagement} flips to {@code "server"} and {@code + * abortable} flips to {@code true}. + */ + @Test + @SuppressWarnings("unchecked") + void defineAgentMetadataReflectsServerManagedState() { + Genkit genkit = experimentalGenkit(); + AgentConfig> config = + AgentConfig.>builder() + .name("metaAgentServer") + .description("server managed agent") + .model("echoModel") + .store(new InMemorySessionStore<>()) + .build(); + + Agent> agent = genkit.beta().defineAgent(config); + + Map metadata = agent.getMetadata(); + Map agentMeta = (Map) metadata.get("agent"); + assertNotNull(agentMeta); + assertEquals("server", agentMeta.get("stateManagement")); + assertEquals(true, agentMeta.get("abortable")); + + assertTrue(agent.serverManaged()); + } + + /** POJO custom-state type used to prove {@code stateSchema} generation end-to-end. */ + public static class TypedState { + private String status; + private int counter; + + public String getStatus() { + return status; + } + + public void setStatus(String status) { + this.status = status; + } + + public int getCounter() { + return counter; + } + + public void setCounter(int counter) { + this.counter = counter; + } + } + + /** + * Fix 4: when a {@code defineAgent} config specifies a non-trivial POJO {@code stateType}, the + * registered agent's metadata must include a generated {@code stateSchema} (under the {@code + * "agent"} sub-map, sibling to {@code stateManagement}/{@code abortable}) describing that type's + * properties. {@code defineAgent} delegates to {@code AgentActions.defineCustomAgent} under the + * hood (see {@code GenkitBeta.defineGenerateBackedAgent}), so this also exercises the + * generate-backed path, not just the low-level {@code defineCustomAgent} entry point. + */ + @Test + @SuppressWarnings("unchecked") + void defineAgentMetadataIncludesStateSchemaForTypedState() { + Genkit genkit = experimentalGenkit(); + AgentConfig config = + AgentConfig.builder() + .name("typedStateAgent") + .model("echoModel") + .stateType(TypedState.class) + .build(); + + Agent agent = genkit.beta().defineAgent(config); + + Map metadata = agent.getMetadata(); + Map agentMeta = (Map) metadata.get("agent"); + assertNotNull(agentMeta); + + Object stateSchemaObj = agentMeta.get("stateSchema"); + assertNotNull(stateSchemaObj, "expected a generated stateSchema for a typed custom state"); + Map stateSchema = (Map) stateSchemaObj; + + Object propertiesObj = stateSchema.get("properties"); + assertNotNull(propertiesObj, "expected stateSchema.properties to be present"); + Map properties = (Map) propertiesObj; + assertTrue(properties.containsKey("status"), "expected a 'status' property in the schema"); + assertTrue(properties.containsKey("counter"), "expected a 'counter' property in the schema"); + } +} diff --git a/genkit/src/test/java/com/google/genkit/ReflectionServerV2BidiTest.java b/genkit/src/test/java/com/google/genkit/ReflectionServerV2BidiTest.java new file mode 100644 index 000000000..3f0a31b7f --- /dev/null +++ b/genkit/src/test/java/com/google/genkit/ReflectionServerV2BidiTest.java @@ -0,0 +1,220 @@ +/* + * 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. + * + * SPDX-License-Identifier: Apache-2.0 + */ + +package com.google.genkit; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import com.fasterxml.jackson.databind.JsonNode; +import com.fasterxml.jackson.databind.ObjectMapper; +import com.fasterxml.jackson.databind.node.ObjectNode; +import com.google.genkit.core.BidiActionImpl; +import com.google.genkit.core.DefaultRegistry; +import com.google.genkit.core.Registry; +import java.util.List; +import java.util.concurrent.CopyOnWriteArrayList; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicReference; +import org.junit.jupiter.api.Test; + +/** + * Tests the Reflection V2 bidirectional agent path: runAction with {@code streamInput}, input + * chunks via {@code sendInputStreamChunk}, end-of-input via {@code endInputStream}, and the + * resulting {@code streamChunk}/{@code result} notifications — i.e. how the Dev UI drives an agent + * chat turn. + */ +class ReflectionServerV2BidiTest { + + private static final ObjectMapper MAPPER = new ObjectMapper(); + + /** Registers a bidi "echo" agent that counts inputs and echoes each one as a stream chunk. */ + private static Registry registryWithEchoAgent() { + Registry registry = new DefaultRegistry(); + BidiActionImpl echo = + BidiActionImpl.builder() + .name("echo") + .inputClass(JsonNode.class) + .outputClass(JsonNode.class) + .streamClass(JsonNode.class) + .initClass(JsonNode.class) + .handler( + (ctx, init, inputs, cb) -> { + int n = 0; + while (true) { + var next = inputs.next(); + if (next.isEmpty()) { + break; + } + n++; + if (cb != null) { + ObjectNode chunk = MAPPER.createObjectNode(); + chunk.set("echo", next.get()); + cb.accept(chunk); + } + } + ObjectNode out = MAPPER.createObjectNode(); + out.put("count", n); + return out; + }) + .build(); + echo.register(registry); + return registry; + } + + /** + * Collects outbound JSON-RPC messages; fires {@code done} when the response for {@code id} lands. + */ + private static final class Collector { + final List messages = new CopyOnWriteArrayList<>(); + final CountDownLatch done = new CountDownLatch(1); + final String awaitId; + + Collector(String awaitId) { + this.awaitId = awaitId; + } + + void accept(String raw) { + try { + JsonNode msg = MAPPER.readTree(raw); + messages.add(msg); + if (msg.has("result") && msg.has("id") && awaitId.equals(msg.get("id").asText())) { + done.countDown(); + } + } catch (Exception e) { + throw new RuntimeException(e); + } + } + + long countNotifications(String method, String requestId) { + return messages.stream() + .filter(m -> m.has("method") && method.equals(m.get("method").asText())) + .filter( + m -> + m.has("params") + && m.get("params").has("requestId") + && requestId.equals(m.get("params").get("requestId").asText())) + .count(); + } + + JsonNode finalResultFor(String id) { + return messages.stream() + .filter(m -> m.has("result") && m.has("id") && id.equals(m.get("id").asText())) + .map(m -> m.get("result")) + .findFirst() + .orElse(null); + } + } + + private static String runActionMsg(String id) { + return "{\"jsonrpc\":\"2.0\",\"method\":\"runAction\",\"params\":{\"key\":\"/agent/echo\"," + + "\"init\":{},\"stream\":true,\"streamInput\":true},\"id\":\"" + + id + + "\"}"; + } + + private static String inputChunkMsg(String requestId, String text) { + return "{\"jsonrpc\":\"2.0\",\"method\":\"sendInputStreamChunk\",\"params\":{\"requestId\":\"" + + requestId + + "\",\"chunk\":{\"message\":\"" + + text + + "\"}}}"; + } + + private static String endInputMsg(String requestId) { + return "{\"jsonrpc\":\"2.0\",\"method\":\"endInputStream\",\"params\":{\"requestId\":\"" + + requestId + + "\"}}"; + } + + @Test + void bidiTurnStreamsInputsAndReturnsResult() throws Exception { + ReflectionServerV2 server = + new ReflectionServerV2(registryWithEchoAgent(), "ws://localhost:1", "test"); + Collector collector = new Collector("b1"); + server.setOutboundSinkForTesting(collector::accept); + + server.handleMessageForTesting(runActionMsg("b1")); + server.handleMessageForTesting(inputChunkMsg("b1", "a")); + server.handleMessageForTesting(inputChunkMsg("b1", "b")); + server.handleMessageForTesting(endInputMsg("b1")); + + assertTrue(collector.done.await(5, TimeUnit.SECONDS), "final result not received in time"); + + // runActionState (with traceId) was emitted. + boolean sawState = + collector.messages.stream() + .anyMatch(m -> m.has("method") && "runActionState".equals(m.get("method").asText())); + assertTrue(sawState, "expected a runActionState notification"); + + // Two input chunks → two streamChunk notifications. + assertEquals(2, collector.countNotifications("streamChunk", "b1")); + + // Final result reflects 2 inputs counted by the bidi handler. + JsonNode result = collector.finalResultFor("b1"); + assertNotNull(result, "expected a final result for b1"); + assertEquals(2, result.get("result").get("count").asInt()); + } + + @Test + void earlyInputChunksAreBufferedBeforeRunAction() throws Exception { + ReflectionServerV2 server = + new ReflectionServerV2(registryWithEchoAgent(), "ws://localhost:1", "test"); + Collector collector = new Collector("b2"); + server.setOutboundSinkForTesting(collector::accept); + + // Inputs arrive BEFORE the runAction handler is dispatched — must be buffered, not dropped. + server.handleMessageForTesting(inputChunkMsg("b2", "x")); + server.handleMessageForTesting(inputChunkMsg("b2", "y")); + server.handleMessageForTesting(inputChunkMsg("b2", "z")); + server.handleMessageForTesting(endInputMsg("b2")); + server.handleMessageForTesting(runActionMsg("b2").replace("\"b1\"", "\"b2\"")); + + assertTrue(collector.done.await(5, TimeUnit.SECONDS), "final result not received in time"); + JsonNode result = collector.finalResultFor("b2"); + assertNotNull(result, "expected a final result for b2"); + assertEquals(3, result.get("result").get("count").asInt()); + } + + @Test + void listActionsExposesBidiMetadata() { + ReflectionServerV2 server = + new ReflectionServerV2(registryWithEchoAgent(), "ws://localhost:1", "test"); + AtomicReference response = new AtomicReference<>(); + server.setOutboundSinkForTesting( + raw -> { + try { + response.set(MAPPER.readTree(raw)); + } catch (Exception e) { + throw new RuntimeException(e); + } + }); + + server.handleMessageForTesting( + "{\"jsonrpc\":\"2.0\",\"method\":\"listActions\",\"id\":\"l1\"}"); + + JsonNode msg = response.get(); + assertNotNull(msg, "expected a listActions response"); + JsonNode actions = msg.get("result").get("actions"); + JsonNode echo = actions.get("/agent/echo"); + assertNotNull(echo, "expected /agent/echo in listActions"); + assertTrue(echo.get("metadata").get("bidi").asBoolean(), "expected metadata.bidi == true"); + } +} diff --git a/genkit/src/test/java/com/google/genkit/ReflectionServerV2MultiTurnTest.java b/genkit/src/test/java/com/google/genkit/ReflectionServerV2MultiTurnTest.java new file mode 100644 index 000000000..39aa51518 --- /dev/null +++ b/genkit/src/test/java/com/google/genkit/ReflectionServerV2MultiTurnTest.java @@ -0,0 +1,365 @@ +/* + * 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. + * + * SPDX-License-Identifier: Apache-2.0 + */ + +package com.google.genkit; + +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import com.fasterxml.jackson.databind.JsonNode; +import com.fasterxml.jackson.databind.ObjectMapper; +import com.google.genkit.ai.Message; +import com.google.genkit.ai.Role; +import com.google.genkit.ai.agent.AgentFinishReason; +import com.google.genkit.ai.agent.AgentResult; +import com.google.genkit.ai.agent.CustomAgentConfig; +import com.google.genkit.ai.agent.internal.AgentActions; +import com.google.genkit.core.DefaultRegistry; +import com.google.genkit.core.Registry; +import java.util.List; +import java.util.Map; +import java.util.concurrent.CopyOnWriteArrayList; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.TimeUnit; +import org.junit.jupiter.api.Test; + +/** + * Reproduction tests for multi-turn agents (client-managed and server-managed) via the Reflection + * V2 wire path. + * + *

Simulates the Dev UI pattern: one {@code runAction} per user message, each carrying {@code + * init}. For client-managed agents, the prior turn's {@code result.state} is resent as the next + * turn's {@code init.state}. For server-managed agents, the prior turn's {@code result.snapshotId} + * is resent as the next turn's {@code init.snapshotId}. Validates that history accumulates across + * turns in both modes. + */ +class ReflectionServerV2MultiTurnTest { + + private static final ObjectMapper MAPPER = new ObjectMapper(); + + /** Builds a registry with a server-managed echo agent (InMemorySessionStore). */ + private static Registry registryWithServerManagedEchoAgent() { + Registry registry = new DefaultRegistry(); + + AgentActions.defineCustomAgent( + registry, + // .store() → server-managed + CustomAgentConfig.>builder() + .name("echoStore") + .store(new com.google.genkit.ai.agent.InMemorySessionStore<>()) + .build(), + (runner, ctx) -> { + List msgs = runner.getMessages(); + int msgCount = msgs.size(); + String userText = ""; + for (int i = msgs.size() - 1; i >= 0; i--) { + Message m = msgs.get(i); + if (Role.USER.equals(m.getRole())) { + userText = m.getText() != null ? m.getText() : ""; + break; + } + } + String replyText = "reply to '" + userText + "'; history=" + msgCount; + return AgentResult.builder() + .message(Message.model(replyText)) + .finishReason(AgentFinishReason.STOP) + .build(); + }); + + return registry; + } + + /** Builds a registry with a client-managed echo agent (no store). */ + private static Registry registryWithClientManagedEchoAgent() { + Registry registry = new DefaultRegistry(); + + // AgentFn: read the latest user message, reply with count of all messages seen so far + AgentActions.defineCustomAgent( + registry, + // NO .store() → client-managed + CustomAgentConfig.>builder().name("echo").build(), + (runner, ctx) -> { + List msgs = runner.getMessages(); + int msgCount = msgs.size(); + // Find the latest user message text using Message.getText() convenience method + String userText = ""; + for (int i = msgs.size() - 1; i >= 0; i--) { + Message m = msgs.get(i); + if (Role.USER.equals(m.getRole())) { + userText = m.getText() != null ? m.getText() : ""; + break; + } + } + String replyText = "reply to '" + userText + "'; history=" + msgCount; + return AgentResult.builder() + .message(Message.model(replyText)) + .finishReason(AgentFinishReason.STOP) + .build(); + }); + + return registry; + } + + /** Collects outbound JSON-RPC messages; tracks final result by id. */ + private static final class Collector { + final List messages = new CopyOnWriteArrayList<>(); + // Per-id latches; we pre-create for ids "1" and "2" + final CountDownLatch done1 = new CountDownLatch(1); + final CountDownLatch done2 = new CountDownLatch(1); + + void accept(String raw) { + try { + JsonNode msg = MAPPER.readTree(raw); + messages.add(msg); + if (msg.has("result") && msg.has("id")) { + String id = msg.get("id").asText(); + if ("1".equals(id)) done1.countDown(); + else if ("2".equals(id)) done2.countDown(); + } + } catch (Exception e) { + throw new RuntimeException(e); + } + } + + JsonNode finalResultFor(String id) { + return messages.stream() + .filter(m -> m.has("result") && m.has("id") && id.equals(m.get("id").asText())) + .map(m -> m.get("result")) + .findFirst() + .orElse(null); + } + } + + // --- JSON-RPC message helpers --- + + private static String runActionMsg(String id, String initJson) { + return runActionMsg(id, "/agent/echo", initJson); + } + + private static String runActionMsg(String id, String key, String initJson) { + return "{" + + "\"jsonrpc\":\"2.0\"," + + "\"method\":\"runAction\"," + + "\"params\":{" + + "\"key\":\"" + + key + + "\"," + + "\"init\":" + + initJson + + "," + + "\"stream\":true," + + "\"streamInput\":true" + + "}," + + "\"id\":\"" + + id + + "\"" + + "}"; + } + + private static String inputChunkMsg(String requestId, String messageText) { + return "{" + + "\"jsonrpc\":\"2.0\"," + + "\"method\":\"sendInputStreamChunk\"," + + "\"params\":{" + + "\"requestId\":\"" + + requestId + + "\"," + + "\"chunk\":{" + + "\"message\":{" + + "\"role\":\"user\"," + + "\"content\":[{\"text\":\"" + + messageText + + "\"}]" + + "}" + + "}" + + "}" + + "}"; + } + + private static String endInputMsg(String requestId) { + return "{" + + "\"jsonrpc\":\"2.0\"," + + "\"method\":\"endInputStream\"," + + "\"params\":{" + + "\"requestId\":\"" + + requestId + + "\"" + + "}" + + "}"; + } + + @Test + void serverManagedMultiTurnAccumulatesHistory() throws Exception { + Registry registry = registryWithServerManagedEchoAgent(); + ReflectionServerV2 server = new ReflectionServerV2(registry, "ws://localhost:1", "test"); + Collector collector = new Collector(); + server.setOutboundSinkForTesting(collector::accept); + + // ── Turn 1: send "hi" ───────────────────────────────────────────────────── + server.handleMessageForTesting(runActionMsg("1", "/agent/echoStore", "{}")); + server.handleMessageForTesting(inputChunkMsg("1", "hi")); + server.handleMessageForTesting(endInputMsg("1")); + + assertTrue(collector.done1.await(10, TimeUnit.SECONDS), "Turn 1 result not received in time"); + + JsonNode turn1Result = collector.finalResultFor("1"); + assertNotNull(turn1Result, "Turn 1: expected a result"); + + String turn1Json = MAPPER.writerWithDefaultPrettyPrinter().writeValueAsString(turn1Result); + System.out.println("=== SERVER-MANAGED TURN 1 RESULT ==="); + System.out.println(turn1Json); + + // For server-managed: result.result must have snapshotId and NO inline state + JsonNode turn1Inner = turn1Result.get("result"); + assertNotNull(turn1Inner, "Turn 1: expected result.result"); + JsonNode snapshotIdNode = turn1Inner.get("snapshotId"); + assertNotNull(snapshotIdNode, "Turn 1: expected result.result.snapshotId (server-managed)"); + assertFalse(snapshotIdNode.isNull(), "Turn 1: snapshotId must not be null"); + assertFalse(snapshotIdNode.asText().isEmpty(), "Turn 1: snapshotId must not be blank"); + assertTrue( + turn1Inner.get("state") == null || turn1Inner.get("state").isNull(), + "Turn 1: server-managed must NOT return inline state"); + + String capturedSnapshotId = snapshotIdNode.asText(); + System.out.println("Turn 1 snapshotId: " + capturedSnapshotId); + + // ── Turn 2: resume with snapshotId from turn 1 ─────────────────────────── + String initJson2 = "{\"snapshotId\":\"" + capturedSnapshotId + "\"}"; + + server.handleMessageForTesting(runActionMsg("2", "/agent/echoStore", initJson2)); + server.handleMessageForTesting(inputChunkMsg("2", "again")); + server.handleMessageForTesting(endInputMsg("2")); + + assertTrue(collector.done2.await(10, TimeUnit.SECONDS), "Turn 2 result not received in time"); + + JsonNode turn2Result = collector.finalResultFor("2"); + assertNotNull(turn2Result, "Turn 2: expected a result"); + + String turn2Json = MAPPER.writerWithDefaultPrettyPrinter().writeValueAsString(turn2Result); + System.out.println("=== SERVER-MANAGED TURN 2 RESULT ==="); + System.out.println(turn2Json); + + JsonNode turn2Inner = turn2Result.get("result"); + assertNotNull(turn2Inner, "Turn 2: expected result.result"); + // Turn 2 should also yield a snapshotId (session continues) + JsonNode snap2Node = turn2Inner.get("snapshotId"); + assertNotNull(snap2Node, "Turn 2: expected result.result.snapshotId"); + assertFalse(snap2Node.isNull(), "Turn 2: snapshotId must not be null"); + + // CORE ASSERTION: the reply text must reflect accumulated history (≥4 messages across 2 turns) + // The AgentFn echoes "history=" where msgCount includes prior messages loaded from + // the store. Turn 2 sees at least 2 msgs from turn 1 + the new user msg = ≥3 total, so + // history= value must be > 1 (which turn 1 would have seen as 1 user msg → history=1). + JsonNode turn2Message = turn2Inner.get("message"); + assertNotNull(turn2Message, "Turn 2: expected a message in result"); + String turn2Text = + turn2Message.get("content") != null + ? turn2Message.get("content").get(0).get("text").asText() + : ""; + System.out.println("Turn 2 reply text: " + turn2Text); + + // Extract the history count from the reply: "reply to 'again'; history=N" + int historyCount = 0; + if (turn2Text.contains("history=")) { + historyCount = Integer.parseInt(turn2Text.substring(turn2Text.indexOf("history=") + 8)); + } + System.out.println("Turn 2 history count seen by AgentFn: " + historyCount); + assertTrue( + historyCount >= 3, + "BUG: Turn 2 AgentFn saw history=" + + historyCount + + " but expected ≥3 (2 msgs from turn 1 + new user msg). " + + "Server-managed history was NOT loaded from the store."); + } + + @Test + void clientManagedMultiTurnAccumulatesHistory() throws Exception { + Registry registry = registryWithClientManagedEchoAgent(); + ReflectionServerV2 server = new ReflectionServerV2(registry, "ws://localhost:1", "test"); + Collector collector = new Collector(); + server.setOutboundSinkForTesting(collector::accept); + + // ── Turn 1: send "hi" ───────────────────────────────────────────────────── + server.handleMessageForTesting(runActionMsg("1", "{}")); + server.handleMessageForTesting(inputChunkMsg("1", "hi")); + server.handleMessageForTesting(endInputMsg("1")); + + assertTrue(collector.done1.await(10, TimeUnit.SECONDS), "Turn 1 result not received in time"); + + JsonNode turn1Result = collector.finalResultFor("1"); + assertNotNull(turn1Result, "Turn 1: expected a result"); + + // Print turn 1 result for evidence + String turn1Json = MAPPER.writerWithDefaultPrettyPrinter().writeValueAsString(turn1Result); + System.out.println("=== TURN 1 RESULT ==="); + System.out.println(turn1Json); + + // For client-managed: result.result.state must be present with messages + JsonNode turn1Inner = turn1Result.get("result"); + assertNotNull(turn1Inner, "Turn 1: expected result.result"); + JsonNode turn1State = turn1Inner.get("state"); + assertNotNull( + turn1State, + "Turn 1: expected result.result.state (client-managed must return inline state)"); + JsonNode turn1Messages = turn1State.get("messages"); + assertNotNull(turn1Messages, "Turn 1: expected state.messages"); + System.out.println("Turn 1 message count in state: " + turn1Messages.size()); + + // ── Turn 2: send "again" with turn-1 state resent as init.state ────────── + // Serialize the captured state to embed into the next runAction init + String stateJson = MAPPER.writeValueAsString(turn1State); + String initJson = "{\"state\":" + stateJson + "}"; + + server.handleMessageForTesting(runActionMsg("2", initJson)); + server.handleMessageForTesting(inputChunkMsg("2", "again")); + server.handleMessageForTesting(endInputMsg("2")); + + assertTrue(collector.done2.await(10, TimeUnit.SECONDS), "Turn 2 result not received in time"); + + JsonNode turn2Result = collector.finalResultFor("2"); + assertNotNull(turn2Result, "Turn 2: expected a result"); + + // Print turn 2 result for evidence + String turn2Json = MAPPER.writerWithDefaultPrettyPrinter().writeValueAsString(turn2Result); + System.out.println("=== TURN 2 RESULT ==="); + System.out.println(turn2Json); + + JsonNode turn2Inner = turn2Result.get("result"); + assertNotNull(turn2Inner, "Turn 2: expected result.result"); + JsonNode turn2State = turn2Inner.get("state"); + assertNotNull(turn2State, "Turn 2: expected result.result.state"); + JsonNode turn2Messages = turn2State.get("messages"); + assertNotNull(turn2Messages, "Turn 2: expected state.messages"); + System.out.println("Turn 2 message count in state: " + turn2Messages.size()); + + // CORE ASSERTION: history must have grown between turn 1 and turn 2 + int turn1Count = turn1Messages.size(); + int turn2Count = turn2Messages.size(); + System.out.println( + "History: turn1=" + turn1Count + " messages, turn2=" + turn2Count + " messages"); + assertTrue( + turn2Count > turn1Count, + "BUG REPRODUCED: Turn 2 state.messages (" + + turn2Count + + ") did NOT grow beyond turn 1 (" + + turn1Count + + ") — history was NOT accumulated across turns. " + + "This means init.state was not hydrated correctly for turn 2."); + } +} diff --git a/genkit/src/test/java/com/google/genkit/conformance/agent/AgentConformanceTest.java b/genkit/src/test/java/com/google/genkit/conformance/agent/AgentConformanceTest.java new file mode 100644 index 000000000..826bfe134 --- /dev/null +++ b/genkit/src/test/java/com/google/genkit/conformance/agent/AgentConformanceTest.java @@ -0,0 +1,784 @@ +/* + * 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. + * + * SPDX-License-Identifier: Apache-2.0 + */ + +package com.google.genkit.conformance.agent; + +import static org.junit.jupiter.api.Assertions.assertFalse; + +import com.fasterxml.jackson.databind.JsonNode; +import com.fasterxml.jackson.databind.ObjectMapper; +import com.fasterxml.jackson.databind.node.ObjectNode; +import com.fasterxml.jackson.dataformat.yaml.YAMLFactory; +import com.google.genkit.ai.agent.Agent; +import com.google.genkit.ai.agent.GetSnapshotRequest; +import com.google.genkit.ai.agent.SessionSnapshot; +import com.google.genkit.ai.agent.SnapshotStatus; +import com.google.genkit.core.ActionContext; +import com.google.genkit.core.BufferedInputSource; +import com.google.genkit.core.GenkitException; +import com.google.genkit.core.JsonUtils; +import java.io.InputStream; +import java.util.ArrayList; +import java.util.HashMap; +import java.util.Iterator; +import java.util.List; +import java.util.Map; +import java.util.regex.Matcher; +import java.util.regex.Pattern; +import org.junit.jupiter.api.AfterAll; +import org.junit.jupiter.api.DynamicTest; +import org.junit.jupiter.api.TestFactory; + +/** + * Cross-language agent conformance harness, driven by {@code conformance/agent.yaml} (a verbatim + * copy of the upstream {@code tests/specs/agent.yaml}). + * + *

The harness ports the JS ({@code js/ai/tests/agents_spec_test.ts}) and Go ({@code + * go/ai/exp/agents_conformance_test.go}) reference runners. It drives the in-process Agent API: + * {@code agent.runBidiJson(ctx, init, inputs, chunkSink)} for {@code send}, {@code + * agent.getSnapshotData(...)} for snapshot lookups, and {@code agent.abort(...)} for aborts. + * Everything is matched as JSON (the wire format), mirroring the Go harness's canonical-JSON + * approach, since every Java agent wire type serializes to the exact spec field names. + * + *

Each case is a dynamic test. A case may PASS, SKIP (a documented first-cut gap — e.g. + * prompt-agent interrupt/restart, tool-response stream chunks), or FAIL (a real mismatch). Skips do + * not fail the suite; the class passes overall and prints a {@code N passed / M skipped / K failed} + * summary with per-case reasons. + */ +class AgentConformanceTest { + + private static final ObjectMapper YAML = new ObjectMapper(new YAMLFactory()); + private static final ObjectMapper JSON = JsonUtils.getObjectMapper(); + + // ── result accounting ──────────────────────────────────────────────────────────── + + private enum Outcome { + PASS, + SKIP, + FAIL + } + + private record CaseResult(String name, Outcome outcome, String reason) {} + + private static final List RESULTS = new ArrayList<>(); + + /** Thrown by a step runner to mark the whole case as a documented skip. */ + private static final class SkipException extends RuntimeException { + SkipException(String reason) { + super(reason); + } + } + + // ── test factory ────────────────────────────────────────────────────────────────── + + @TestFactory + Iterable agentConformance() throws Exception { + JsonNode suite; + try (InputStream in = + AgentConformanceTest.class.getResourceAsStream("/conformance/agent.yaml")) { + if (in == null) { + throw new IllegalStateException("conformance/agent.yaml not found on the test classpath"); + } + suite = YAML.readTree(in); + } + JsonNode tests = suite.get("tests"); + assertFalse(tests == null || tests.isEmpty(), "spec contains no tests"); + + List dynamic = new ArrayList<>(); + for (JsonNode tc : tests) { + String name = tc.get("name").asText(); + dynamic.add( + DynamicTest.dynamicTest( + name, + () -> { + try { + runCase(tc); + RESULTS.add(new CaseResult(name, Outcome.PASS, null)); + } catch (SkipException se) { + RESULTS.add(new CaseResult(name, Outcome.SKIP, se.getMessage())); + } catch (AssertionError | Exception e) { + RESULTS.add(new CaseResult(name, Outcome.FAIL, oneLine(e.getMessage()))); + // Re-throw so the dynamic test is reported as failed individually, but the + // suite-level @AfterAll guard below decides whether the build fails. + throw e; + } + })); + } + return dynamic; + } + + /** + * Runs one case's steps in order, stopping at the first failed step (later steps need captures). + */ + private void runCase(JsonNode tc) throws Exception { + String agentName = tc.get("agent").asText(); + Fixtures fixtures = new Fixtures(); + Agent> agent = fixtures.agent(agentName); + if (agent == null) { + throw new AssertionError("unknown agent '" + agentName + "'"); + } + ActionContext ctx = new ActionContext(fixtures.genkit().getRegistry()); + Map captures = new HashMap<>(); + + JsonNode steps = tc.get("steps"); + for (int i = 0; i < steps.size(); i++) { + JsonNode step = steps.get(i); + String type = step.get("type").asText(); + String label = "step[" + i + "] (" + type + ")"; + switch (type) { + case "send" -> runSend(label, fixtures, agent, ctx, step, captures); + case "getSnapshotData" -> runGetSnapshotData(label, agent, step, captures); + case "abort" -> runAbort(label, agent, step, captures); + case "waitUntilCompleted" -> runWaitUntilCompleted(label, agent, step, captures); + default -> throw new AssertionError(label + ": unknown step type '" + type + "'"); + } + } + } + + // ── send ──────────────────────────────────────────────────────────────────────── + + private void runSend( + String label, + Fixtures fixtures, + Agent> agent, + ActionContext ctx, + JsonNode rawStep, + Map captures) { + JsonNode step = resolve(rawStep, captures); + + // Program the model for this step. + List modelResponses = toNodeList(step.get("modelResponses")); + List> streamChunks = new ArrayList<>(); + if (step.has("streamChunks")) { + for (JsonNode perCall : step.get("streamChunks")) { + streamChunks.add(toNodeList(perCall)); + } + } + fixtures.programmableModel().program(modelResponses, streamChunks); + + // Build init JSON from the step's init (empty object when absent). + JsonNode init = + step.has("init") && !step.get("init").isNull() ? step.get("init") : JSON.createObjectNode(); + + // Feed inputs. + BufferedInputSource inputs = new BufferedInputSource<>(); + if (step.has("inputs")) { + for (JsonNode input : step.get("inputs")) { + inputs.offer(input); + } + } + inputs.end(); + + List chunks = new ArrayList<>(); + JsonNode output = null; + GenkitException thrown = null; + try { + output = agent.runBidiJson(ctx, init, inputs, chunks::add); + } catch (GenkitException ge) { + thrown = ge; + } + + // expectError: API misuse — the turn must throw; assert the status (not the message). + if (step.has("expectError")) { + JsonNode ee = step.get("expectError"); + if (thrown == null) { + throw new AssertionError(label + ": expected the turn to throw, but it resolved"); + } + if (ee.has("status")) { + String want = ee.get("status").asText(); + String got = thrown.getErrorCode(); + if (!want.equals(got)) { + throw new AssertionError( + label + + ": expectError.status: want '" + + want + + "', got '" + + got + + "' (msg: " + + thrown.getMessage() + + ")"); + } + } + return; + } + + if (thrown != null) { + throw new AssertionError(label + ": invocation threw unexpectedly: " + thrown.getMessage()); + } + + // expectChunks: semi-strict ordered comparison. + if (step.has("expectChunks")) { + assertChunks(label, chunks, step.get("expectChunks")); + } + + // expectOutput. + if (step.has("expectOutput")) { + assertOutput(label, output, step.get("expectOutput")); + } + + // captures. + final JsonNode out = output; + if (out != null) { + capture( + rawStep, "captureSnapshotId", () -> textOrNull(out.get("snapshotId")), captures, label); + capture(rawStep, "captureSessionId", () -> stateSessionId(out), captures, label); + if (rawStep.has("captureState")) { + JsonNode state = out.get("state"); + if (state == null || state.isNull()) { + throw new AssertionError(label + ": captureState requested but output has no state"); + } + captures.put(rawStep.get("captureState").asText(), state); + } + } + } + + // ── getSnapshotData ─────────────────────────────────────────────────────────────── + + private void runGetSnapshotData( + String label, + Agent> agent, + JsonNode rawStep, + Map captures) { + JsonNode step = resolve(rawStep, captures); + String snapshotId = textOrNull(step.get("snapshotId")); + String sessionId = textOrNull(step.get("sessionId")); + if ((snapshotId == null) == (sessionId == null)) { + throw new AssertionError(label + ": requires exactly one of snapshotId / sessionId"); + } + + GetSnapshotRequest.Builder req = GetSnapshotRequest.builder(); + if (snapshotId != null) { + req.snapshotId(snapshotId); + } else { + req.sessionId(sessionId); + } + + SessionSnapshot> snap; + try { + snap = agent.getSnapshotData(req.build()); + } catch (RuntimeException e) { + if (step.has("expectError")) { + return; // expected to throw; message wording is not asserted cross-language. + } + throw new AssertionError(label + ": getSnapshotData threw: " + e.getMessage()); + } + + if (step.has("expectError")) { + throw new AssertionError(label + ": expected getSnapshotData to throw, but it succeeded"); + } + if (snap == null) { + throw new AssertionError(label + ": snapshot not found"); + } + if (step.has("expectSnapshot")) { + assertSnapshot(label, JsonUtils.toJsonNode(snap), step.get("expectSnapshot")); + } + } + + // ── abort ────────────────────────────────────────────────────────────────────── + + private void runAbort( + String label, + Agent> agent, + JsonNode rawStep, + Map captures) { + JsonNode step = resolve(rawStep, captures); + String snapshotId = textOrNull(step.get("snapshotId")); + if (snapshotId == null) { + throw new AssertionError(label + ": abort requires snapshotId"); + } + + // The spec's expectPreviousStatus is the status *before* the abort. agent.abort() returns the + // status after the attempt: for a PENDING snapshot it returns ABORTED (so the previous was + // PENDING); for terminal snapshots it returns the unchanged terminal status (== previous); for + // a non-existent snapshot it returns null (== previous absent). Read the prior status directly + // so the assertion matches the spec's semantics exactly. + SnapshotStatus previous = null; + SessionSnapshot> before = + agent.getSnapshotData(GetSnapshotRequest.builder().snapshotId(snapshotId).build()); + if (before != null) { + previous = before.getStatus() != null ? before.getStatus() : SnapshotStatus.COMPLETED; + } + agent.abort(snapshotId); + + if (rawStep.has("expectPreviousStatus")) { + JsonNode want = rawStep.get("expectPreviousStatus"); + String wantStr = (want == null || want.isNull()) ? null : want.asText(); + String gotStr = previous != null ? previous.getValue() : null; + if (!java.util.Objects.equals(wantStr, gotStr)) { + throw new AssertionError( + label + ": expectPreviousStatus: want '" + wantStr + "', got '" + gotStr + "'"); + } + } + } + + // ── waitUntilCompleted ──────────────────────────────────────────────────────────── + + private void runWaitUntilCompleted( + String label, + Agent> agent, + JsonNode rawStep, + Map captures) + throws InterruptedException { + JsonNode step = resolve(rawStep, captures); + String snapshotId = textOrNull(step.get("snapshotId")); + if (snapshotId == null) { + throw new AssertionError(label + ": waitUntilCompleted requires snapshotId"); + } + long timeoutMs = step.has("timeoutMs") ? step.get("timeoutMs").asLong() : 5000L; + + long deadline = System.currentTimeMillis() + timeoutMs; + SessionSnapshot> snap = null; + while (System.currentTimeMillis() < deadline) { + snap = agent.getSnapshotData(GetSnapshotRequest.builder().snapshotId(snapshotId).build()); + if (snap != null && isTerminal(snap.getStatus())) { + break; + } + Thread.sleep(50); + } + if (snap == null || !isTerminal(snap.getStatus())) { + throw new AssertionError( + label + + ": snapshot did not reach a terminal status within " + + timeoutMs + + "ms (status=" + + (snap == null ? "" : snap.getStatus()) + + ")"); + } + if (step.has("expectSnapshot")) { + assertSnapshot(label, JsonUtils.toJsonNode(snap), step.get("expectSnapshot")); + } + } + + private static boolean isTerminal(SnapshotStatus s) { + return s == SnapshotStatus.COMPLETED + || s == SnapshotStatus.FAILED + || s == SnapshotStatus.ABORTED; + } + + // ── chunk assertions ────────────────────────────────────────────────────────────── + + /** + * Semi-strict ordered chunk comparison: same length and order, with type-aware matching per chunk + * (turnEnd asserts presence and, when specified, finishReason ignoring the dynamic snapshotId; + * modelChunk / artifact / customPatch use CONTAINS). + */ + private void assertChunks(String label, List actual, JsonNode expected) { + if (expected.size() != actual.size()) { + throw new AssertionError( + label + + ": expected " + + expected.size() + + " chunks, got " + + actual.size() + + "\n expected: " + + expected + + "\n actual: " + + actual); + } + for (int i = 0; i < expected.size(); i++) { + JsonNode exp = expected.get(i); + JsonNode act = actual.get(i); + String err = matchChunk(act, exp); + if (err != null) { + throw new AssertionError( + label + ": chunk[" + i + "]: " + err + "\n expected: " + exp + "\n actual: " + act); + } + } + } + + private String matchChunk(JsonNode actual, JsonNode expected) { + if (expected.has("turnEnd")) { + JsonNode te = actual.get("turnEnd"); + if (te == null || te.isNull()) { + return "expected a turnEnd chunk"; + } + JsonNode exTe = expected.get("turnEnd"); + if (exTe != null && exTe.has("finishReason")) { + // snapshotId is dynamic; only finishReason is asserted when specified. + if (!nodeEquals(te.get("finishReason"), exTe.get("finishReason"))) { + return "turnEnd.finishReason: want " + + exTe.get("finishReason") + + ", got " + + te.get("finishReason"); + } + } + return null; + } + if (expected.has("modelChunk")) { + JsonNode exMc = expected.get("modelChunk"); + JsonNode acMc = actual.get("modelChunk"); + return contains(acMc, exMc, "modelChunk"); + } + for (String key : new String[] {"artifact", "customPatch"}) { + if (expected.has(key)) { + return contains(actual.get(key), expected.get(key), key); + } + } + return contains(actual, expected, "chunk"); + } + + // ── output assertions ───────────────────────────────────────────────────────────── + + private void assertOutput(String label, JsonNode output, JsonNode expect) { + if (output == null) { + throw new AssertionError(label + ": no output"); + } + if (expect.has("message")) { + String err = contains(output.get("message"), expect.get("message"), "output.message"); + if (err != null) { + throw new AssertionError(label + ": " + err); + } + } + if (boolField(expect, "hasSnapshotId") && isBlank(output.get("snapshotId"))) { + throw new AssertionError(label + ": expected output.snapshotId to be non-empty"); + } + if (boolField(expect, "hasSessionId") && stateSessionId(output) == null) { + throw new AssertionError(label + ": expected output.state.sessionId to be non-empty"); + } + if (expect.has("stateContains")) { + String err = contains(output.get("state"), expect.get("stateContains"), "output.state"); + if (err != null) { + throw new AssertionError(label + ": " + err); + } + } + if (expect.has("artifactsContain")) { + assertArtifactsContain(label, output.get("artifacts"), expect.get("artifactsContain")); + } + if (expect.has("finishReason")) { + if (!nodeEquals(output.get("finishReason"), expect.get("finishReason"))) { + throw new AssertionError( + label + + ": output.finishReason: want " + + expect.get("finishReason") + + ", got " + + output.get("finishReason")); + } + } + if (expect.has("errorContains")) { + assertErrorContains(label, "output.error", output.get("error"), expect.get("errorContains")); + } + } + + private void assertSnapshot(String label, JsonNode snap, JsonNode expect) { + if (expect.has("parentId")) { + JsonNode actualParent = snap.get("parentId"); + if (!nodeEquals(actualParent, expect.get("parentId"))) { + throw new AssertionError( + label + + ": snapshot.parentId: want " + + expect.get("parentId") + + ", got " + + actualParent); + } + } + if (expect.has("status")) { + JsonNode status = snap.get("status"); + // A null/absent status is COMPLETED on the wire. + String got = (status == null || status.isNull()) ? "completed" : status.asText(); + if (!expect.get("status").asText().equals(got)) { + throw new AssertionError( + label + ": snapshot.status: want " + expect.get("status").asText() + ", got " + got); + } + } + if (expect.has("finishReason")) { + if (!nodeEquals(snap.get("finishReason"), expect.get("finishReason"))) { + throw new AssertionError( + label + + ": snapshot.finishReason: want " + + expect.get("finishReason") + + ", got " + + snap.get("finishReason")); + } + } + if (boolField(expect, "hasSessionId") && stateSessionId(snap) == null) { + throw new AssertionError(label + ": expected snapshot.state.sessionId to be non-empty"); + } + if (expect.has("stateContains")) { + String err = contains(snap.get("state"), expect.get("stateContains"), "snapshot.state"); + if (err != null) { + throw new AssertionError(label + ": " + err); + } + } + if (expect.has("errorContains")) { + assertErrorContains(label, "snapshot.error", snap.get("error"), expect.get("errorContains")); + } + } + + private void assertArtifactsContain(String label, JsonNode actual, JsonNode expected) { + if (actual == null || !actual.isArray()) { + throw new AssertionError(label + ": expected output.artifacts to be a list, got " + actual); + } + for (JsonNode ea : expected) { + String name = ea.path("name").asText(null); + JsonNode found = null; + for (JsonNode a : actual) { + if (name != null && name.equals(a.path("name").asText(null))) { + found = a; + break; + } + } + if (found == null) { + throw new AssertionError( + label + ": expected artifact '" + name + "' not found in " + actual); + } + String err = contains(found, ea, "artifact(" + name + ")"); + if (err != null) { + throw new AssertionError(label + ": " + err); + } + } + } + + /** + * Matches a structured error: presence + status (exactly). The message is intentionally NOT + * asserted — wording is implementation-specific; the cross-language contract is the status. + */ + private void assertErrorContains(String label, String path, JsonNode actual, JsonNode expect) { + if (actual == null || actual.isNull() || !actual.isObject()) { + throw new AssertionError(label + ": expected " + path + " to be present, got " + actual); + } + if (expect.has("status")) { + if (!nodeEquals(actual.get("status"), expect.get("status"))) { + throw new AssertionError( + label + + ": " + + path + + ".status: want " + + expect.get("status") + + ", got " + + actual.get("status")); + } + } + } + + // ── contains / subsequence matchers ──────────────────────────────────────────────── + + /** + * Asserts that {@code actual} contains all fields specified in {@code expected}. Objects match + * key-by-key (extra actual keys ignored); arrays match as an ordered subsequence; scalars must be + * deep-equal. Returns {@code null} on match, or an error message describing the first mismatch. + */ + private String contains(JsonNode actual, JsonNode expected, String path) { + if (expected == null || expected.isNull()) { + return null; + } + if (expected.isArray()) { + if (actual == null || !actual.isArray()) { + return path + ": expected array, got " + actual; + } + return subsequence(actual, expected, path); + } + if (expected.isObject()) { + if (actual == null || !actual.isObject()) { + return path + ": expected object, got " + actual; + } + Iterator fields = expected.fieldNames(); + while (fields.hasNext()) { + String k = fields.next(); + String err = contains(actual.get(k), expected.get(k), path + "." + k); + if (err != null) { + return err; + } + } + return null; + } + if (!nodeEquals(actual, expected)) { + return path + ": want " + expected + ", got " + actual; + } + return null; + } + + /** Each expected item must appear in {@code actual} in order (not necessarily contiguous). */ + private String subsequence(JsonNode actual, JsonNode expected, String path) { + int idx = 0; + for (int i = 0; i < expected.size(); i++) { + JsonNode want = expected.get(i); + boolean found = false; + while (idx < actual.size()) { + if (contains(actual.get(idx), want, path + "[" + idx + "]") == null) { + found = true; + idx++; + break; + } + idx++; + } + if (!found) { + return path + ": expected item " + i + " not found in order: " + want; + } + } + return null; + } + + /** Numeric-tolerant deep equality (1 == 1.0; matches YAML ints vs JSON longs/doubles). */ + private static boolean nodeEquals(JsonNode a, JsonNode b) { + if (a == null || a.isNull()) { + return b == null || b.isNull(); + } + if (b == null || b.isNull()) { + return false; + } + if (a.isNumber() && b.isNumber()) { + return a.asDouble() == b.asDouble(); + } + return a.equals(b); + } + + // ── template resolution ──────────────────────────────────────────────────────────── + + private static final Pattern FULL = Pattern.compile("^\\{\\{(\\w+)\\}\\}$"); + private static final Pattern INLINE = Pattern.compile("\\{\\{(\\w+)\\}\\}"); + + /** + * Recursively replaces {@code {{name}}} references with previously captured values. A value that + * is exactly {@code {{name}}} is replaced by the captured node (which may be a non-string, e.g. a + * captured state object); inline occurrences are string-substituted. + */ + private JsonNode resolve(JsonNode v, Map captures) { + if (v == null) { + return null; + } + if (v.isTextual()) { + String s = v.asText(); + Matcher full = FULL.matcher(s); + if (full.matches()) { + JsonNode val = captures.get(full.group(1)); + if (val == null) { + throw new AssertionError("template reference {{" + full.group(1) + "}} not found"); + } + return val; + } + Matcher inline = INLINE.matcher(s); + StringBuilder out = new StringBuilder(); + while (inline.find()) { + JsonNode val = captures.get(inline.group(1)); + if (val == null) { + throw new AssertionError("template reference {{" + inline.group(1) + "}} not found"); + } + String rep = val.isTextual() ? val.asText() : val.toString(); + inline.appendReplacement(out, Matcher.quoteReplacement(rep)); + } + inline.appendTail(out); + return JSON.getNodeFactory().textNode(out.toString()); + } + if (v.isObject()) { + ObjectNode out = JSON.createObjectNode(); + Iterator> it = v.fields(); + while (it.hasNext()) { + Map.Entry e = it.next(); + out.set(e.getKey(), resolve(e.getValue(), captures)); + } + return out; + } + if (v.isArray()) { + var out = JSON.createArrayNode(); + for (JsonNode e : v) { + out.add(resolve(e, captures)); + } + return out; + } + return v; + } + + // ── small helpers ────────────────────────────────────────────────────────────────── + + private interface Supplier { + String get(); + } + + private void capture( + JsonNode rawStep, String key, Supplier value, Map captures, String label) { + if (!rawStep.has(key)) { + return; + } + String v = value.get(); + if (v == null) { + throw new AssertionError(label + ": " + key + " requested but the value was absent"); + } + captures.put(rawStep.get(key).asText(), JSON.getNodeFactory().textNode(v)); + } + + private static String stateSessionId(JsonNode container) { + if (container == null) { + return null; + } + JsonNode state = container.get("state"); + if (state == null || state.isNull()) { + return null; + } + return isBlank(state.get("sessionId")) ? null : state.get("sessionId").asText(); + } + + private static String textOrNull(JsonNode n) { + return (n == null || n.isNull()) ? null : n.asText(); + } + + private static boolean isBlank(JsonNode n) { + return n == null || n.isNull() || n.asText().isEmpty(); + } + + private static boolean boolField(JsonNode obj, String key) { + JsonNode n = obj.get(key); + return n != null && n.asBoolean(); + } + + private static List toNodeList(JsonNode arr) { + List out = new ArrayList<>(); + if (arr != null && arr.isArray()) { + arr.forEach(out::add); + } + return out; + } + + private static String oneLine(String s) { + if (s == null) { + return "(no message)"; + } + String t = s.replaceAll("\\s+", " ").trim(); + return t.length() > 200 ? t.substring(0, 200) + "…" : t; + } + + // ── summary ────────────────────────────────────────────────────────────────────── + + @AfterAll + static void printSummary() { + long passed = RESULTS.stream().filter(r -> r.outcome() == Outcome.PASS).count(); + long skipped = RESULTS.stream().filter(r -> r.outcome() == Outcome.SKIP).count(); + long failed = RESULTS.stream().filter(r -> r.outcome() == Outcome.FAIL).count(); + + StringBuilder sb = new StringBuilder(); + sb.append("\n========================================================================\n"); + sb.append("Agent conformance summary: ") + .append(passed) + .append(" passed / ") + .append(skipped) + .append(" skipped / ") + .append(failed) + .append(" failed (of ") + .append(RESULTS.size()) + .append(")\n"); + sb.append("------------------------------------------------------------------------\n"); + for (CaseResult r : RESULTS) { + sb.append(String.format(" %-7s %s", r.outcome(), r.name())); + if (r.reason() != null) { + sb.append(" — ").append(r.reason()); + } + sb.append('\n'); + } + sb.append("========================================================================\n"); + System.out.println(sb); + } +} diff --git a/genkit/src/test/java/com/google/genkit/conformance/agent/Fixtures.java b/genkit/src/test/java/com/google/genkit/conformance/agent/Fixtures.java new file mode 100644 index 000000000..b8d624b45 --- /dev/null +++ b/genkit/src/test/java/com/google/genkit/conformance/agent/Fixtures.java @@ -0,0 +1,334 @@ +/* + * 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. + * + * SPDX-License-Identifier: Apache-2.0 + */ + +package com.google.genkit.conformance.agent; + +import com.google.genkit.Genkit; +import com.google.genkit.GenkitOptions; +import com.google.genkit.agent.AgentConfig; +import com.google.genkit.ai.Message; +import com.google.genkit.ai.Part; +import com.google.genkit.ai.ToolInterruptException; +import com.google.genkit.ai.agent.Agent; +import com.google.genkit.ai.agent.AgentFinishReason; +import com.google.genkit.ai.agent.AgentFn; +import com.google.genkit.ai.agent.AgentResult; +import com.google.genkit.ai.agent.Artifact; +import com.google.genkit.ai.agent.CustomAgentConfig; +import com.google.genkit.ai.agent.InMemorySessionStore; +import com.google.genkit.ai.agent.internal.AgentActions; +import java.util.HashMap; +import java.util.List; +import java.util.Map; + +/** + * Builds the harness's Genkit instance with the agents, tools and programmable model that the + * {@code tests/specs/agent.yaml} conformance suite drives. + * + *

This mirrors {@code setupHarness} in the JS ({@code js/ai/tests/agents_spec_test.ts}) and Go + * ({@code go/ai/exp/agents_conformance_test.go}) reference harnesses: the same single dynamic + * (free-form {@code Map}) custom-state type serves every agent (prompt agents ignore custom state; + * the custom-state agents manipulate it). + * + *

A fresh {@link Fixtures} is created per test case (the JS/Go harnesses rebuild the registry + * before every test) so server-managed stores never leak snapshots across cases. + */ +final class Fixtures { + + /** + * Safety cap (ms) on how long {@code customAgentBlocking} blocks a turn, so a missing abort + * signal can never hang the suite. The abort-pending cases observe the PENDING snapshot and abort + * well within this window. + */ + private static final long BLOCKING_SAFETY_MILLIS = 3000L; + + private final Genkit genkit; + private final ProgrammableModel programmableModel; + private final Map>> agents = new HashMap<>(); + + Fixtures() { + this.genkit = new Genkit(GenkitOptions.builder().experimental(true).build()); + this.programmableModel = new ProgrammableModel("programmableModel"); + genkit.registerModel(programmableModel); + registerTools(); + registerPromptAgents(); + registerCustomAgents(); + } + + Genkit genkit() { + return genkit; + } + + ProgrammableModel programmableModel() { + return programmableModel; + } + + Agent> agent(String name) { + return agents.get(name); + } + + // ── tools ────────────────────────────────────────────────────────────────────── + + private com.google.genkit.ai.Tool testTool; + private com.google.genkit.ai.Tool interruptTool; + private com.google.genkit.ai.Tool restartTool; + + private void registerTools() { + // testTool: {} -> "tool called" + testTool = + genkit.defineTool( + "testTool", + "A simple test tool", + (ctx, in) -> "tool called", + Object.class, + String.class); + + // interruptTool: always pauses the turn, returning the tool request to the client for + // external resolution (resume.respond). + interruptTool = + genkit.defineTool( + "interruptTool", + "An interrupt tool", + (ctx, in) -> { + throw new ToolInterruptException(); + }, + Object.class, + Object.class); + + // restartTool: interrupts on first call; succeeds when restarted with resumed metadata. A + // restart-aware tool reads ctx.isResumed()/getResumed() to distinguish the first (interrupting) + // call from a re-invocation after resume.restart, exactly like a real human-in-the-loop tool. + restartTool = + genkit.defineTool( + "restartTool", + "A tool that requires confirmation before executing", + (ctx, in) -> { + if (ctx.isResumed()) { + // Re-invoked via resume.restart with the client's approval payload in getResumed(). + return "restarted: " + ctx.getResumed(); + } + throw new ToolInterruptException(Map.of("requiresConfirmation", true)); + }, + Object.class, + Object.class); + } + + // ── prompt-backed agents (use the programmable model) ──────────────────────────── + + private void registerPromptAgents() { + register( + "promptAgent", + genkit + .beta() + .defineAgent( + AgentConfig.>builder() + .name("promptAgent") + .model("programmableModel") + .build())); + + register( + "promptAgentWithStore", + genkit + .beta() + .defineAgent( + AgentConfig.>builder() + .name("promptAgentWithStore") + .model("programmableModel") + .store(new InMemorySessionStore<>()) + .build())); + + register( + "promptAgentWithTools", + genkit + .beta() + .defineAgent( + AgentConfig.>builder() + .name("promptAgentWithTools") + .model("programmableModel") + .tools(List.of(testTool)) + .build())); + + register( + "promptAgentWithInterrupt", + genkit + .beta() + .defineAgent( + AgentConfig.>builder() + .name("promptAgentWithInterrupt") + .model("programmableModel") + .tools(List.of(interruptTool)) + .store(new InMemorySessionStore<>()) + .build())); + + register( + "promptAgentWithRestartTool", + genkit + .beta() + .defineAgent( + AgentConfig.>builder() + .name("promptAgentWithRestartTool") + .model("programmableModel") + .tools(List.of(restartTool)) + .store(new InMemorySessionStore<>()) + .build())); + } + + // ── custom agents (deterministic, via defineCustomAgent) ───────────────────────── + + private void registerCustomAgents() { + // customAgentBlocking: server-managed, blocks until the caller requests cancellation. The + // conformance suite only ever drives it detached + abort (the abort flips the snapshot status + // in the store directly). A bounded safety timeout guarantees the harness can never hang even + // though the in-process abort does not currently propagate the signal to a running turn fn. + register( + "customAgentBlocking", + defineCustom( + "customAgentBlocking", + true, + (sess, ctx) -> { + long deadline = System.currentTimeMillis() + BLOCKING_SAFETY_MILLIS; + while (!ctx.isAborted() && System.currentTimeMillis() < deadline) { + Thread.sleep(20); + } + return AgentResult.builder() + .message(Message.model("unblocked")) + .finishReason(AgentFinishReason.STOP) + .build(); + })); + + // customAgentFailing: server-managed, fails during processing. + register( + "customAgentFailing", + defineCustom( + "customAgentFailing", + true, + (sess, ctx) -> { + throw new RuntimeException("intentional failure"); + })); + + // customAgentWithArtifacts: client-managed, streams and dedupes artifacts. + register( + "customAgentWithArtifacts", + defineCustom( + "customAgentWithArtifacts", + false, + (sess, ctx) -> { + sess.addArtifacts(artifact("doc1", "v1")); + sess.addArtifacts(artifact("doc1", "v2")); + sess.addArtifacts(artifact("doc2", "other")); + return done(); + })); + + // customAgentWithCustomState: client-managed, increments custom.counter. + register( + "customAgentWithCustomState", defineCustom("customAgentWithCustomState", false, COUNTER)); + + // customAgentWithMultiCustomState: client-managed, three sequential custom-state updates. + register( + "customAgentWithMultiCustomState", + defineCustom( + "customAgentWithMultiCustomState", + false, + (sess, ctx) -> { + sess.updateCustom(prev -> mapOf("counter", 1, "status", "working")); + sess.updateCustom( + prev -> { + Map out = new HashMap<>(prev); + out.put("counter", 2); + return out; + }); + sess.updateCustom( + prev -> { + Map out = new HashMap<>(prev); + out.put("status", "done"); + return out; + }); + return done(); + })); + + // customAgentWithArtifactsStore: server-managed, adds a numbered artifact per invocation. + register( + "customAgentWithArtifactsStore", + defineCustom( + "customAgentWithArtifactsStore", + true, + (sess, ctx) -> { + int count = sess.getArtifacts().size() + 1; + sess.addArtifacts(artifact("doc" + count, "content" + count)); + return done(); + })); + + // customAgentWithCustomStateStore: server-managed counter agent. + register( + "customAgentWithCustomStateStore", + defineCustom("customAgentWithCustomStateStore", true, COUNTER)); + } + + /** Counter agent func: increments custom.counter by 1 each turn (default 0 -> 1). */ + private static final AgentFn> COUNTER = + (sess, ctx) -> { + Map prev = sess.getCustom(); + long counter = 0; + if (prev != null && prev.get("counter") instanceof Number n) { + counter = n.longValue(); + } + long next = counter + 1; + sess.updateCustom(p -> mapOf("counter", next)); + return done(); + }; + + // ── helpers ────────────────────────────────────────────────────────────────────── + + private Agent> defineCustom( + String name, boolean serverManaged, AgentFn> fn) { + CustomAgentConfig.Builder> cfg = + CustomAgentConfig.>builder().name(name); + if (serverManaged) { + cfg.store(new InMemorySessionStore<>()); + } + return AgentActions.defineCustomAgent(genkit.getRegistry(), cfg.build(), fn); + } + + private void register(String name, Agent> agent) { + agents.put(name, agent); + } + + private static AgentResult done() { + return AgentResult.builder() + .message(Message.model("done")) + .finishReason(AgentFinishReason.STOP) + .build(); + } + + private static Artifact artifact(String name, String text) { + return Artifact.builder().name(name).parts(List.of(Part.text(text))).build(); + } + + private static Map mapOf(String k, Object v) { + Map m = new HashMap<>(); + m.put(k, v); + return m; + } + + private static Map mapOf(String k1, Object v1, String k2, Object v2) { + Map m = new HashMap<>(); + m.put(k1, v1); + m.put(k2, v2); + return m; + } +} diff --git a/genkit/src/test/java/com/google/genkit/conformance/agent/ProgrammableModel.java b/genkit/src/test/java/com/google/genkit/conformance/agent/ProgrammableModel.java new file mode 100644 index 000000000..aa7bad68b --- /dev/null +++ b/genkit/src/test/java/com/google/genkit/conformance/agent/ProgrammableModel.java @@ -0,0 +1,148 @@ +/* + * 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. + * + * SPDX-License-Identifier: Apache-2.0 + */ + +package com.google.genkit.conformance.agent; + +import com.fasterxml.jackson.databind.JsonNode; +import com.google.genkit.ai.Candidate; +import com.google.genkit.ai.FinishReason; +import com.google.genkit.ai.Message; +import com.google.genkit.ai.Model; +import com.google.genkit.ai.ModelInfo; +import com.google.genkit.ai.ModelRequest; +import com.google.genkit.ai.ModelResponse; +import com.google.genkit.ai.ModelResponseChunk; +import com.google.genkit.core.ActionContext; +import com.google.genkit.core.JsonUtils; +import java.util.List; +import java.util.function.Consumer; + +/** + * A {@link Model} whose per-call behaviour is programmed per {@code send} step. + * + *

Mirrors the JS {@code defineProgrammableModel} / Go {@code programmableModel} helpers: for the + * {@code i}-th generate call of a step it streams {@code streamChunks[i]} (if present) then returns + * {@code modelResponses[i]}. The request is echoed onto the response ({@link + * ModelResponse#setRequest}) so the generate tool loop can thread intermediate tool-request / + * tool-response messages into the session history (the prompt agent relies on this). + * + *

Responses and chunks are supplied as raw {@link JsonNode} (parsed straight from the spec YAML) + * and converted to typed objects on demand, so the harness never has to hand-build model types. + */ +final class ProgrammableModel implements Model { + + private final String name; + + /** Per-call responses, set by {@link #program} at the start of each {@code send} step. */ + private List modelResponses = List.of(); + + /** Per-call streaming chunks (outer index = call, inner list = chunks for that call). */ + private List> streamChunks = List.of(); + + /** How many generate calls have happened in the current step. */ + private int callCount; + + ProgrammableModel(String name) { + this.name = name; + } + + /** + * Programs the model for the next {@code send} step. Resets the per-step call counter. + * + * @param modelResponses one {@link ModelResponse}-shaped node per generate call (may be empty) + * @param streamChunks per-call lists of {@link ModelResponseChunk}-shaped nodes (may be empty) + */ + void program(List modelResponses, List> streamChunks) { + this.modelResponses = modelResponses != null ? modelResponses : List.of(); + this.streamChunks = streamChunks != null ? streamChunks : List.of(); + this.callCount = 0; + } + + @Override + public String getName() { + return name; + } + + @Override + public ModelInfo getInfo() { + ModelInfo info = new ModelInfo(); + ModelInfo.ModelCapabilities caps = new ModelInfo.ModelCapabilities(); + caps.setMultiturn(true); + caps.setTools(true); + caps.setSystemRole(true); + info.setSupports(caps); + return info; + } + + @Override + public boolean supportsStreaming() { + return true; + } + + @Override + public ModelResponse run(ActionContext ctx, ModelRequest request) { + return run(ctx, request, null); + } + + @Override + public ModelResponse run( + ActionContext ctx, ModelRequest request, Consumer streamCallback) { + int i = callCount++; + + if (streamCallback != null && i < streamChunks.size()) { + for (JsonNode chunkNode : streamChunks.get(i)) { + ModelResponseChunk chunk = JsonUtils.fromJsonNode(chunkNode, ModelResponseChunk.class); + streamCallback.accept(chunk); + } + } + + if (i >= modelResponses.size()) { + throw new IllegalStateException( + "programmableModel: no response programmed for generate call " + i); + } + + ModelResponse response = buildResponse(modelResponses.get(i)); + // Echo the request back, as every real model does. The tool loop relies on + // response.getRequest() + // to thread intermediate tool-request / tool-response messages into history. + response.setRequest(request); + return response; + } + + /** + * Builds a {@link ModelResponse} from a spec {@code GenerateResponse}-shaped node ({@code + * {message, finishReason}}). The Java {@link ModelResponse} carries the message under a {@code + * candidate}, so the message is lifted into a single {@link Candidate} (a plain {@code + * treeToValue} would silently drop the top-level {@code message} field). + */ + private static ModelResponse buildResponse(JsonNode node) { + Message message = null; + if (node.hasNonNull("message")) { + message = JsonUtils.fromJsonNode(node.get("message"), Message.class); + } + FinishReason finishReason = null; + if (node.hasNonNull("finishReason")) { + finishReason = + JsonUtils.getObjectMapper().convertValue(node.get("finishReason"), FinishReason.class); + } + Candidate candidate = new Candidate(message, finishReason); + ModelResponse response = new ModelResponse(List.of(candidate)); + response.setFinishReason(finishReason); + return response; + } +} diff --git a/genkit/src/test/resources/conformance/agent.yaml b/genkit/src/test/resources/conformance/agent.yaml new file mode 100644 index 000000000..1aab42efe --- /dev/null +++ b/genkit/src/test/resources/conformance/agent.yaml @@ -0,0 +1,1396 @@ +# Copyright 2025 Google LLC +# SPDX-License-Identifier: Apache-2.0 + +# This file describes the behavioral specification for the Agent API. +# It is designed to be consumed by conformance test harnesses in any +# language (JS, Go, Dart, Python, etc.) to ensure cross-language +# compatibility of the Agent abstraction. +# +# See docs/agents-conformance-testing.md for harness requirements and +# full spec format reference. + +tests: + # --------------------------------------------------------------------------- + # Basic single-turn + # --------------------------------------------------------------------------- + - name: simple single turn - client managed + description: > + A single user message is sent to a client-managed agent. + The agent generates one model response and returns it along with + the accumulated session state. A sessionId must be generated. + agent: promptAgent + steps: + - type: send + init: {} + inputs: + - message: { role: user, content: [{ text: hi }] } + modelResponses: + - message: { role: model, content: [{ text: hello back }] } + finishReason: stop + expectChunks: + # A normal completion carries the model's finishReason on turnEnd. + - turnEnd: { finishReason: stop } + expectOutput: + message: { role: model, content: [{ text: hello back }] } + # A normal completion surfaces finishReason 'stop' on the output. + finishReason: stop + hasSessionId: true + stateContains: + messages: + - { role: user, content: [{ text: hi }] } + - { role: model, content: [{ text: hello back }] } + + - name: simple single turn - server managed + description: > + A single user message is sent to a server-managed agent. + The output should contain a snapshotId and no inline state. + The snapshot state must contain a generated sessionId. + agent: promptAgentWithStore + steps: + - type: send + init: {} + inputs: + - message: { role: user, content: [{ text: hi }] } + modelResponses: + - message: { role: model, content: [{ text: hello back }] } + finishReason: stop + expectChunks: + - turnEnd: { finishReason: stop } + expectOutput: + message: { role: model, content: [{ text: hello back }] } + finishReason: stop + hasSnapshotId: true + captureSnapshotId: snap1 + + - type: getSnapshotData + snapshotId: '{{snap1}}' + expectSnapshot: + status: completed + hasSessionId: true + + # --------------------------------------------------------------------------- + # Streaming + # --------------------------------------------------------------------------- + - name: streaming model chunks + description: > + Model emits streaming chunks during generation. They should be + forwarded as modelChunk stream events, followed by a turnEnd. + agent: promptAgent + steps: + - type: send + init: {} + inputs: + - message: { role: user, content: [{ text: hi }] } + streamChunks: + - [ + { index: 0, role: model, content: [{ text: 'hel' }] }, + { index: 0, role: model, content: [{ text: 'lo' }] }, + ] + modelResponses: + - message: { role: model, content: [{ text: hello }] } + finishReason: stop + expectChunks: + - modelChunk: { index: 0, role: model, content: [{ text: 'hel' }] } + - modelChunk: { index: 0, role: model, content: [{ text: 'lo' }] } + - turnEnd: { finishReason: stop } + expectOutput: + message: { role: model, content: [{ text: hello }] } + finishReason: stop + stateContains: + messages: + - { role: user, content: [{ text: hi }] } + - { role: model, content: [{ text: hello }] } + + # --------------------------------------------------------------------------- + # Multi-turn in one invocation + # --------------------------------------------------------------------------- + - name: multi-turn in one invocation + description: > + Two user messages are sent sequentially in one invocation. + Both should be processed as separate turns. History must + accumulate so the second model call sees the first exchange. + agent: promptAgent + steps: + - type: send + init: {} + inputs: + - message: { role: user, content: [{ text: turn1 }] } + - message: { role: user, content: [{ text: turn2 }] } + modelResponses: + - message: { role: model, content: [{ text: reply1 }] } + finishReason: stop + - message: { role: model, content: [{ text: reply2 }] } + finishReason: stop + expectChunks: + # Each turn ends with the model's finishReason. + - turnEnd: { finishReason: stop } + - turnEnd: { finishReason: stop } + expectOutput: + message: { role: model, content: [{ text: reply2 }] } + finishReason: stop + stateContains: + messages: + - { role: user, content: [{ text: turn1 }] } + - { role: model, content: [{ text: reply1 }] } + - { role: user, content: [{ text: turn2 }] } + - { role: model, content: [{ text: reply2 }] } + + # --------------------------------------------------------------------------- + # Tool calling + # --------------------------------------------------------------------------- + - name: agent calls tools + description: > + Model issues a tool request, tool executes automatically, and + the tool response is fed back to the model which then produces + a final text response. + agent: promptAgentWithTools + steps: + - type: send + init: {} + inputs: + - message: { role: user, content: [{ text: do it }] } + modelResponses: + - message: + role: model + content: + - toolRequest: { name: testTool, input: {}, ref: ref1 } + finishReason: stop + - message: { role: model, content: [{ text: done }] } + finishReason: stop + expectChunks: + - modelChunk: + role: tool + content: + - toolResponse: + { name: testTool, ref: ref1, output: 'tool called' } + - turnEnd: {} + expectOutput: + message: { role: model, content: [{ text: done }] } + stateContains: + messages: + - { role: user, content: [{ text: do it }] } + - role: model + content: + - toolRequest: { name: testTool, input: {}, ref: ref1 } + - role: tool + content: + - toolResponse: + { name: testTool, output: 'tool called', ref: ref1 } + - { role: model, content: [{ text: done }] } + + # --------------------------------------------------------------------------- + # Interrupt and resume (multi-invocation) + # --------------------------------------------------------------------------- + - name: interrupt and resume + description: > + Model returns an interrupt tool request. The agent saves state + and returns the tool request as the output message. The client + then resumes with a new invocation providing the tool response + via snapshotId. The model receives the full history including + the tool response and produces a final answer. + agent: promptAgentWithInterrupt + steps: + # Phase 1: model interrupts + - type: send + init: {} + inputs: + - message: { role: user, content: [{ text: hello }] } + modelResponses: + - message: + role: model + content: + - toolRequest: + { + name: interruptTool, + input: { query: 'yes?' }, + ref: '123', + } + finishReason: stop + # A tool interrupt pauses the turn — finishReason is 'interrupted'. + expectOutput: + message: + content: + - toolRequest: + { name: interruptTool, input: { query: 'yes?' }, ref: '123' } + finishReason: interrupted + hasSnapshotId: true + captureSnapshotId: snap1 + + # Phase 2: client resumes with resume.respond + - type: send + init: { snapshotId: '{{snap1}}' } + inputs: + - resume: + respond: + - toolResponse: + { + name: interruptTool, + ref: '123', + output: { answer: 'yes indeed' }, + } + modelResponses: + - message: { role: model, content: [{ text: completed }] } + finishReason: stop + expectOutput: + message: { role: model, content: [{ text: completed }] } + + # --------------------------------------------------------------------------- + # Interrupt and restart (resume.restart) + # --------------------------------------------------------------------------- + - name: interrupt and restart + description: > + Model requests a tool that throws ToolInterruptError on first call. + The agent saves state and returns the tool request as output. The + client resumes with resume.restart (same input + metadata). The + tool re-executes successfully with the resumed metadata and the + model produces a final answer. + agent: promptAgentWithRestartTool + steps: + # Phase 1: model requests restartTool, tool interrupts + - type: send + init: {} + inputs: + - message: { role: user, content: [{ text: do it }] } + modelResponses: + - message: + role: model + content: + - toolRequest: + { + name: restartTool, + input: { action: 'delete' }, + ref: 'r1', + } + finishReason: stop + # ToolInterruptError pauses the turn — finishReason is 'interrupted'. + expectOutput: + message: + content: + - toolRequest: + { name: restartTool, input: { action: 'delete' }, ref: 'r1' } + finishReason: interrupted + hasSnapshotId: true + captureSnapshotId: snap1 + + # Phase 2: client resumes with resume.restart + - type: send + init: { snapshotId: '{{snap1}}' } + inputs: + - resume: + restart: + - toolRequest: + { + name: restartTool, + input: { action: 'delete' }, + ref: 'r1', + } + metadata: { resumed: { approved: true } } + modelResponses: + - message: { role: model, content: [{ text: deleted }] } + finishReason: stop + expectOutput: + message: { role: model, content: [{ text: deleted }] } + + # --------------------------------------------------------------------------- + # Resume validation — forged restart rejected + # --------------------------------------------------------------------------- + - name: restart with forged inputs rejected + description: > + A malicious client attempts to restart a tool with modified inputs. + The agent must reject the restart because the input does not match + the original tool request in session history. + agent: promptAgentWithRestartTool + steps: + # Phase 1: model requests restartTool with safe input + - type: send + init: {} + inputs: + - message: { role: user, content: [{ text: do it }] } + modelResponses: + - message: + role: model + content: + - toolRequest: + { + name: restartTool, + input: { action: 'safe-op' }, + ref: 'r1', + } + finishReason: stop + captureSnapshotId: snap1 + + # Phase 2: client forges restart with different input + - type: send + init: { snapshotId: '{{snap1}}' } + inputs: + - resume: + restart: + - toolRequest: + { + name: restartTool, + input: { action: '/etc/passwd' }, + ref: 'r1', + } + metadata: { resumed: { approved: true } } + modelResponses: + - message: { role: model, content: [{ text: should not reach }] } + finishReason: stop + # The agent does not throw — it resolves gracefully with + # finishReason 'failed', preserving the original error status. + expectOutput: + finishReason: failed + errorContains: + status: INVALID_ARGUMENT + message: modified inputs + + # --------------------------------------------------------------------------- + # Resume validation — respond referencing non-existent tool + # --------------------------------------------------------------------------- + - name: respond referencing non-existent tool rejected + description: > + A client attempts to respond with a tool name/ref that does not + match any tool request in the session history. The agent must + reject with INVALID_ARGUMENT. + agent: promptAgentWithInterrupt + steps: + # Phase 1: model requests interruptTool + - type: send + init: {} + inputs: + - message: { role: user, content: [{ text: hello }] } + modelResponses: + - message: + role: model + content: + - toolRequest: + { + name: interruptTool, + input: { query: 'confirm?' }, + ref: 'i1', + } + finishReason: stop + captureSnapshotId: snap1 + + # Phase 2: client responds with a fabricated tool name + - type: send + init: { snapshotId: '{{snap1}}' } + inputs: + - resume: + respond: + - toolResponse: + { + name: fakeTool, + ref: 'fake-ref', + output: { answer: 'hacked' }, + } + modelResponses: + - message: { role: model, content: [{ text: should not reach }] } + finishReason: stop + # Resolves gracefully with finishReason 'failed' and INVALID_ARGUMENT. + expectOutput: + finishReason: failed + errorContains: + status: INVALID_ARGUMENT + message: not found in session history + + # --------------------------------------------------------------------------- + # Snapshot chaining + # --------------------------------------------------------------------------- + - name: snapshot chaining across invocations + description: > + Two sequential invocations against a server-managed agent. + The second invocation resumes from the first snapshot. After + both complete, getSnapshotData verifies the parent chain, + accumulated history, and that a sessionId is present. + agent: promptAgentWithStore + steps: + - type: send + init: {} + inputs: + - message: { role: user, content: [{ text: first }] } + modelResponses: + - message: { role: model, content: [{ text: reply1 }] } + finishReason: stop + captureSnapshotId: snap1 + + - type: send + init: { snapshotId: '{{snap1}}' } + inputs: + - message: { role: user, content: [{ text: second }] } + modelResponses: + - message: { role: model, content: [{ text: reply2 }] } + finishReason: stop + captureSnapshotId: snap2 + + - type: getSnapshotData + snapshotId: '{{snap2}}' + expectSnapshot: + parentId: '{{snap1}}' + status: completed + hasSessionId: true + stateContains: + messages: + - { role: user, content: [{ text: first }] } + - { role: model, content: [{ text: reply1 }] } + - { role: user, content: [{ text: second }] } + - { role: model, content: [{ text: reply2 }] } + + # --------------------------------------------------------------------------- + # Client-managed state across invocations + # --------------------------------------------------------------------------- + - name: client-managed state across invocations + description: > + Two sequential invocations against a client-managed agent. + The second invocation seeds session state from the first + invocation's output state. History should accumulate and the + sessionId must be preserved across invocations. + agent: promptAgent + steps: + - type: send + init: {} + inputs: + - message: { role: user, content: [{ text: first }] } + modelResponses: + - message: { role: model, content: [{ text: reply1 }] } + finishReason: stop + captureState: state1 + captureSessionId: sid1 + + - type: send + init: { state: '{{state1}}' } + inputs: + - message: { role: user, content: [{ text: second }] } + modelResponses: + - message: { role: model, content: [{ text: reply2 }] } + finishReason: stop + expectOutput: + message: { role: model, content: [{ text: reply2 }] } + stateContains: + sessionId: '{{sid1}}' + messages: + - { role: user, content: [{ text: first }] } + - { role: model, content: [{ text: reply1 }] } + - { role: user, content: [{ text: second }] } + - { role: model, content: [{ text: reply2 }] } + + # =========================================================================== + # Phase 2: Detach, Abort, Artifacts, Custom State + # =========================================================================== + + # --------------------------------------------------------------------------- + # Detach & background execution + # --------------------------------------------------------------------------- + - name: detach and background completion + description: > + A detach flag causes the agent to return immediately with a + pending snapshot. The background continues processing and + eventually the snapshot reaches "done" status with accumulated + state. + agent: promptAgentWithStore + steps: + - type: send + init: {} + inputs: + - message: { role: user, content: [{ text: process this }] } + detach: true + modelResponses: + - message: { role: model, content: [{ text: done in background }] } + finishReason: stop + expectOutput: + hasSnapshotId: true + captureSnapshotId: snap1 + + - type: waitUntilCompleted + snapshotId: '{{snap1}}' + expectSnapshot: + status: completed + stateContains: + messages: + - { role: user, content: [{ text: process this }] } + - { role: model, content: [{ text: done in background }] } + + - name: detach with background failure + description: > + When a detached agent fails in the background, the snapshot + status should be set to "failed". + agent: customAgentFailing + steps: + - type: send + init: {} + inputs: + - message: { role: user, content: [{ text: fail }] } + detach: true + expectOutput: + hasSnapshotId: true + captureSnapshotId: snap1 + + - type: waitUntilCompleted + snapshotId: '{{snap1}}' + expectSnapshot: + status: failed + # A failed run records finishReason 'failed' on the snapshot, + # distinct from the snapshot status. + finishReason: failed + + # --------------------------------------------------------------------------- + # Abort + # --------------------------------------------------------------------------- + - name: abort pending agent + description: > + Abort a detached agent that is still processing. The abort + should return "pending" as the previous status and the snapshot + should be set to "aborted". + agent: customAgentBlocking + steps: + - type: send + init: {} + inputs: + - message: { role: user, content: [{ text: do work }] } + detach: true + expectOutput: + hasSnapshotId: true + captureSnapshotId: snap1 + + - type: abort + snapshotId: '{{snap1}}' + expectPreviousStatus: pending + + - type: getSnapshotData + snapshotId: '{{snap1}}' + expectSnapshot: + status: aborted + + - name: abort completed agent + description: > + Abort an agent that has already finished. The abort returns + "done" as previous status but the snapshot remains "done" + because terminal states (done, failed, aborted) cannot be + overridden — only "pending" can transition to "aborted". + agent: promptAgentWithStore + steps: + - type: send + init: {} + inputs: + - message: { role: user, content: [{ text: hi }] } + modelResponses: + - message: { role: model, content: [{ text: hello }] } + finishReason: stop + captureSnapshotId: snap1 + + - type: abort + snapshotId: '{{snap1}}' + expectPreviousStatus: completed + + - type: getSnapshotData + snapshotId: '{{snap1}}' + expectSnapshot: + status: completed + + - name: server-managed agent rejects init state + description: > + When a server-managed agent receives init.state, it must throw a + FAILED_PRECONDITION error (mapped to an HTTP status by the server + handler). Server-managed agents expect a snapshotId, not the full state + blob. This is API misuse, so it is a thrown error rather than a graceful + 'failed' output. + agent: promptAgentWithStore + steps: + - type: send + init: + state: + messages: + - { role: user, content: [{ text: stale history }] } + custom: { shouldBeIgnored: true } + artifacts: [] + inputs: + - message: { role: user, content: [{ text: fresh message }] } + modelResponses: + - message: { role: model, content: [{ text: reply }] } + finishReason: stop + # API misuse: the turn throws rather than resolving with a graceful + # 'failed' output. + expectError: + status: FAILED_PRECONDITION + message: Cannot send 'state' to agent + + - name: pure detach without payload + description: > + A detach-only message (no messages or resume payloads) sent as a + separate input should detach immediately without processing an + extra turn. The agent should return a pending snapshot. + agent: customAgentBlocking + steps: + - type: send + init: {} + inputs: + - message: { role: user, content: [{ text: work }] } + - detach: true + expectOutput: + hasSnapshotId: true + captureSnapshotId: snap1 + + - type: getSnapshotData + snapshotId: '{{snap1}}' + expectSnapshot: + status: pending + + - type: abort + snapshotId: '{{snap1}}' + expectPreviousStatus: pending + + - name: failed snapshot includes error details + description: > + When a detached agent fails in the background, the snapshot + should include error details with the failure message. + agent: customAgentFailing + steps: + - type: send + init: {} + inputs: + - message: { role: user, content: [{ text: fail }] } + detach: true + expectOutput: + hasSnapshotId: true + captureSnapshotId: snap1 + + - type: waitUntilCompleted + snapshotId: '{{snap1}}' + expectSnapshot: + status: failed + finishReason: failed + errorContains: + message: intentional failure + + - name: abort non-existent snapshot + description: > + Aborting a snapshot that does not exist should not throw and + should return no previous status. + agent: promptAgentWithStore + steps: + - type: abort + snapshotId: non-existent-id + expectPreviousStatus: ~ + + # --------------------------------------------------------------------------- + # Artifacts + # --------------------------------------------------------------------------- + - name: artifacts streamed and deduplicated + description: > + Custom agent adds artifacts during execution. Artifact chunks + are streamed in order. Same-named artifacts are deduplicated + (updated) so the final output contains only the latest version. + agent: customAgentWithArtifacts + steps: + - type: send + init: {} + inputs: + - message: { role: user, content: [{ text: go }] } + expectChunks: + - artifact: { name: doc1, parts: [{ text: v1 }] } + - artifact: { name: doc1, parts: [{ text: v2 }] } + - artifact: { name: doc2, parts: [{ text: other }] } + - turnEnd: {} + expectOutput: + message: { role: model, content: [{ text: done }] } + artifactsContain: + - name: doc1 + parts: [{ text: v2 }] + - name: doc2 + parts: [{ text: other }] + + # --------------------------------------------------------------------------- + # Custom state + # --------------------------------------------------------------------------- + - name: custom state updated during execution + description: > + Custom agent updates custom state during processing. The output + state should contain the updated custom data. + agent: customAgentWithCustomState + steps: + - type: send + init: {} + inputs: + - message: { role: user, content: [{ text: go }] } + expectOutput: + message: { role: model, content: [{ text: done }] } + stateContains: + custom: { counter: 1 } + + - name: custom state persisted across invocations + description: > + Custom state is preserved when seeded from a previous + invocation's output. The counter should increment across + invocations. + agent: customAgentWithCustomState + steps: + - type: send + init: {} + inputs: + - message: { role: user, content: [{ text: first }] } + expectOutput: + stateContains: + custom: { counter: 1 } + captureState: state1 + + - type: send + init: { state: '{{state1}}' } + inputs: + - message: { role: user, content: [{ text: second }] } + expectOutput: + stateContains: + custom: { counter: 2 } + messages: + - { role: user, content: [{ text: first }] } + - { role: user, content: [{ text: second }] } + + # --------------------------------------------------------------------------- + # Custom state streamed live as customPatch chunks + # --------------------------------------------------------------------------- + - name: custom state streamed as customPatch chunk + description: > + A single custom-state mutation during a turn is auto-emitted to the + client as a `customPatch` stream chunk. The first (and here only) patch + of a turn is a whole-document replace at the root pointer ('') carrying + the full transformed custom state. The final output state still reflects + the same custom data. + agent: customAgentWithCustomState + steps: + - type: send + init: {} + inputs: + - message: { role: user, content: [{ text: go }] } + expectChunks: + - customPatch: + - { op: replace, path: '', value: { counter: 1 } } + - turnEnd: {} + expectOutput: + message: { role: model, content: [{ text: done }] } + stateContains: + custom: { counter: 1 } + + - name: customPatch first chunk is whole-document replace then incremental + description: > + Multiple custom-state mutations within a single turn produce multiple + `customPatch` chunks. The first patch is a whole-document replace at the + root pointer ('') re-basing the client with the full custom state. + Subsequent patches are incremental RFC 6902 diffs targeting only the + changed members (matched on op + path; values are partially asserted). + The final output state contains the fully accumulated custom data. + agent: customAgentWithMultiCustomState + steps: + - type: send + init: {} + inputs: + - message: { role: user, content: [{ text: go }] } + expectChunks: + # First mutation -> whole-document replace at root. + - customPatch: + - op: replace + path: '' + value: { counter: 1, status: working } + # Second mutation -> incremental replace of /counter. + - customPatch: + - op: replace + path: /counter + value: 2 + # Third mutation -> incremental replace of /status. + - customPatch: + - op: replace + path: /status + value: done + - turnEnd: {} + expectOutput: + message: { role: model, content: [{ text: done }] } + stateContains: + custom: { counter: 2, status: done } + + - name: detached run emits no customPatch chunks + description: > + A detached run streams no chunks to the connection (it returns a pending + snapshot immediately and continues in the background). In particular no + `customPatch` chunks are emitted even though the agent mutates custom + state. The mutation is still persisted and observable on the completed + snapshot. + agent: customAgentWithCustomStateStore + steps: + - type: send + init: {} + inputs: + - message: { role: user, content: [{ text: go }] } + detach: true + # No customPatch (or any) chunks are streamed for a detached run. + expectChunks: [] + expectOutput: + hasSnapshotId: true + captureSnapshotId: snap1 + + - type: waitUntilCompleted + snapshotId: '{{snap1}}' + expectSnapshot: + status: completed + stateContains: + custom: { counter: 1 } + + # =========================================================================== + # Phase 3: Additional API Coverage + # =========================================================================== + + # --------------------------------------------------------------------------- + # Snapshot branching + # --------------------------------------------------------------------------- + - name: snapshot branching across invocations + description: > + Two different continuations from the same snapshot produce + independent histories. Each child snapshot must have the same + parentId but divergent message histories after the branch point. + agent: promptAgentWithStore + steps: + - type: send + init: {} + inputs: + - message: { role: user, content: [{ text: root }] } + modelResponses: + - message: { role: model, content: [{ text: rootReply }] } + finishReason: stop + captureSnapshotId: snap1 + + - type: send + init: { snapshotId: '{{snap1}}' } + inputs: + - message: { role: user, content: [{ text: branch-a }] } + modelResponses: + - message: { role: model, content: [{ text: reply-a }] } + finishReason: stop + captureSnapshotId: snap2a + + - type: send + init: { snapshotId: '{{snap1}}' } + inputs: + - message: { role: user, content: [{ text: branch-b }] } + modelResponses: + - message: { role: model, content: [{ text: reply-b }] } + finishReason: stop + captureSnapshotId: snap2b + + - type: getSnapshotData + snapshotId: '{{snap2a}}' + expectSnapshot: + parentId: '{{snap1}}' + status: completed + stateContains: + messages: + - { role: user, content: [{ text: root }] } + - { role: model, content: [{ text: rootReply }] } + - { role: user, content: [{ text: branch-a }] } + - { role: model, content: [{ text: reply-a }] } + + - type: getSnapshotData + snapshotId: '{{snap2b}}' + expectSnapshot: + parentId: '{{snap1}}' + status: completed + stateContains: + messages: + - { role: user, content: [{ text: root }] } + - { role: model, content: [{ text: rootReply }] } + - { role: user, content: [{ text: branch-b }] } + - { role: model, content: [{ text: reply-b }] } + + # --------------------------------------------------------------------------- + # Multiple tool calls in one model response + # --------------------------------------------------------------------------- + - name: multiple tool calls in one model response + description: > + Model issues two tool requests in a single message. Both tools + execute automatically and their responses are fed back to the + model which then produces a final text response. + agent: promptAgentWithTools + steps: + - type: send + init: {} + inputs: + - message: { role: user, content: [{ text: do both }] } + modelResponses: + - message: + role: model + content: + - toolRequest: { name: testTool, input: {}, ref: ref1 } + - toolRequest: { name: testTool, input: {}, ref: ref2 } + finishReason: stop + - message: { role: model, content: [{ text: all done }] } + finishReason: stop + expectOutput: + message: { role: model, content: [{ text: all done }] } + stateContains: + messages: + - { role: user, content: [{ text: do both }] } + - role: model + content: + - toolRequest: { name: testTool, input: {}, ref: ref1 } + - toolRequest: { name: testTool, input: {}, ref: ref2 } + - role: tool + content: + - toolResponse: + { name: testTool, output: 'tool called', ref: ref1 } + - toolResponse: + { name: testTool, output: 'tool called', ref: ref2 } + - { role: model, content: [{ text: all done }] } + + # --------------------------------------------------------------------------- + # Multiple interrupt tool requests + # --------------------------------------------------------------------------- + - name: interrupt with multiple tool requests + description: > + Model returns two interrupt tool requests in a single message. + The agent returns both tool requests to the client as the output + message. The client resumes by providing responses for both tools. + agent: promptAgentWithInterrupt + steps: + # Phase 1: model returns two interrupt requests + - type: send + init: {} + inputs: + - message: { role: user, content: [{ text: confirm both }] } + modelResponses: + - message: + role: model + content: + - toolRequest: + { name: interruptTool, input: { query: 'q1?' }, ref: 'i1' } + - toolRequest: + { name: interruptTool, input: { query: 'q2?' }, ref: 'i2' } + finishReason: stop + expectOutput: + message: + content: + - toolRequest: + { name: interruptTool, input: { query: 'q1?' }, ref: 'i1' } + - toolRequest: + { name: interruptTool, input: { query: 'q2?' }, ref: 'i2' } + hasSnapshotId: true + captureSnapshotId: snap1 + + # Phase 2: client responds to both + - type: send + init: { snapshotId: '{{snap1}}' } + inputs: + - resume: + respond: + - toolResponse: + { + name: interruptTool, + ref: 'i1', + output: { answer: 'yes1' }, + } + - toolResponse: + { + name: interruptTool, + ref: 'i2', + output: { answer: 'yes2' }, + } + modelResponses: + - message: { role: model, content: [{ text: both confirmed }] } + finishReason: stop + expectOutput: + message: { role: model, content: [{ text: both confirmed }] } + + # --------------------------------------------------------------------------- + # Interrupt resume — full state accumulation + # --------------------------------------------------------------------------- + - name: interrupt resume state accumulation + description: > + After interrupt and resume, the accumulated state must contain + the full exchange: user message, model tool request, client + tool response, and the final model response. + agent: promptAgentWithInterrupt + steps: + - type: send + init: {} + inputs: + - message: { role: user, content: [{ text: check }] } + modelResponses: + - message: + role: model + content: + - toolRequest: + { name: interruptTool, input: { query: 'ok?' }, ref: 't1' } + finishReason: stop + captureSnapshotId: snap1 + + - type: send + init: { snapshotId: '{{snap1}}' } + inputs: + - resume: + respond: + - toolResponse: + { + name: interruptTool, + ref: 't1', + output: { answer: 'confirmed' }, + } + modelResponses: + - message: { role: model, content: [{ text: done }] } + finishReason: stop + captureSnapshotId: snap2 + + - type: getSnapshotData + snapshotId: '{{snap2}}' + expectSnapshot: + status: completed + stateContains: + messages: + - { role: user, content: [{ text: check }] } + - role: model + content: + - toolRequest: + { + name: interruptTool, + input: { query: 'ok?' }, + ref: 't1', + } + - role: tool + content: + - toolResponse: + { + name: interruptTool, + ref: 't1', + output: { answer: 'confirmed' }, + } + - { role: model, content: [{ text: done }] } + + # --------------------------------------------------------------------------- + # Artifacts across invocations (server-managed) + # --------------------------------------------------------------------------- + - name: artifacts persisted across invocations + description: > + Artifacts added in the first invocation persist in the server-managed + snapshot. The second invocation loads the snapshot, adds a new artifact, + and the output contains both the original and new artifacts. + agent: customAgentWithArtifactsStore + steps: + - type: send + init: {} + inputs: + - message: { role: user, content: [{ text: first }] } + expectOutput: + hasSnapshotId: true + artifactsContain: + - name: doc1 + parts: [{ text: content1 }] + captureSnapshotId: snap1 + + - type: send + init: { snapshotId: '{{snap1}}' } + inputs: + - message: { role: user, content: [{ text: second }] } + expectOutput: + artifactsContain: + - name: doc1 + parts: [{ text: content1 }] + - name: doc2 + parts: [{ text: content2 }] + + # --------------------------------------------------------------------------- + # Custom state via server-managed store + # --------------------------------------------------------------------------- + - name: custom state persisted via server-managed store + description: > + Custom state is preserved when resuming from a server-managed + snapshot. The counter should increment across invocations and + be visible in the snapshot state. + agent: customAgentWithCustomStateStore + steps: + - type: send + init: {} + inputs: + - message: { role: user, content: [{ text: first }] } + captureSnapshotId: snap1 + + - type: send + init: { snapshotId: '{{snap1}}' } + inputs: + - message: { role: user, content: [{ text: second }] } + captureSnapshotId: snap2 + + - type: getSnapshotData + snapshotId: '{{snap2}}' + expectSnapshot: + status: completed + stateContains: + custom: { counter: 2 } + messages: + - { role: user, content: [{ text: first }] } + - { role: user, content: [{ text: second }] } + + # --------------------------------------------------------------------------- + # Abort terminal states — failed and aborted + # --------------------------------------------------------------------------- + - name: abort failed agent + description: > + Abort an agent that has already failed. The abort returns "failed" + as the previous status but the snapshot remains "failed" because + terminal states cannot be overridden. + agent: customAgentFailing + steps: + - type: send + init: {} + inputs: + - message: { role: user, content: [{ text: fail }] } + detach: true + expectOutput: + hasSnapshotId: true + captureSnapshotId: snap1 + + - type: waitUntilCompleted + snapshotId: '{{snap1}}' + expectSnapshot: + status: failed + finishReason: failed + + - type: abort + snapshotId: '{{snap1}}' + expectPreviousStatus: failed + + - type: getSnapshotData + snapshotId: '{{snap1}}' + expectSnapshot: + status: failed + + - name: abort already aborted agent + description: > + Abort an agent that was already aborted. The abort returns + "aborted" as previous status and the snapshot remains "aborted". + agent: customAgentBlocking + steps: + - type: send + init: {} + inputs: + - message: { role: user, content: [{ text: work }] } + detach: true + expectOutput: + hasSnapshotId: true + captureSnapshotId: snap1 + + - type: abort + snapshotId: '{{snap1}}' + expectPreviousStatus: pending + + - type: abort + snapshotId: '{{snap1}}' + expectPreviousStatus: aborted + + - type: getSnapshotData + snapshotId: '{{snap1}}' + expectSnapshot: + status: aborted + + # =========================================================================== + # Phase 4: Server-managed sessions by sessionId + # =========================================================================== + # + # Server-managed agents can be resumed by a bare `sessionId` (a UUID) in + # addition to an exact `snapshotId`. This is the simple case used by + # `useChat`-style clients: the client tracks only a stable session id and the + # store resolves the session's latest (leaf) snapshot on each turn. + + # --------------------------------------------------------------------------- + # Resume a server-managed session by sessionId + # --------------------------------------------------------------------------- + - name: resume server-managed session by sessionId + description: > + A server-managed agent is invoked twice using the same caller-provided + sessionId (a UUID). The first turn seeds a fresh session bound to that + sessionId; the second turn resumes the session's latest snapshot without a + snapshotId. History must accumulate and the resulting snapshot must carry + the same sessionId and chain to the first snapshot. + agent: promptAgentWithStore + steps: + - type: send + init: { sessionId: 11111111-1111-4111-8111-111111111111 } + inputs: + - message: { role: user, content: [{ text: first }] } + modelResponses: + - message: { role: model, content: [{ text: reply1 }] } + finishReason: stop + expectOutput: + hasSnapshotId: true + captureSnapshotId: snap1 + + - type: send + init: { sessionId: 11111111-1111-4111-8111-111111111111 } + inputs: + - message: { role: user, content: [{ text: second }] } + modelResponses: + - message: { role: model, content: [{ text: reply2 }] } + finishReason: stop + captureSnapshotId: snap2 + + - type: getSnapshotData + snapshotId: '{{snap2}}' + expectSnapshot: + parentId: '{{snap1}}' + status: completed + stateContains: + sessionId: 11111111-1111-4111-8111-111111111111 + messages: + - { role: user, content: [{ text: first }] } + - { role: model, content: [{ text: reply1 }] } + - { role: user, content: [{ text: second }] } + - { role: model, content: [{ text: reply2 }] } + + # --------------------------------------------------------------------------- + # Fetch the latest snapshot by sessionId + # --------------------------------------------------------------------------- + - name: fetch latest snapshot by sessionId + description: > + getSnapshotData can resolve a snapshot by sessionId, returning the + session's latest (leaf) snapshot. After two linear turns the leaf is the + second snapshot, whose parent is the first. + agent: promptAgentWithStore + steps: + - type: send + init: { sessionId: 22222222-2222-4222-8222-222222222222 } + inputs: + - message: { role: user, content: [{ text: first }] } + modelResponses: + - message: { role: model, content: [{ text: reply1 }] } + finishReason: stop + captureSnapshotId: snap1 + + - type: send + init: { sessionId: 22222222-2222-4222-8222-222222222222 } + inputs: + - message: { role: user, content: [{ text: second }] } + modelResponses: + - message: { role: model, content: [{ text: reply2 }] } + finishReason: stop + + - type: getSnapshotData + sessionId: 22222222-2222-4222-8222-222222222222 + expectSnapshot: + parentId: '{{snap1}}' + status: completed + hasSessionId: true + stateContains: + sessionId: 22222222-2222-4222-8222-222222222222 + messages: + - { role: user, content: [{ text: first }] } + - { role: model, content: [{ text: reply1 }] } + - { role: user, content: [{ text: second }] } + - { role: model, content: [{ text: reply2 }] } + + # --------------------------------------------------------------------------- + # Non-UUID sessionId accepted + # --------------------------------------------------------------------------- + - name: non-UUID sessionId accepted + description: > + A sessionId can be any non-empty string (not necessarily a UUID). + Sending an application-specific (non-UUID) sessionId to a server-managed + agent is accepted; the session is bound to that id and is resumable by it. + agent: promptAgentWithStore + steps: + - type: send + init: { sessionId: my-app-session-123 } + inputs: + - message: { role: user, content: [{ text: first }] } + modelResponses: + - message: { role: model, content: [{ text: reply1 }] } + finishReason: stop + expectOutput: + finishReason: stop + hasSnapshotId: true + captureSnapshotId: snap1 + + - type: getSnapshotData + sessionId: my-app-session-123 + expectSnapshot: + status: completed + stateContains: + sessionId: my-app-session-123 + messages: + - { role: user, content: [{ text: first }] } + - { role: model, content: [{ text: reply1 }] } + + # --------------------------------------------------------------------------- + # Client-managed agent rejects sessionId + # --------------------------------------------------------------------------- + - name: client-managed agent rejects sessionId + description: > + A client-managed agent (no store) cannot resume by sessionId because it + has nowhere to load the session from. Sending init.sessionId must be + rejected with FAILED_PRECONDITION. + agent: promptAgent + steps: + - type: send + init: { sessionId: 44444444-4444-4444-8444-444444444444 } + inputs: + - message: { role: user, content: [{ text: hi }] } + modelResponses: + - message: { role: model, content: [{ text: hello }] } + finishReason: stop + # API misuse: the turn throws with the original FAILED_PRECONDITION + # status rather than resolving with a graceful 'failed' output. + expectError: + status: FAILED_PRECONDITION + message: "Cannot use 'sessionId'" + + # --------------------------------------------------------------------------- + # snapshotId + sessionId: snapshotId selects, sessionId guards ownership + # --------------------------------------------------------------------------- + - name: snapshotId and matching sessionId together resume + description: > + init may carry both a snapshotId (the exact snapshot to resume) and a + sessionId. When the snapshot belongs to that session, the sessionId acts + as an ownership guard and the resume proceeds normally. + agent: promptAgentWithStore + steps: + - type: send + init: { sessionId: 55555555-5555-4555-8555-555555555555 } + inputs: + - message: { role: user, content: [{ text: first }] } + modelResponses: + - message: { role: model, content: [{ text: reply1 }] } + finishReason: stop + captureSnapshotId: snap1 + + - type: send + init: + snapshotId: '{{snap1}}' + sessionId: 55555555-5555-4555-8555-555555555555 + inputs: + - message: { role: user, content: [{ text: second }] } + modelResponses: + - message: { role: model, content: [{ text: reply2 }] } + finishReason: stop + expectOutput: + finishReason: stop + + - name: snapshotId with mismatched sessionId rejected + description: > + When init carries both a snapshotId and a sessionId, the snapshot must + belong to that session. A mismatch is rejected. + agent: promptAgentWithStore + steps: + - type: send + init: { sessionId: 55555555-5555-4555-8555-555555555555 } + inputs: + - message: { role: user, content: [{ text: first }] } + modelResponses: + - message: { role: model, content: [{ text: reply1 }] } + finishReason: stop + captureSnapshotId: snap1 + + - type: send + init: + snapshotId: '{{snap1}}' + sessionId: 99999999-9999-4999-8999-999999999999 + inputs: + - message: { role: user, content: [{ text: second }] } + modelResponses: + - message: { role: model, content: [{ text: reply2 }] } + finishReason: stop + # API misuse: the turn throws with the original INVALID_ARGUMENT status + # rather than resolving with a graceful 'failed' output. + expectError: + status: INVALID_ARGUMENT + message: 'does not belong to session' diff --git a/genkit/src/test/resources/logback-test.xml b/genkit/src/test/resources/logback-test.xml new file mode 100644 index 000000000..0561a8ec1 --- /dev/null +++ b/genkit/src/test/resources/logback-test.xml @@ -0,0 +1,18 @@ + + + + + + %d{HH:mm:ss.SSS} %-5level %logger{20} - %msg%n + + + + + + diff --git a/genkit/src/test/resources/prompts/topicAgent.prompt b/genkit/src/test/resources/prompts/topicAgent.prompt new file mode 100644 index 000000000..df9f13ede --- /dev/null +++ b/genkit/src/test/resources/prompts/topicAgent.prompt @@ -0,0 +1,4 @@ +--- +model: echoModel +--- +You are an expert on {{topic}}. Answer questions about {{topic}} for {{userName}}. diff --git a/plugins/aws-bedrock/pom.xml b/plugins/aws-bedrock/pom.xml index 91616b099..17fb43ba1 100644 --- a/plugins/aws-bedrock/pom.xml +++ b/plugins/aws-bedrock/pom.xml @@ -45,6 +45,16 @@ genkit ${project.version} + + com.google.genkit + genkit-core + ${project.version} + + + com.google.genkit + genkit-ai + ${project.version} + @@ -52,19 +62,26 @@ auth ${aws.sdk.version} - + software.amazon.awssdk http-client-spi ${aws.sdk.version} - + software.amazon.awssdk regions ${aws.sdk.version} + + + software.amazon.awssdk + dynamodb + ${aws.sdk.version} + + com.squareup.okhttp3 diff --git a/plugins/aws-bedrock/src/main/java/com/google/genkit/plugins/awsbedrock/AwsBedrockPlugin.java b/plugins/aws-bedrock/src/main/java/com/google/genkit/plugins/awsbedrock/AwsBedrockPlugin.java index 961c2995c..fa3f0a066 100644 --- a/plugins/aws-bedrock/src/main/java/com/google/genkit/plugins/awsbedrock/AwsBedrockPlugin.java +++ b/plugins/aws-bedrock/src/main/java/com/google/genkit/plugins/awsbedrock/AwsBedrockPlugin.java @@ -77,6 +77,15 @@ public class AwsBedrockPlugin implements Plugin { "anthropic.claude-haiku-4-5-20251001-v1:0", "anthropic.claude-opus-4-1-20250805-v1:0", "anthropic.claude-opus-4-5-20251101-v1:0", + // Anthropic Claude 4.6+ / 5 models (INFERENCE_PROFILE; short-form IDs). + // In regions without in-region support, use the geo inference profile id instead, + // e.g. "aws-bedrock/us.anthropic.claude-opus-4-8". + "anthropic.claude-sonnet-5", + "anthropic.claude-fable-5", + "anthropic.claude-opus-4-8", + "anthropic.claude-opus-4-7", + "anthropic.claude-opus-4-6-v1", + "anthropic.claude-sonnet-4-6", // Anthropic Claude 3.x models "anthropic.claude-3-7-sonnet-20250219-v1:0", "anthropic.claude-3-5-sonnet-20241022-v2:0", diff --git a/plugins/aws-bedrock/src/main/java/com/google/genkit/plugins/awsbedrock/session/DynamoDbSessionStore.java b/plugins/aws-bedrock/src/main/java/com/google/genkit/plugins/awsbedrock/session/DynamoDbSessionStore.java new file mode 100644 index 000000000..af7d57bb2 --- /dev/null +++ b/plugins/aws-bedrock/src/main/java/com/google/genkit/plugins/awsbedrock/session/DynamoDbSessionStore.java @@ -0,0 +1,739 @@ +/* + * 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. + * + * SPDX-License-Identifier: Apache-2.0 + */ + +package com.google.genkit.plugins.awsbedrock.session; + +import com.fasterxml.jackson.databind.JsonNode; +import com.fasterxml.jackson.databind.ObjectMapper; +import com.fasterxml.jackson.databind.node.NullNode; +import com.google.genkit.ai.agent.AgentFinishReason; +import com.google.genkit.ai.agent.GetSnapshotOptions; +import com.google.genkit.ai.agent.RuntimeError; +import com.google.genkit.ai.agent.SessionSnapshot; +import com.google.genkit.ai.agent.SessionState; +import com.google.genkit.ai.agent.SessionStore; +import com.google.genkit.ai.agent.SessionStoreOptions; +import com.google.genkit.ai.agent.SnapshotMutator; +import com.google.genkit.ai.agent.SnapshotStatus; +import com.google.genkit.ai.agent.SnapshotSubscriber; +import com.google.genkit.ai.agent.internal.SnapshotSharding; +import com.google.genkit.core.GenkitException; +import com.google.genkit.core.JsonUtils; +import com.google.genkit.core.jsonpatch.JsonPatch; +import java.nio.charset.StandardCharsets; +import java.util.ArrayList; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.UUID; +import java.util.concurrent.Executors; +import java.util.concurrent.ScheduledExecutorService; +import java.util.concurrent.ScheduledFuture; +import java.util.concurrent.TimeUnit; +import java.util.function.Consumer; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; +import software.amazon.awssdk.services.dynamodb.DynamoDbClient; +import software.amazon.awssdk.services.dynamodb.model.AttributeDefinition; +import software.amazon.awssdk.services.dynamodb.model.AttributeValue; +import software.amazon.awssdk.services.dynamodb.model.BillingMode; +import software.amazon.awssdk.services.dynamodb.model.ConditionalCheckFailedException; +import software.amazon.awssdk.services.dynamodb.model.CreateTableRequest; +import software.amazon.awssdk.services.dynamodb.model.DescribeTableRequest; +import software.amazon.awssdk.services.dynamodb.model.GetItemRequest; +import software.amazon.awssdk.services.dynamodb.model.KeySchemaElement; +import software.amazon.awssdk.services.dynamodb.model.KeyType; +import software.amazon.awssdk.services.dynamodb.model.PutItemRequest; +import software.amazon.awssdk.services.dynamodb.model.ResourceNotFoundException; +import software.amazon.awssdk.services.dynamodb.model.ScalarAttributeType; +import software.amazon.awssdk.services.dynamodb.model.UpdateItemRequest; + +/** + * DynamoDB-backed implementation of {@link SessionStore} and {@link SnapshotSubscriber}. + * + *

Persists session snapshots with the same sharded checkpoint + diff + pointer layout as the + * Firestore backend (see {@code FirestoreSessionStore}), sharing the pure logic in {@link + * SnapshotSharding}. + * + *

Storage layout (single table)

+ * + *

All rows live in one table (default {@code genkit-sessions}) with a composite key: partition + * key {@code pk} = the per-tenant prefix (default {@code "global"}), sort key {@code sk}: + * + *

    + *
  • {@code SNAP#} — one metadata row per snapshot. {@code kind} is {@code + * "checkpoint"} or {@code "diff"}; carries {@code checkpointId}, {@code + * checkpointShardCount}, {@code segmentPath}, {@code statePatch} (RFC-6902 as a JSON string + * for diffs), {@code error} (JSON string), and a numeric {@code version} for optimistic + * concurrency. + *
  • {@code SHARD##} — a shard of the checkpoint state JSON. + *
  • {@code PTR#} — the current leaf pointer for a session. + *
+ * + *

Concurrency

+ * + *

DynamoDB has no interactive read-then-write transaction, so the observable {@code + * saveSnapshot} contract is preserved with idempotent unique-key writes plus a conditional, + * monotonic pointer advance: shards are written first, then the snapshot row, then the pointer + * flips — so a reader following the pointer always sees complete data. Updating an existing + * snapshot id uses a {@code version} conditional put and re-applies the (pure) mutator on conflict. + * + * @param the type of custom session state + */ +public final class DynamoDbSessionStore implements SessionStore, SnapshotSubscriber { + + private static final Logger logger = LoggerFactory.getLogger(DynamoDbSessionStore.class); + private static final ObjectMapper MAPPER = JsonUtils.getObjectMapper(); + + static final String KIND_CHECKPOINT = "checkpoint"; + static final String KIND_DIFF = "diff"; + + private static final int MAX_ATTEMPTS = 5; + + private final DynamoDbClient db; + private final DynamoDbSessionStoreOptions options; + private final ScheduledExecutorService scheduler; + + /** + * Creates a store with default options. + * + * @param db the DynamoDB client + */ + public DynamoDbSessionStore(DynamoDbClient db) { + this(db, DynamoDbSessionStoreOptions.defaults()); + } + + /** + * Creates a store. + * + * @param db the DynamoDB client + * @param options the store options + */ + public DynamoDbSessionStore(DynamoDbClient db, DynamoDbSessionStoreOptions options) { + if (db == null) { + throw new IllegalArgumentException("DynamoDbClient must be non-null"); + } + if (options == null) { + throw new IllegalArgumentException("options must be non-null"); + } + this.db = db; + this.options = options; + this.scheduler = + Executors.newSingleThreadScheduledExecutor( + r -> { + Thread t = new Thread(r, "genkit-dynamodb-session-store-poll"); + t.setDaemon(true); + return t; + }); + if (options.isCreateTableIfNotExists()) { + ensureTable(); + } + } + + // ────────────────────────────────────────────────────────────────────────── + // Key helpers + // ────────────────────────────────────────────────────────────────────────── + + private String prefix(SessionStoreOptions opts) { + String p = + options.getSnapshotPathPrefix().apply(opts != null ? opts : SessionStoreOptions.empty()); + return (p == null || p.isBlank()) ? "global" : p; + } + + private static String snapSk(String snapshotId) { + return "SNAP#" + snapshotId; + } + + private static String shardSk(String checkpointId, int index) { + return "SHARD#" + checkpointId + "#" + index; + } + + private static String ptrSk(String sessionId) { + return "PTR#" + sessionId; + } + + private static Map key(String prefix, String sk) { + Map k = new HashMap<>(); + k.put("pk", s(prefix)); + k.put("sk", s(sk)); + return k; + } + + private static AttributeValue s(String value) { + return AttributeValue.builder().s(value).build(); + } + + private static AttributeValue n(long value) { + return AttributeValue.builder().n(Long.toString(value)).build(); + } + + // ────────────────────────────────────────────────────────────────────────── + // Table management + // ────────────────────────────────────────────────────────────────────────── + + private void ensureTable() { + String table = options.getTableName(); + try { + db.describeTable(DescribeTableRequest.builder().tableName(table).build()); + } catch (ResourceNotFoundException e) { + db.createTable( + CreateTableRequest.builder() + .tableName(table) + .keySchema( + KeySchemaElement.builder().attributeName("pk").keyType(KeyType.HASH).build(), + KeySchemaElement.builder().attributeName("sk").keyType(KeyType.RANGE).build()) + .attributeDefinitions( + AttributeDefinition.builder() + .attributeName("pk") + .attributeType(ScalarAttributeType.S) + .build(), + AttributeDefinition.builder() + .attributeName("sk") + .attributeType(ScalarAttributeType.S) + .build()) + .billingMode(BillingMode.PAY_PER_REQUEST) + .build()); + db.waiter().waitUntilTableExists(DescribeTableRequest.builder().tableName(table).build()); + } + } + + // ────────────────────────────────────────────────────────────────────────── + // SnapshotWriter + // ────────────────────────────────────────────────────────────────────────── + + /** + * {@inheritDoc} + * + *

Implements the same identity/sessionId/status defaulting contract as the reference stores, + * then writes the snapshot (checkpoint shards + metadata row) and advances the session pointer. + */ + @Override + public String saveSnapshot( + String snapshotId, SnapshotMutator mutator, SessionStoreOptions storeOpts) { + String prefix = prefix(storeOpts); + String table = options.getTableName(); + + for (int attempt = 1; ; attempt++) { + // 1. Read existing snapshot + reconstruct state. + SessionSnapshot existing = null; + String existingSessionId = null; + Long existingVersion = null; + if (snapshotId != null) { + SnapshotSharding.validateId(snapshotId); + Map item = getItem(key(prefix, snapSk(snapshotId))); + if (item != null) { + existing = readSnapshot(prefix, item); + existingSessionId = existing.getSessionId(); + existingVersion = getLong(item, "version"); + } + } + + // 2. Apply mutator (pure — safe to re-run on conflict retry). + SessionSnapshot result = mutator.apply(existing); + if (result == null) { + return null; + } + + // 3. Identity / sessionId / status defaulting (mirror the reference stores). + String finalId; + if (snapshotId != null) { + finalId = snapshotId; + } else if (result.getSnapshotId() != null && !result.getSnapshotId().isBlank()) { + finalId = result.getSnapshotId(); + } else { + finalId = UUID.randomUUID().toString(); + } + SnapshotSharding.validateId(finalId); + result.setSnapshotId(finalId); + + if (existingSessionId != null) { + result.setSessionId(existingSessionId); + } + if (result.getSessionId() == null && result.getState() != null) { + result.setSessionId(result.getState().getSessionId()); + } + if (result.getSessionId() == null || result.getSessionId().isBlank()) { + throw GenkitException.builder() + .message("snapshot requires sessionId") + .errorCode("INVALID_ARGUMENT") + .build(); + } + if (result.getStatus() == null) { + result.setStatus(SnapshotStatus.COMPLETED); + } + + // 4. Resolve parent metadata to decide checkpoint vs diff. + String parentId = result.getParentId(); + JsonNode newState = stateToJson(result.getState()); + + ParentInfo parent = null; + if (parentId != null && !parentId.isBlank()) { + Map parentItem = getItem(key(prefix, snapSk(parentId))); + if (parentItem != null) { + parent = loadParentInfo(prefix, parentItem); + } + } + + boolean parentExists = parent != null; + int depthFromCheckpoint = 0; + JsonNode statePatch = null; + int diffSizeBytes = 0; + if (parent != null) { + depthFromCheckpoint = parent.segmentPath.size() + 1; + statePatch = JsonPatch.diff(parent.state, newState); + diffSizeBytes = jsonBytes(statePatch); + } + + boolean checkpoint = + SnapshotSharding.shouldCheckpoint( + parentExists, + depthFromCheckpoint, + options.getCheckpointInterval(), + diffSizeBytes, + options.getShardSize()); + + // 5. Build the snapshot metadata row (+ shard rows for a checkpoint). + Map row = baseRow(prefix, result); + String checkpointId; + int checkpointShardCount; + List segmentPath; + + if (checkpoint) { + checkpointId = finalId; + String stateJson = writeJson(newState); + List shards = SnapshotSharding.shardString(stateJson, options.getShardSize()); + checkpointShardCount = shards.size(); + segmentPath = new ArrayList<>(); + for (int i = 0; i < shards.size(); i++) { + Map shardRow = new HashMap<>(); + shardRow.put("pk", s(prefix)); + shardRow.put("sk", s(shardSk(checkpointId, i))); + shardRow.put("checkpointId", s(checkpointId)); + shardRow.put("index", n(i)); + shardRow.put("data", s(shards.get(i))); + db.putItem(PutItemRequest.builder().tableName(table).item(shardRow).build()); + } + row.put("kind", s(KIND_CHECKPOINT)); + } else { + ParentInfo p = parent; + if (p == null) { + throw new GenkitException("internal: diff path without parent"); + } + checkpointId = p.checkpointId; + checkpointShardCount = p.checkpointShardCount; + segmentPath = new ArrayList<>(p.segmentPath); + segmentPath.add(finalId); + row.put("kind", s(KIND_DIFF)); + row.put("statePatch", s(writeJson(statePatch))); + } + + row.put("checkpointId", s(checkpointId)); + row.put("checkpointShardCount", n(checkpointShardCount)); + row.put("segmentPath", stringList(segmentPath)); + long newVersion = (existingVersion == null ? 0 : existingVersion) + 1; + row.put("version", n(newVersion)); + + // 6. Write the snapshot row with optimistic concurrency. + PutItemRequest.Builder put = PutItemRequest.builder().tableName(table).item(row); + if (existingVersion == null) { + put.conditionExpression("attribute_not_exists(sk)"); + } else { + put.conditionExpression("version = :ev") + .expressionAttributeValues(Map.of(":ev", n(existingVersion))); + } + try { + db.putItem(put.build()); + } catch (ConditionalCheckFailedException e) { + if (attempt < MAX_ATTEMPTS) { + continue; // concurrent writer won the race; re-read and retry the pure mutator. + } + throw new GenkitException("Failed to save snapshot after " + MAX_ATTEMPTS + " attempts", e); + } + + // 7. Advance the session pointer (never backward). A lost race here is benign. + advancePointer( + prefix, + result.getSessionId(), + finalId, + result.getCreatedAt(), + result.getUpdatedAt(), + checkpointId, + checkpointShardCount, + segmentPath); + + return finalId; + } + } + + /** Advances the session pointer to the new leaf unless the stored pointer is already newer. */ + private void advancePointer( + String prefix, + String sessionId, + String snapshotId, + String createdAt, + String updatedAt, + String checkpointId, + int checkpointShardCount, + List segmentPath) { + String table = options.getTableName(); + + Map values = new HashMap<>(); + values.put(":sid", s(snapshotId)); + values.put(":cp", s(checkpointId)); + values.put(":sc", n(checkpointShardCount)); + values.put(":sp", stringList(segmentPath)); + + StringBuilder setExpr = + new StringBuilder( + "SET currentSnapshotId = :sid, checkpointId = :cp, checkpointShardCount = :sc," + + " segmentPath = :sp"); + if (createdAt != null) { + setExpr.append(", currentCreatedAt = :cc"); + values.put(":cc", s(createdAt)); + } + if (updatedAt != null) { + setExpr.append(", updatedAt = :ua"); + values.put(":ua", s(updatedAt)); + } + + UpdateItemRequest.Builder update = + UpdateItemRequest.builder() + .tableName(table) + .key(key(prefix, ptrSk(sessionId))) + .updateExpression(setExpr.toString()) + .expressionAttributeValues(values); + + // Only guard against moving backward when we have a comparable createdAt. When createdAt is + // null we always advance (matching the reference stores, which treat null createdAt as newer). + if (createdAt != null) { + update.conditionExpression( + "attribute_not_exists(currentSnapshotId) OR attribute_not_exists(currentCreatedAt) OR" + + " currentCreatedAt <= :cc OR currentSnapshotId = :sid"); + } + + try { + db.updateItem(update.build()); + } catch (ConditionalCheckFailedException e) { + // Stored pointer is already newer; leave it in place. + logger.debug("Pointer for session {} not advanced (stored leaf is newer)", sessionId); + } + } + + // ────────────────────────────────────────────────────────────────────────── + // SnapshotReader + // ────────────────────────────────────────────────────────────────────────── + + /** + * {@inheritDoc} + * + *

By {@code snapshotId}: loads the snapshot row and reconstructs its state from the checkpoint + * shards + ordered {@code segmentPath} diffs. By {@code sessionId}: reads the pointer, then loads + * and reconstructs the pointed snapshot. + */ + @Override + public SessionSnapshot getSnapshot(GetSnapshotOptions opts) { + if (opts == null) { + return null; + } + SessionStoreOptions storeOpts = SessionStoreOptions.empty(); + String prefix = prefix(storeOpts); + + if (opts.getSnapshotId() != null) { + SnapshotSharding.validateId(opts.getSnapshotId()); + Map item = getItem(key(prefix, snapSk(opts.getSnapshotId()))); + return item == null ? null : readSnapshot(prefix, item); + } + if (opts.getSessionId() != null) { + SnapshotSharding.validateId(opts.getSessionId()); + Map pointer = getItem(key(prefix, ptrSk(opts.getSessionId()))); + if (pointer == null) { + return null; + } + String leafId = getString(pointer, "currentSnapshotId"); + if (leafId == null) { + return null; + } + Map item = getItem(key(prefix, snapSk(leafId))); + return item == null ? null : readSnapshot(prefix, item); + } + return null; + } + + // ────────────────────────────────────────────────────────────────────────── + // SnapshotSubscriber (polling) + // ────────────────────────────────────────────────────────────────────────── + + /** + * {@inheritDoc} + * + *

DynamoDB has no built-in per-item change notification usable here, so the subscription polls + * {@link #getSnapshot} on a shared daemon scheduler and fires the callback whenever the + * serialized snapshot content changes. The callback also fires immediately if the snapshot + * already exists. + */ + @Override + public AutoCloseable onSnapshotStateChange( + String snapshotId, Consumer> cb, SessionStoreOptions storeOpts) { + SnapshotSharding.validateId(snapshotId); + GetSnapshotOptions get = GetSnapshotOptions.builder().snapshotId(snapshotId).build(); + + final String[] lastContent = {null}; + SessionSnapshot initial = getSnapshot(get); + if (initial != null) { + lastContent[0] = serializeQuietly(initial); + cb.accept(initial); + } + + ScheduledFuture future = + scheduler.scheduleAtFixedRate( + () -> { + try { + SessionSnapshot snap = getSnapshot(get); + if (snap == null) { + return; + } + String content = serializeQuietly(snap); + if (!content.equals(lastContent[0])) { + lastContent[0] = content; + cb.accept(snap); + } + } catch (Exception e) { + // Swallow poll errors — don't kill the scheduler thread. + } + }, + options.getPollIntervalMs(), + options.getPollIntervalMs(), + TimeUnit.MILLISECONDS); + + return () -> future.cancel(false); + } + + // ────────────────────────────────────────────────────────────────────────── + // Row (de)serialization + // ────────────────────────────────────────────────────────────────────────── + + /** Builds the non-state base row for a snapshot (pk/sk + metadata). */ + private Map baseRow(String prefix, SessionSnapshot snap) { + Map data = new HashMap<>(); + data.put("pk", s(prefix)); + data.put("sk", s(snapSk(snap.getSnapshotId()))); + data.put("snapshotId", s(snap.getSnapshotId())); + data.put("sessionId", s(snap.getSessionId())); + if (snap.getParentId() != null) { + data.put("parentId", s(snap.getParentId())); + } + if (snap.getCreatedAt() != null) { + data.put("createdAt", s(snap.getCreatedAt())); + } + if (snap.getUpdatedAt() != null) { + data.put("updatedAt", s(snap.getUpdatedAt())); + } + if (snap.getHeartbeatAt() != null) { + data.put("heartbeatAt", s(snap.getHeartbeatAt())); + } + if (snap.getStatus() != null) { + data.put("status", s(snap.getStatus().getValue())); + } + if (snap.getFinishReason() != null) { + data.put("finishReason", s(snap.getFinishReason().getValue())); + } + if (snap.getError() != null) { + data.put("error", s(writeJson(snap.getError()))); + } + return data; + } + + /** Holds the reconstructed parent state and its checkpoint lineage. */ + private static final class ParentInfo { + JsonNode state; + String checkpointId; + int checkpointShardCount; + List segmentPath; + } + + /** Loads a parent snapshot's checkpoint lineage and reconstructs its state. */ + private ParentInfo loadParentInfo(String prefix, Map item) { + ParentInfo info = new ParentInfo(); + info.checkpointId = getString(item, "checkpointId"); + Long shardCount = getLong(item, "checkpointShardCount"); + info.checkpointShardCount = shardCount != null ? shardCount.intValue() : 0; + info.segmentPath = getStringList(item, "segmentPath"); + info.state = + reconstructFullState( + prefix, info.checkpointId, info.checkpointShardCount, info.segmentPath); + return info; + } + + /** Reads and fully reconstructs a snapshot from its metadata row. */ + private SessionSnapshot readSnapshot(String prefix, Map item) { + String checkpointId = getString(item, "checkpointId"); + Long shardCount = getLong(item, "checkpointShardCount"); + int checkpointShardCount = shardCount != null ? shardCount.intValue() : 0; + List segmentPath = getStringList(item, "segmentPath"); + JsonNode state = reconstructFullState(prefix, checkpointId, checkpointShardCount, segmentPath); + return rowToSnapshot(item, state); + } + + /** + * Reconstructs full state: loads the checkpoint shards (concatenate, parse) then applies the + * {@code segmentPath} diffs in order. + */ + private JsonNode reconstructFullState( + String prefix, String checkpointId, int checkpointShardCount, List segmentPath) { + if (checkpointId == null) { + return NullNode.getInstance(); + } + List shardContents = new ArrayList<>(); + for (int i = 0; i < checkpointShardCount; i++) { + Map shard = getItem(key(prefix, shardSk(checkpointId, i))); + shardContents.add(shard != null ? getString(shard, "data") : ""); + } + String checkpointJson = SnapshotSharding.reassembleShards(shardContents); + + List diffs = new ArrayList<>(); + for (String diffId : segmentPath) { + Map diffItem = getItem(key(prefix, snapSk(diffId))); + if (diffItem != null) { + String patch = getString(diffItem, "statePatch"); + if (patch != null) { + diffs.add(patch); + } + } + } + try { + return SnapshotSharding.reconstructState(checkpointJson, diffs); + } catch (Exception e) { + throw new GenkitException("Failed to reconstruct session state: " + e.getMessage(), e); + } + } + + /** Builds a {@link SessionSnapshot} from a metadata row and reconstructed state. */ + @SuppressWarnings("unchecked") + private SessionSnapshot rowToSnapshot(Map item, JsonNode state) { + SessionSnapshot.Builder builder = SessionSnapshot.builder(); + builder.snapshotId(getString(item, "snapshotId")); + builder.sessionId(getString(item, "sessionId")); + builder.parentId(getString(item, "parentId")); + builder.createdAt(getString(item, "createdAt")); + builder.updatedAt(getString(item, "updatedAt")); + builder.heartbeatAt(getString(item, "heartbeatAt")); + String status = getString(item, "status"); + if (status != null) { + builder.status(SnapshotStatus.fromValueOrCompleted(status)); + } + String finishReason = getString(item, "finishReason"); + if (finishReason != null) { + builder.finishReason(AgentFinishReason.fromValue(finishReason)); + } + String errorJson = getString(item, "error"); + if (errorJson != null) { + try { + builder.error(MAPPER.readValue(errorJson, RuntimeError.class)); + } catch (Exception e) { + throw new GenkitException("Failed to parse snapshot error: " + e.getMessage(), e); + } + } + if (state != null && !state.isNull()) { + try { + builder.state((SessionState) MAPPER.treeToValue(state, SessionState.class)); + } catch (Exception e) { + throw new GenkitException("Failed to parse session state: " + e.getMessage(), e); + } + } + return builder.build(); + } + + /** Serializes session state to a JSON node (null state → JSON null). */ + private JsonNode stateToJson(SessionState state) { + if (state == null) { + return NullNode.getInstance(); + } + return MAPPER.valueToTree(state); + } + + // ────────────────────────────────────────────────────────────────────────── + // Low-level DynamoDB + attribute helpers + // ────────────────────────────────────────────────────────────────────────── + + /** Consistent-read GetItem; returns {@code null} when the item does not exist. */ + private Map getItem(Map key) { + Map item = + db.getItem( + GetItemRequest.builder() + .tableName(options.getTableName()) + .key(key) + .consistentRead(true) + .build()) + .item(); + return (item == null || item.isEmpty()) ? null : item; + } + + private static AttributeValue stringList(List values) { + List list = new ArrayList<>(); + for (String v : values) { + list.add(s(v)); + } + return AttributeValue.builder().l(list).build(); + } + + private static String getString(Map item, String name) { + AttributeValue av = item.get(name); + return (av == null || av.s() == null) ? null : av.s(); + } + + private static Long getLong(Map item, String name) { + AttributeValue av = item.get(name); + if (av == null || av.n() == null) { + return null; + } + return Long.parseLong(av.n()); + } + + private static List getStringList(Map item, String name) { + AttributeValue av = item.get(name); + List out = new ArrayList<>(); + if (av != null && av.hasL()) { + for (AttributeValue e : av.l()) { + if (e.s() != null) { + out.add(e.s()); + } + } + } + return out; + } + + private static String writeJson(Object value) { + try { + return MAPPER.writeValueAsString(value); + } catch (Exception e) { + throw new GenkitException("Failed to serialize value: " + e.getMessage(), e); + } + } + + private static int jsonBytes(Object value) { + return writeJson(value).getBytes(StandardCharsets.UTF_8).length; + } + + private static String serializeQuietly(SessionSnapshot snap) { + try { + return MAPPER.writeValueAsString(MAPPER.valueToTree(snap)); + } catch (Exception e) { + return ""; + } + } +} diff --git a/plugins/aws-bedrock/src/main/java/com/google/genkit/plugins/awsbedrock/session/DynamoDbSessionStoreOptions.java b/plugins/aws-bedrock/src/main/java/com/google/genkit/plugins/awsbedrock/session/DynamoDbSessionStoreOptions.java new file mode 100644 index 000000000..3db0e9c33 --- /dev/null +++ b/plugins/aws-bedrock/src/main/java/com/google/genkit/plugins/awsbedrock/session/DynamoDbSessionStoreOptions.java @@ -0,0 +1,243 @@ +/* + * 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. + * + * SPDX-License-Identifier: Apache-2.0 + */ + +package com.google.genkit.plugins.awsbedrock.session; + +import com.google.genkit.ai.agent.SessionStoreOptions; +import java.util.function.Function; + +/** + * Configuration for {@link DynamoDbSessionStore}. + * + *

The store persists all rows in a single DynamoDB table (default {@value #DEFAULT_TABLE}) keyed + * by a partition key {@code pk} (the per-tenant prefix) and a sort key {@code sk} that + * discriminates snapshot, shard, and pointer rows. All rows for a tenant share the same {@code pk}, + * derived from {@link #getSnapshotPathPrefix()} (default {@code "global"}). + * + *

The default {@link #getShardSize()} is {@value #DEFAULT_SHARD_SIZE} bytes, kept safely under + * the DynamoDB 400 KB item-size limit; because the store forces a checkpoint whenever a diff + * would exceed the shard size, diff rows stay bounded too. + */ +public final class DynamoDbSessionStoreOptions { + + /** Default table name. */ + public static final String DEFAULT_TABLE = "genkit-sessions"; + + /** Default number of turns between full checkpoints. */ + public static final int DEFAULT_CHECKPOINT_INTERVAL = 25; + + /** Default shard size in bytes for checkpoint state (350 KiB, under the 400 KB item cap). */ + public static final int DEFAULT_SHARD_SIZE = 350 * 1024; + + /** Default subscription poll interval in milliseconds. */ + public static final long DEFAULT_POLL_INTERVAL_MS = 2000L; + + private final String tableName; + private final int checkpointInterval; + private final int shardSize; + private final Function snapshotPathPrefix; + private final boolean createTableIfNotExists; + private final long pollIntervalMs; + + private DynamoDbSessionStoreOptions(Builder builder) { + this.tableName = builder.tableName; + this.checkpointInterval = builder.checkpointInterval; + this.shardSize = builder.shardSize; + this.snapshotPathPrefix = builder.snapshotPathPrefix; + this.createTableIfNotExists = builder.createTableIfNotExists; + this.pollIntervalMs = builder.pollIntervalMs; + } + + /** + * Returns the DynamoDB table name (default {@value #DEFAULT_TABLE}). + * + * @return the table name + */ + public String getTableName() { + return tableName; + } + + /** + * Returns the number of turns between full checkpoints (default {@value + * #DEFAULT_CHECKPOINT_INTERVAL}). + * + * @return the checkpoint interval + */ + public int getCheckpointInterval() { + return checkpointInterval; + } + + /** + * Returns the shard size in bytes for checkpoint state (default {@value #DEFAULT_SHARD_SIZE}). + * + * @return the shard size in bytes + */ + public int getShardSize() { + return shardSize; + } + + /** + * Returns the function that derives the per-tenant partition-key prefix from the per-request + * store options (default {@code o -> "global"}). + * + * @return the prefix function + */ + public Function getSnapshotPathPrefix() { + return snapshotPathPrefix; + } + + /** + * Returns whether the store should create the table on first use if it does not exist (default + * {@code false}). + * + * @return {@code true} if the table should be created when missing + */ + public boolean isCreateTableIfNotExists() { + return createTableIfNotExists; + } + + /** + * Returns the subscription poll interval in milliseconds (default {@value + * #DEFAULT_POLL_INTERVAL_MS}). + * + * @return the poll interval in milliseconds + */ + public long getPollIntervalMs() { + return pollIntervalMs; + } + + /** + * Returns default options. + * + * @return a {@code DynamoDbSessionStoreOptions} with all defaults + */ + public static DynamoDbSessionStoreOptions defaults() { + return builder().build(); + } + + /** + * Creates a new builder. + * + * @return a new builder + */ + public static Builder builder() { + return new Builder(); + } + + /** Builder for {@link DynamoDbSessionStoreOptions}. */ + public static final class Builder { + private String tableName = DEFAULT_TABLE; + private int checkpointInterval = DEFAULT_CHECKPOINT_INTERVAL; + private int shardSize = DEFAULT_SHARD_SIZE; + private Function snapshotPathPrefix = o -> "global"; + private boolean createTableIfNotExists = false; + private long pollIntervalMs = DEFAULT_POLL_INTERVAL_MS; + + private Builder() {} + + /** + * Sets the DynamoDB table name. + * + * @param tableName the table name + * @return this builder + */ + public Builder tableName(String tableName) { + this.tableName = tableName; + return this; + } + + /** + * Sets the number of turns between full checkpoints. + * + * @param checkpointInterval the checkpoint interval (must be {@code >= 1}) + * @return this builder + */ + public Builder checkpointInterval(int checkpointInterval) { + this.checkpointInterval = checkpointInterval; + return this; + } + + /** + * Sets the shard size in bytes for checkpoint state (must stay under the 400 KB item cap). + * + * @param shardSize the shard size in bytes (must be {@code >= 1}) + * @return this builder + */ + public Builder shardSize(int shardSize) { + this.shardSize = shardSize; + return this; + } + + /** + * Sets the function that derives the per-tenant partition-key prefix. + * + * @param snapshotPathPrefix the prefix function + * @return this builder + */ + public Builder snapshotPathPrefix(Function snapshotPathPrefix) { + this.snapshotPathPrefix = snapshotPathPrefix; + return this; + } + + /** + * Sets whether to create the table on first use if it does not exist. + * + * @param createTableIfNotExists whether to create the table when missing + * @return this builder + */ + public Builder createTableIfNotExists(boolean createTableIfNotExists) { + this.createTableIfNotExists = createTableIfNotExists; + return this; + } + + /** + * Sets the subscription poll interval in milliseconds. + * + * @param pollIntervalMs the poll interval (must be {@code >= 1}) + * @return this builder + */ + public Builder pollIntervalMs(long pollIntervalMs) { + this.pollIntervalMs = pollIntervalMs; + return this; + } + + /** + * Builds a new {@code DynamoDbSessionStoreOptions}. + * + * @return a new options instance + */ + public DynamoDbSessionStoreOptions build() { + if (tableName == null || tableName.isBlank()) { + throw new IllegalArgumentException("tableName must be non-empty"); + } + if (checkpointInterval < 1) { + throw new IllegalArgumentException("checkpointInterval must be >= 1"); + } + if (shardSize < 1) { + throw new IllegalArgumentException("shardSize must be >= 1"); + } + if (snapshotPathPrefix == null) { + throw new IllegalArgumentException("snapshotPathPrefix must be non-null"); + } + if (pollIntervalMs < 1) { + throw new IllegalArgumentException("pollIntervalMs must be >= 1"); + } + return new DynamoDbSessionStoreOptions(this); + } + } +} diff --git a/plugins/aws-bedrock/src/main/java/com/google/genkit/plugins/awsbedrock/session/package-info.java b/plugins/aws-bedrock/src/main/java/com/google/genkit/plugins/awsbedrock/session/package-info.java new file mode 100644 index 000000000..92cbfe293 --- /dev/null +++ b/plugins/aws-bedrock/src/main/java/com/google/genkit/plugins/awsbedrock/session/package-info.java @@ -0,0 +1,27 @@ +/* + * 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. + * + * SPDX-License-Identifier: Apache-2.0 + */ + +/** + * DynamoDB-backed agent session persistence. + * + *

{@link com.google.genkit.plugins.awsbedrock.session.DynamoDbSessionStore} implements the + * Genkit {@code SessionStore} contract using the sharded checkpoint + RFC-6902 diff + pointer + * layout shared with the Firestore and Cosmos DB backends. Construct it directly and pass it to an + * agent via {@code AgentConfig.store(...)}. + */ +package com.google.genkit.plugins.awsbedrock.session; diff --git a/plugins/aws-bedrock/src/test/java/com/google/genkit/plugins/awsbedrock/session/DynamoDbSessionStoreOptionsTest.java b/plugins/aws-bedrock/src/test/java/com/google/genkit/plugins/awsbedrock/session/DynamoDbSessionStoreOptionsTest.java new file mode 100644 index 000000000..5e06fcc46 --- /dev/null +++ b/plugins/aws-bedrock/src/test/java/com/google/genkit/plugins/awsbedrock/session/DynamoDbSessionStoreOptionsTest.java @@ -0,0 +1,76 @@ +/* + * 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. + * + * SPDX-License-Identifier: Apache-2.0 + */ + +package com.google.genkit.plugins.awsbedrock.session; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertThrows; + +import com.google.genkit.ai.agent.SessionStoreOptions; +import org.junit.jupiter.api.Test; + +/** Unit tests for {@link DynamoDbSessionStoreOptions}. */ +class DynamoDbSessionStoreOptionsTest { + + @Test + void defaultsAreSane() { + DynamoDbSessionStoreOptions o = DynamoDbSessionStoreOptions.defaults(); + assertEquals("genkit-sessions", o.getTableName()); + assertEquals(25, o.getCheckpointInterval()); + assertEquals(350 * 1024, o.getShardSize()); + assertEquals("global", o.getSnapshotPathPrefix().apply(SessionStoreOptions.empty())); + assertFalse(o.isCreateTableIfNotExists()); + assertEquals(2000L, o.getPollIntervalMs()); + } + + @Test + void customBuilder() { + DynamoDbSessionStoreOptions o = + DynamoDbSessionStoreOptions.builder() + .tableName("my-sessions") + .checkpointInterval(10) + .shardSize(1024) + .snapshotPathPrefix(so -> "tenant-1") + .createTableIfNotExists(true) + .pollIntervalMs(500) + .build(); + assertEquals("my-sessions", o.getTableName()); + assertEquals(10, o.getCheckpointInterval()); + assertEquals(1024, o.getShardSize()); + assertEquals("tenant-1", o.getSnapshotPathPrefix().apply(SessionStoreOptions.empty())); + assertEquals(true, o.isCreateTableIfNotExists()); + assertEquals(500L, o.getPollIntervalMs()); + } + + @Test + void builderValidates() { + assertThrows( + IllegalArgumentException.class, + () -> DynamoDbSessionStoreOptions.builder().tableName("").build()); + assertThrows( + IllegalArgumentException.class, + () -> DynamoDbSessionStoreOptions.builder().shardSize(0).build()); + assertThrows( + IllegalArgumentException.class, + () -> DynamoDbSessionStoreOptions.builder().checkpointInterval(0).build()); + assertThrows( + IllegalArgumentException.class, + () -> DynamoDbSessionStoreOptions.builder().pollIntervalMs(0).build()); + } +} diff --git a/plugins/aws-bedrock/src/test/java/com/google/genkit/plugins/awsbedrock/session/DynamoDbSessionStoreTest.java b/plugins/aws-bedrock/src/test/java/com/google/genkit/plugins/awsbedrock/session/DynamoDbSessionStoreTest.java new file mode 100644 index 000000000..539d44def --- /dev/null +++ b/plugins/aws-bedrock/src/test/java/com/google/genkit/plugins/awsbedrock/session/DynamoDbSessionStoreTest.java @@ -0,0 +1,193 @@ +/* + * 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. + * + * SPDX-License-Identifier: Apache-2.0 + */ + +package com.google.genkit.plugins.awsbedrock.session; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertNull; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assumptions.assumeTrue; + +import com.google.genkit.ai.Message; +import com.google.genkit.ai.agent.GetSnapshotOptions; +import com.google.genkit.ai.agent.SessionSnapshot; +import com.google.genkit.ai.agent.SessionState; +import com.google.genkit.ai.agent.SessionStoreOptions; +import com.google.genkit.ai.agent.SnapshotStatus; +import com.google.genkit.core.GenkitException; +import java.net.URI; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.UUID; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import software.amazon.awssdk.auth.credentials.AwsBasicCredentials; +import software.amazon.awssdk.auth.credentials.StaticCredentialsProvider; +import software.amazon.awssdk.regions.Region; +import software.amazon.awssdk.services.dynamodb.DynamoDbClient; + +/** + * Tests for {@link DynamoDbSessionStore}. + * + *

Integration tests are gated on the {@code DYNAMODB_LOCAL_ENDPOINT} environment variable (e.g. + * {@code http://localhost:8000} for a DynamoDB Local container). When it is unset the tests are + * skipped via {@link org.junit.jupiter.api.Assumptions}. + */ +class DynamoDbSessionStoreTest { + + private static final String ENDPOINT = System.getenv("DYNAMODB_LOCAL_ENDPOINT"); + + private DynamoDbClient client; + private DynamoDbSessionStore> store; + + @BeforeEach + void setUp() { + if (ENDPOINT == null || ENDPOINT.isEmpty()) { + return; // integration tests skip via assumeTrue + } + client = + DynamoDbClient.builder() + .endpointOverride(URI.create(ENDPOINT)) + .region(Region.US_EAST_1) + .credentialsProvider( + StaticCredentialsProvider.create(AwsBasicCredentials.create("local", "local"))) + .build(); + String table = "genkit-sessions-test-" + UUID.randomUUID().toString().substring(0, 8); + store = + new DynamoDbSessionStore<>( + client, + DynamoDbSessionStoreOptions.builder() + .tableName(table) + .checkpointInterval(3) + .createTableIfNotExists(true) + .build()); + } + + @AfterEach + void tearDown() { + if (client != null) { + client.close(); + } + } + + private static SessionSnapshot> snapshotWithState( + String sessionId, String parentId, Map custom) { + SessionState> state = + SessionState.>builder() + .sessionId(sessionId) + .messages(List.of(Message.user("hello"))) + .custom(custom) + .build(); + return SessionSnapshot.>builder() + .sessionId(sessionId) + .parentId(parentId) + .status(SnapshotStatus.COMPLETED) + .state(state) + .build(); + } + + @Test + void saveThenGetBySnapshotIdRoundTrips() { + assumeTrue(ENDPOINT != null && !ENDPOINT.isEmpty()); + String sessionId = "s-" + UUID.randomUUID(); + Map custom = new HashMap<>(); + custom.put("count", 1); + + String id = + store.saveSnapshot( + null, e -> snapshotWithState(sessionId, null, custom), SessionStoreOptions.empty()); + assertNotNull(id); + + SessionSnapshot> got = + store.getSnapshot(GetSnapshotOptions.builder().snapshotId(id).build()); + assertNotNull(got); + assertEquals(sessionId, got.getSessionId()); + assertEquals(SnapshotStatus.COMPLETED, got.getStatus()); + assertEquals(1, got.getState().getMessages().size()); + assertEquals(1, ((Number) got.getState().getCustom().get("count")).intValue()); + } + + @Test + void getBySessionIdReturnsLeaf() { + assumeTrue(ENDPOINT != null && !ENDPOINT.isEmpty()); + String sessionId = "s-" + UUID.randomUUID(); + + Map c1 = new HashMap<>(); + c1.put("count", 1); + String id1 = + store.saveSnapshot( + null, e -> snapshotWithState(sessionId, null, c1), SessionStoreOptions.empty()); + Map c2 = new HashMap<>(); + c2.put("count", 2); + String id2 = + store.saveSnapshot( + null, e -> snapshotWithState(sessionId, id1, c2), SessionStoreOptions.empty()); + + SessionSnapshot> latest = + store.getSnapshot(GetSnapshotOptions.builder().sessionId(sessionId).build()); + assertNotNull(latest); + assertEquals(id2, latest.getSnapshotId()); + assertEquals(2, ((Number) latest.getState().getCustom().get("count")).intValue()); + } + + @Test + void diffThenCheckpointReconstructs() { + assumeTrue(ENDPOINT != null && !ENDPOINT.isEmpty()); + // checkpointInterval is 3; save 5 turns across checkpoint boundaries and confirm the leaf + // reconstructs correctly (checkpoint shards + segment-path diffs). + String sessionId = "s-" + UUID.randomUUID(); + String parent = null; + String lastId = null; + for (int i = 1; i <= 5; i++) { + Map c = new HashMap<>(); + c.put("count", i); + final String p = parent; + lastId = + store.saveSnapshot( + null, e -> snapshotWithState(sessionId, p, c), SessionStoreOptions.empty()); + parent = lastId; + } + SessionSnapshot> got = + store.getSnapshot(GetSnapshotOptions.builder().snapshotId(lastId).build()); + assertNotNull(got); + assertEquals(5, ((Number) got.getState().getCustom().get("count")).intValue()); + } + + @Test + void rejectsEmptySessionId() { + assumeTrue(ENDPOINT != null && !ENDPOINT.isEmpty()); + GenkitException ex = + assertThrows( + GenkitException.class, + () -> + store.saveSnapshot( + null, + e -> snapshotWithState("", null, new HashMap<>()), + SessionStoreOptions.empty())); + assertEquals("INVALID_ARGUMENT", ex.getErrorCode()); + } + + @Test + void mutatorNullIsNoOp() { + assumeTrue(ENDPOINT != null && !ENDPOINT.isEmpty()); + assertNull(store.saveSnapshot(null, e -> null, SessionStoreOptions.empty())); + } +} diff --git a/plugins/azure-foundry/pom.xml b/plugins/azure-foundry/pom.xml index bbc1df5da..d8fb93a59 100644 --- a/plugins/azure-foundry/pom.xml +++ b/plugins/azure-foundry/pom.xml @@ -37,6 +37,7 @@ false 1.18.4 + 4.71.0 @@ -45,6 +46,16 @@ genkit ${project.version} + + com.google.genkit + genkit-core + ${project.version} + + + com.google.genkit + genkit-ai + ${project.version} + com.google.genkit @@ -59,6 +70,19 @@ ${azure.identity.version} + + + com.azure + azure-cosmos + ${azure.cosmos.version} + + + + + com.fasterxml.jackson.core + jackson-databind + + org.slf4j diff --git a/plugins/azure-foundry/src/main/java/com/google/genkit/plugins/azurefoundry/AzureFoundryPlugin.java b/plugins/azure-foundry/src/main/java/com/google/genkit/plugins/azurefoundry/AzureFoundryPlugin.java index 200a80ea0..5e3827d5d 100644 --- a/plugins/azure-foundry/src/main/java/com/google/genkit/plugins/azurefoundry/AzureFoundryPlugin.java +++ b/plugins/azure-foundry/src/main/java/com/google/genkit/plugins/azurefoundry/AzureFoundryPlugin.java @@ -93,6 +93,14 @@ public AzureFoundryPlugin(AzureFoundryPluginOptions options) { CompatOAIPluginOptions.Builder compatBuilder = CompatOAIPluginOptions.builder().baseUrl(buildBaseUrl(options)); + // Pass api-version as a real query parameter so compat-oai appends it AFTER the + // /chat/completions path. Baking it into the base URL string (as buildBaseUrl used to) corrupts + // the request URL for non-Azure-OpenAI hosts, e.g. AI Foundry v1 endpoints on + // services.ai.azure.com. + if (options.getApiVersion() != null) { + compatBuilder.queryParams(java.util.Map.of("api-version", options.getApiVersion())); + } + // Handle authentication if (options.getApiKey() != null) { compatBuilder.apiKey(options.getApiKey()); @@ -154,8 +162,8 @@ private String buildBaseUrl(AzureFoundryPluginOptions options) { options.getEndpoint().contains("openai.azure.com") || options.getEndpoint().contains("cognitiveservices.azure.com"); - if (!options.getEndpoint().contains("/inference") - && !options.getEndpoint().contains("/openai")) { + String path = endpointPath(options.getEndpoint()); + if (!path.contains("inference") && !path.contains("openai")) { if (isAzureOpenAI) { // Azure OpenAI Service uses /openai/deployments/{deployment}/ path // The deployment name will be added by the model, so we just set the base @@ -166,14 +174,27 @@ private String buildBaseUrl(AzureFoundryPluginOptions options) { } } - // Add API version as query parameter - if (options.getApiVersion() != null) { - url.append("?api-version=").append(options.getApiVersion()); - } - + // NOTE: api-version is NOT appended here — it is added as a query parameter by the caller so it + // lands after the /chat/completions path segment (see the constructor and + // buildAzureOpenAIOptions). return url.toString(); } + /** + * Returns the path component of an endpoint URL (empty string when there is none or it can't be + * parsed). Used to detect whether the endpoint already carries the {@code /openai} or {@code + * /inference} path, without being fooled by hostnames that contain those words (e.g. a resource + * named {@code openai-foo}). + */ + private static String endpointPath(String endpoint) { + try { + String p = java.net.URI.create(endpoint).getPath(); + return p == null ? "" : p; + } catch (RuntimeException e) { + return ""; + } + } + /** * Builds CompatOAI options for Azure OpenAI deployments. Azure OpenAI requires the deployment * name in the URL path. @@ -184,8 +205,11 @@ private CompatOAIPluginOptions buildAzureOpenAIOptions(String deploymentName) { url.append("/"); } - // Azure OpenAI path: /openai/deployments/{deployment-id} - if (!options.getEndpoint().contains("/openai")) { + // Azure OpenAI path: /openai/deployments/{deployment-id}. Inspect the URL *path* (not the whole + // endpoint) so a resource whose host contains "openai" (e.g. + // https://openai-foo.openai.azure.com) isn't mistaken for an endpoint that already includes the + // /openai path. + if (!endpointPath(options.getEndpoint()).contains("openai")) { url.append("openai/deployments/").append(deploymentName); } diff --git a/plugins/azure-foundry/src/main/java/com/google/genkit/plugins/azurefoundry/session/CosmosSessionStore.java b/plugins/azure-foundry/src/main/java/com/google/genkit/plugins/azurefoundry/session/CosmosSessionStore.java new file mode 100644 index 000000000..5ddcdeb7c --- /dev/null +++ b/plugins/azure-foundry/src/main/java/com/google/genkit/plugins/azurefoundry/session/CosmosSessionStore.java @@ -0,0 +1,713 @@ +/* + * 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. + * + * SPDX-License-Identifier: Apache-2.0 + */ + +package com.google.genkit.plugins.azurefoundry.session; + +import com.azure.cosmos.CosmosClient; +import com.azure.cosmos.CosmosContainer; +import com.azure.cosmos.CosmosDatabase; +import com.azure.cosmos.CosmosException; +import com.azure.cosmos.models.CosmosItemRequestOptions; +import com.azure.cosmos.models.CosmosItemResponse; +import com.azure.cosmos.models.PartitionKey; +import com.fasterxml.jackson.databind.JsonNode; +import com.fasterxml.jackson.databind.ObjectMapper; +import com.fasterxml.jackson.databind.node.ArrayNode; +import com.fasterxml.jackson.databind.node.NullNode; +import com.fasterxml.jackson.databind.node.ObjectNode; +import com.google.genkit.ai.agent.AgentFinishReason; +import com.google.genkit.ai.agent.GetSnapshotOptions; +import com.google.genkit.ai.agent.RuntimeError; +import com.google.genkit.ai.agent.SessionSnapshot; +import com.google.genkit.ai.agent.SessionState; +import com.google.genkit.ai.agent.SessionStore; +import com.google.genkit.ai.agent.SessionStoreOptions; +import com.google.genkit.ai.agent.SnapshotMutator; +import com.google.genkit.ai.agent.SnapshotStatus; +import com.google.genkit.ai.agent.SnapshotSubscriber; +import com.google.genkit.ai.agent.internal.SnapshotSharding; +import com.google.genkit.core.GenkitException; +import com.google.genkit.core.JsonUtils; +import com.google.genkit.core.jsonpatch.JsonPatch; +import java.nio.charset.StandardCharsets; +import java.util.ArrayList; +import java.util.List; +import java.util.UUID; +import java.util.concurrent.Executors; +import java.util.concurrent.ScheduledExecutorService; +import java.util.concurrent.ScheduledFuture; +import java.util.concurrent.TimeUnit; +import java.util.function.Consumer; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +/** + * Azure Cosmos DB-backed implementation of {@link SessionStore} and {@link SnapshotSubscriber}. + * + *

Persists session snapshots with the same sharded checkpoint + diff + pointer layout as the + * Firestore backend (see {@code FirestoreSessionStore}), sharing the pure logic in {@link + * SnapshotSharding}. + * + *

Storage layout (single container)

+ * + *

All documents live in one container (default database {@code genkit}, container {@code + * genkit-sessions}) partitioned by {@code /pk} (the per-tenant prefix, default {@code "global"}). + * Each document's {@code id} discriminates the record kind: + * + *

    + *
  • {@code SNAP_} — one metadata document per snapshot. {@code kind} is {@code + * "checkpoint"} or {@code "diff"}; carries {@code checkpointId}, {@code + * checkpointShardCount}, {@code segmentPath}, {@code statePatch} (RFC-6902 as a JSON string + * for diffs) and {@code error} (JSON string). + *
  • {@code SHARD__} — a shard of the checkpoint state JSON. + *
  • {@code PTR_} — the current leaf pointer for a session. + *
+ * + *

Cosmos document ids may not contain {@code #}, so the layout uses {@code _} separators; ids + * are only ever constructed and looked up (never parsed back into components). + * + *

Concurrency

+ * + *

Snapshot and shard documents are idempotent by id (upserted); the snapshot document uses ETag + * ({@code If-Match}) optimistic concurrency when updating an existing id, re-applying the (pure) + * mutator on {@code 412}. The session pointer is advanced monotonically (never backward) under ETag + * concurrency. Shards are written before the snapshot document, which is written before the pointer + * flips, so a reader following the pointer always sees complete data. + * + * @param the type of custom session state + */ +public final class CosmosSessionStore implements SessionStore, SnapshotSubscriber { + + private static final Logger logger = LoggerFactory.getLogger(CosmosSessionStore.class); + private static final ObjectMapper MAPPER = JsonUtils.getObjectMapper(); + + static final String KIND_CHECKPOINT = "checkpoint"; + static final String KIND_DIFF = "diff"; + + private static final int MAX_ATTEMPTS = 5; + private static final int STATUS_NOT_FOUND = 404; + private static final int STATUS_CONFLICT = 409; + private static final int STATUS_PRECONDITION_FAILED = 412; + + private final CosmosContainer container; + private final CosmosSessionStoreOptions options; + private final ScheduledExecutorService scheduler; + + /** + * Creates a store with default options. + * + * @param client the Cosmos DB client + */ + public CosmosSessionStore(CosmosClient client) { + this(client, CosmosSessionStoreOptions.defaults()); + } + + /** + * Creates a store. + * + * @param client the Cosmos DB client + * @param options the store options + */ + public CosmosSessionStore(CosmosClient client, CosmosSessionStoreOptions options) { + if (client == null) { + throw new IllegalArgumentException("CosmosClient must be non-null"); + } + if (options == null) { + throw new IllegalArgumentException("options must be non-null"); + } + this.options = options; + if (options.isCreateIfNotExists()) { + client.createDatabaseIfNotExists(options.getDatabaseName()); + CosmosDatabase database = client.getDatabase(options.getDatabaseName()); + database.createContainerIfNotExists(options.getContainerName(), "/pk"); + this.container = database.getContainer(options.getContainerName()); + } else { + this.container = + client.getDatabase(options.getDatabaseName()).getContainer(options.getContainerName()); + } + this.scheduler = + Executors.newSingleThreadScheduledExecutor( + r -> { + Thread t = new Thread(r, "genkit-cosmos-session-store-poll"); + t.setDaemon(true); + return t; + }); + } + + // ────────────────────────────────────────────────────────────────────────── + // Id helpers + // ────────────────────────────────────────────────────────────────────────── + + private String prefix(SessionStoreOptions opts) { + String p = + options.getSnapshotPathPrefix().apply(opts != null ? opts : SessionStoreOptions.empty()); + return (p == null || p.isBlank()) ? "global" : p; + } + + private static String snapId(String snapshotId) { + return "SNAP_" + snapshotId; + } + + private static String shardId(String checkpointId, int index) { + return "SHARD_" + checkpointId + "_" + index; + } + + private static String ptrId(String sessionId) { + return "PTR_" + sessionId; + } + + // ────────────────────────────────────────────────────────────────────────── + // SnapshotWriter + // ────────────────────────────────────────────────────────────────────────── + + /** + * {@inheritDoc} + * + *

Implements the same identity/sessionId/status defaulting contract as the reference stores, + * then writes the snapshot (checkpoint shards + metadata document) and advances the session + * pointer. + */ + @Override + public String saveSnapshot( + String snapshotId, SnapshotMutator mutator, SessionStoreOptions storeOpts) { + String prefix = prefix(storeOpts); + PartitionKey pk = new PartitionKey(prefix); + + for (int attempt = 1; ; attempt++) { + // 1. Read existing snapshot (+ etag) and reconstruct state. + SessionSnapshot existing = null; + String existingSessionId = null; + String existingEtag = null; + if (snapshotId != null) { + SnapshotSharding.validateId(snapshotId); + CosmosItemResponse resp = readItemOrNull(pk, snapId(snapshotId)); + if (resp != null) { + ObjectNode item = resp.getItem(); + existing = readSnapshot(prefix, item); + existingSessionId = existing.getSessionId(); + existingEtag = resp.getETag(); + } + } + + // 2. Apply mutator (pure — safe to re-run on conflict retry). + SessionSnapshot result = mutator.apply(existing); + if (result == null) { + return null; + } + + // 3. Identity / sessionId / status defaulting (mirror the reference stores). + String finalId; + if (snapshotId != null) { + finalId = snapshotId; + } else if (result.getSnapshotId() != null && !result.getSnapshotId().isBlank()) { + finalId = result.getSnapshotId(); + } else { + finalId = UUID.randomUUID().toString(); + } + SnapshotSharding.validateId(finalId); + result.setSnapshotId(finalId); + + if (existingSessionId != null) { + result.setSessionId(existingSessionId); + } + if (result.getSessionId() == null && result.getState() != null) { + result.setSessionId(result.getState().getSessionId()); + } + if (result.getSessionId() == null || result.getSessionId().isBlank()) { + throw GenkitException.builder() + .message("snapshot requires sessionId") + .errorCode("INVALID_ARGUMENT") + .build(); + } + if (result.getStatus() == null) { + result.setStatus(SnapshotStatus.COMPLETED); + } + + // 4. Resolve parent metadata to decide checkpoint vs diff. + String parentId = result.getParentId(); + JsonNode newState = stateToJson(result.getState()); + + ParentInfo parent = null; + if (parentId != null && !parentId.isBlank()) { + ObjectNode parentDoc = readDoc(pk, snapId(parentId)); + if (parentDoc != null) { + parent = loadParentInfo(prefix, parentDoc); + } + } + + boolean parentExists = parent != null; + int depthFromCheckpoint = 0; + JsonNode statePatch = null; + int diffSizeBytes = 0; + if (parent != null) { + depthFromCheckpoint = parent.segmentPath.size() + 1; + statePatch = JsonPatch.diff(parent.state, newState); + diffSizeBytes = jsonBytes(statePatch); + } + + boolean checkpoint = + SnapshotSharding.shouldCheckpoint( + parentExists, + depthFromCheckpoint, + options.getCheckpointInterval(), + diffSizeBytes, + options.getShardSize()); + + // 5. Build the snapshot metadata document (+ shard docs for a checkpoint). + ObjectNode doc = baseDoc(prefix, result); + String checkpointId; + int checkpointShardCount; + List segmentPath; + + if (checkpoint) { + checkpointId = finalId; + String stateJson = writeJson(newState); + List shards = SnapshotSharding.shardString(stateJson, options.getShardSize()); + checkpointShardCount = shards.size(); + segmentPath = new ArrayList<>(); + for (int i = 0; i < shards.size(); i++) { + ObjectNode shardDoc = MAPPER.createObjectNode(); + shardDoc.put("id", shardId(checkpointId, i)); + shardDoc.put("pk", prefix); + shardDoc.put("checkpointId", checkpointId); + shardDoc.put("index", i); + shardDoc.put("data", shards.get(i)); + container.upsertItem(shardDoc, pk, new CosmosItemRequestOptions()); + } + doc.put("kind", KIND_CHECKPOINT); + } else { + ParentInfo p = parent; + if (p == null) { + throw new GenkitException("internal: diff path without parent"); + } + checkpointId = p.checkpointId; + checkpointShardCount = p.checkpointShardCount; + segmentPath = new ArrayList<>(p.segmentPath); + segmentPath.add(finalId); + doc.put("kind", KIND_DIFF); + doc.put("statePatch", writeJson(statePatch)); + } + + doc.put("checkpointId", checkpointId); + doc.put("checkpointShardCount", checkpointShardCount); + putStringArray(doc, "segmentPath", segmentPath); + + // 6. Write the snapshot document with optimistic concurrency. + try { + if (existingEtag == null) { + container.createItem(doc, pk, new CosmosItemRequestOptions()); + } else { + container.replaceItem( + doc, + snapId(finalId), + pk, + new CosmosItemRequestOptions().setIfMatchETag(existingEtag)); + } + } catch (CosmosException e) { + if ((e.getStatusCode() == STATUS_CONFLICT + || e.getStatusCode() == STATUS_PRECONDITION_FAILED) + && attempt < MAX_ATTEMPTS) { + continue; // concurrent writer won the race; re-read and retry the pure mutator. + } + throw new GenkitException("Failed to save snapshot: " + e.getMessage(), e); + } + + // 7. Advance the session pointer (never backward). + advancePointer( + prefix, + result.getSessionId(), + finalId, + result.getCreatedAt(), + result.getUpdatedAt(), + checkpointId, + checkpointShardCount, + segmentPath); + + return finalId; + } + } + + /** Advances the session pointer to the new leaf unless the stored pointer is already newer. */ + private void advancePointer( + String prefix, + String sessionId, + String snapshotId, + String createdAt, + String updatedAt, + String checkpointId, + int checkpointShardCount, + List segmentPath) { + PartitionKey pk = new PartitionKey(prefix); + String id = ptrId(sessionId); + + for (int attempt = 1; ; attempt++) { + ObjectNode existing = null; + String etag = null; + CosmosItemResponse resp = readItemOrNull(pk, id); + if (resp != null) { + existing = resp.getItem(); + etag = resp.getETag(); + String currentLeaf = getString(existing, "currentSnapshotId"); + String currentCreatedAt = getString(existing, "currentCreatedAt"); + boolean sameLeaf = snapshotId.equals(currentLeaf); + boolean newer = + createdAt == null + || currentCreatedAt == null + || createdAt.compareTo(currentCreatedAt) >= 0; + if (!sameLeaf && !newer) { + return; // stored pointer is already newer; don't move backward. + } + } + + ObjectNode ptr = MAPPER.createObjectNode(); + ptr.put("id", id); + ptr.put("pk", prefix); + ptr.put("currentSnapshotId", snapshotId); + ptr.put("checkpointId", checkpointId); + ptr.put("checkpointShardCount", checkpointShardCount); + putStringArray(ptr, "segmentPath", segmentPath); + if (createdAt != null) { + ptr.put("currentCreatedAt", createdAt); + } + if (updatedAt != null) { + ptr.put("updatedAt", updatedAt); + } + + try { + if (existing == null) { + container.createItem(ptr, pk, new CosmosItemRequestOptions()); + } else { + container.replaceItem(ptr, id, pk, new CosmosItemRequestOptions().setIfMatchETag(etag)); + } + return; + } catch (CosmosException e) { + boolean conflict = + e.getStatusCode() == STATUS_CONFLICT || e.getStatusCode() == STATUS_PRECONDITION_FAILED; + if (conflict && attempt < MAX_ATTEMPTS) { + continue; // concurrent pointer write; re-read and re-evaluate never-backward. + } + if (conflict) { + logger.debug("Pointer for session {} not advanced (lost concurrency race)", sessionId); + return; + } + throw new GenkitException("Failed to advance session pointer: " + e.getMessage(), e); + } + } + } + + // ────────────────────────────────────────────────────────────────────────── + // SnapshotReader + // ────────────────────────────────────────────────────────────────────────── + + /** + * {@inheritDoc} + * + *

By {@code snapshotId}: loads the snapshot document and reconstructs its state from the + * checkpoint shards + ordered {@code segmentPath} diffs. By {@code sessionId}: reads the pointer, + * then loads and reconstructs the pointed snapshot. + */ + @Override + public SessionSnapshot getSnapshot(GetSnapshotOptions opts) { + if (opts == null) { + return null; + } + SessionStoreOptions storeOpts = SessionStoreOptions.empty(); + String prefix = prefix(storeOpts); + PartitionKey pk = new PartitionKey(prefix); + + if (opts.getSnapshotId() != null) { + SnapshotSharding.validateId(opts.getSnapshotId()); + ObjectNode doc = readDoc(pk, snapId(opts.getSnapshotId())); + return doc == null ? null : readSnapshot(prefix, doc); + } + if (opts.getSessionId() != null) { + SnapshotSharding.validateId(opts.getSessionId()); + ObjectNode pointer = readDoc(pk, ptrId(opts.getSessionId())); + if (pointer == null) { + return null; + } + String leafId = getString(pointer, "currentSnapshotId"); + if (leafId == null) { + return null; + } + ObjectNode doc = readDoc(pk, snapId(leafId)); + return doc == null ? null : readSnapshot(prefix, doc); + } + return null; + } + + // ────────────────────────────────────────────────────────────────────────── + // SnapshotSubscriber (polling) + // ────────────────────────────────────────────────────────────────────────── + + /** + * {@inheritDoc} + * + *

The subscription polls {@link #getSnapshot} on a shared daemon scheduler and fires the + * callback whenever the serialized snapshot content changes. The callback also fires immediately + * if the snapshot already exists. + */ + @Override + public AutoCloseable onSnapshotStateChange( + String snapshotId, Consumer> cb, SessionStoreOptions storeOpts) { + SnapshotSharding.validateId(snapshotId); + GetSnapshotOptions get = GetSnapshotOptions.builder().snapshotId(snapshotId).build(); + + final String[] lastContent = {null}; + SessionSnapshot initial = getSnapshot(get); + if (initial != null) { + lastContent[0] = serializeQuietly(initial); + cb.accept(initial); + } + + ScheduledFuture future = + scheduler.scheduleAtFixedRate( + () -> { + try { + SessionSnapshot snap = getSnapshot(get); + if (snap == null) { + return; + } + String content = serializeQuietly(snap); + if (!content.equals(lastContent[0])) { + lastContent[0] = content; + cb.accept(snap); + } + } catch (Exception e) { + // Swallow poll errors — don't kill the scheduler thread. + } + }, + options.getPollIntervalMs(), + options.getPollIntervalMs(), + TimeUnit.MILLISECONDS); + + return () -> future.cancel(false); + } + + // ────────────────────────────────────────────────────────────────────────── + // Document (de)serialization + // ────────────────────────────────────────────────────────────────────────── + + /** Builds the non-state base document for a snapshot (id/pk + metadata). */ + private ObjectNode baseDoc(String prefix, SessionSnapshot snap) { + ObjectNode data = MAPPER.createObjectNode(); + data.put("id", snapId(snap.getSnapshotId())); + data.put("pk", prefix); + data.put("snapshotId", snap.getSnapshotId()); + data.put("sessionId", snap.getSessionId()); + if (snap.getParentId() != null) { + data.put("parentId", snap.getParentId()); + } + if (snap.getCreatedAt() != null) { + data.put("createdAt", snap.getCreatedAt()); + } + if (snap.getUpdatedAt() != null) { + data.put("updatedAt", snap.getUpdatedAt()); + } + if (snap.getHeartbeatAt() != null) { + data.put("heartbeatAt", snap.getHeartbeatAt()); + } + if (snap.getStatus() != null) { + data.put("status", snap.getStatus().getValue()); + } + if (snap.getFinishReason() != null) { + data.put("finishReason", snap.getFinishReason().getValue()); + } + if (snap.getError() != null) { + data.put("error", writeJson(snap.getError())); + } + return data; + } + + /** Holds the reconstructed parent state and its checkpoint lineage. */ + private static final class ParentInfo { + JsonNode state; + String checkpointId; + int checkpointShardCount; + List segmentPath; + } + + /** Loads a parent snapshot's checkpoint lineage and reconstructs its state. */ + private ParentInfo loadParentInfo(String prefix, ObjectNode doc) { + ParentInfo info = new ParentInfo(); + info.checkpointId = getString(doc, "checkpointId"); + info.checkpointShardCount = getInt(doc, "checkpointShardCount"); + info.segmentPath = getStringList(doc, "segmentPath"); + info.state = + reconstructFullState( + prefix, info.checkpointId, info.checkpointShardCount, info.segmentPath); + return info; + } + + /** Reads and fully reconstructs a snapshot from its metadata document. */ + private SessionSnapshot readSnapshot(String prefix, ObjectNode doc) { + String checkpointId = getString(doc, "checkpointId"); + int checkpointShardCount = getInt(doc, "checkpointShardCount"); + List segmentPath = getStringList(doc, "segmentPath"); + JsonNode state = reconstructFullState(prefix, checkpointId, checkpointShardCount, segmentPath); + return docToSnapshot(doc, state); + } + + /** + * Reconstructs full state: loads the checkpoint shards (concatenate, parse) then applies the + * {@code segmentPath} diffs in order. + */ + private JsonNode reconstructFullState( + String prefix, String checkpointId, int checkpointShardCount, List segmentPath) { + if (checkpointId == null) { + return NullNode.getInstance(); + } + PartitionKey pk = new PartitionKey(prefix); + List shardContents = new ArrayList<>(); + for (int i = 0; i < checkpointShardCount; i++) { + ObjectNode shard = readDoc(pk, shardId(checkpointId, i)); + shardContents.add(shard != null ? getString(shard, "data") : ""); + } + String checkpointJson = SnapshotSharding.reassembleShards(shardContents); + + List diffs = new ArrayList<>(); + for (String diffId : segmentPath) { + ObjectNode diffDoc = readDoc(pk, snapId(diffId)); + if (diffDoc != null) { + String patch = getString(diffDoc, "statePatch"); + if (patch != null) { + diffs.add(patch); + } + } + } + try { + return SnapshotSharding.reconstructState(checkpointJson, diffs); + } catch (Exception e) { + throw new GenkitException("Failed to reconstruct session state: " + e.getMessage(), e); + } + } + + /** Builds a {@link SessionSnapshot} from a metadata document and reconstructed state. */ + @SuppressWarnings("unchecked") + private SessionSnapshot docToSnapshot(ObjectNode doc, JsonNode state) { + SessionSnapshot.Builder builder = SessionSnapshot.builder(); + builder.snapshotId(getString(doc, "snapshotId")); + builder.sessionId(getString(doc, "sessionId")); + builder.parentId(getString(doc, "parentId")); + builder.createdAt(getString(doc, "createdAt")); + builder.updatedAt(getString(doc, "updatedAt")); + builder.heartbeatAt(getString(doc, "heartbeatAt")); + String status = getString(doc, "status"); + if (status != null) { + builder.status(SnapshotStatus.fromValueOrCompleted(status)); + } + String finishReason = getString(doc, "finishReason"); + if (finishReason != null) { + builder.finishReason(AgentFinishReason.fromValue(finishReason)); + } + String errorJson = getString(doc, "error"); + if (errorJson != null) { + try { + builder.error(MAPPER.readValue(errorJson, RuntimeError.class)); + } catch (Exception e) { + throw new GenkitException("Failed to parse snapshot error: " + e.getMessage(), e); + } + } + if (state != null && !state.isNull()) { + try { + builder.state((SessionState) MAPPER.treeToValue(state, SessionState.class)); + } catch (Exception e) { + throw new GenkitException("Failed to parse session state: " + e.getMessage(), e); + } + } + return builder.build(); + } + + /** Serializes session state to a JSON node (null state → JSON null). */ + private JsonNode stateToJson(SessionState state) { + if (state == null) { + return NullNode.getInstance(); + } + return MAPPER.valueToTree(state); + } + + // ────────────────────────────────────────────────────────────────────────── + // Low-level Cosmos + node helpers + // ────────────────────────────────────────────────────────────────────────── + + /** Reads an item and returns its full response, or {@code null} when it does not exist. */ + private CosmosItemResponse readItemOrNull(PartitionKey pk, String id) { + try { + return container.readItem(id, pk, ObjectNode.class); + } catch (CosmosException e) { + if (e.getStatusCode() == STATUS_NOT_FOUND) { + return null; + } + throw new GenkitException("Failed to read Cosmos item " + id + ": " + e.getMessage(), e); + } + } + + /** Reads an item body, or {@code null} when it does not exist. */ + private ObjectNode readDoc(PartitionKey pk, String id) { + CosmosItemResponse resp = readItemOrNull(pk, id); + return resp == null ? null : resp.getItem(); + } + + private static void putStringArray(ObjectNode node, String name, List values) { + ArrayNode arr = node.putArray(name); + for (String v : values) { + arr.add(v); + } + } + + private static String getString(ObjectNode node, String name) { + JsonNode v = node.get(name); + return (v == null || v.isNull()) ? null : v.asText(); + } + + private static int getInt(ObjectNode node, String name) { + JsonNode v = node.get(name); + return (v == null || v.isNull()) ? 0 : v.asInt(); + } + + private static List getStringList(ObjectNode node, String name) { + List out = new ArrayList<>(); + JsonNode v = node.get(name); + if (v != null && v.isArray()) { + for (JsonNode e : v) { + out.add(e.asText()); + } + } + return out; + } + + private static String writeJson(Object value) { + try { + return MAPPER.writeValueAsString(value); + } catch (Exception e) { + throw new GenkitException("Failed to serialize value: " + e.getMessage(), e); + } + } + + private static int jsonBytes(Object value) { + return writeJson(value).getBytes(StandardCharsets.UTF_8).length; + } + + private static String serializeQuietly(SessionSnapshot snap) { + try { + return MAPPER.writeValueAsString(MAPPER.valueToTree(snap)); + } catch (Exception e) { + return ""; + } + } +} diff --git a/plugins/azure-foundry/src/main/java/com/google/genkit/plugins/azurefoundry/session/CosmosSessionStoreOptions.java b/plugins/azure-foundry/src/main/java/com/google/genkit/plugins/azurefoundry/session/CosmosSessionStoreOptions.java new file mode 100644 index 000000000..1e2fa5617 --- /dev/null +++ b/plugins/azure-foundry/src/main/java/com/google/genkit/plugins/azurefoundry/session/CosmosSessionStoreOptions.java @@ -0,0 +1,273 @@ +/* + * 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. + * + * SPDX-License-Identifier: Apache-2.0 + */ + +package com.google.genkit.plugins.azurefoundry.session; + +import com.google.genkit.ai.agent.SessionStoreOptions; +import java.util.function.Function; + +/** + * Configuration for {@link CosmosSessionStore}. + * + *

The store persists all documents in a single Cosmos DB container (default database {@value + * #DEFAULT_DATABASE}, container {@value #DEFAULT_CONTAINER}) partitioned by {@code /pk} (the + * per-tenant prefix). Each document's {@code id} discriminates snapshot, shard, and pointer + * records. All documents for a tenant share the same {@code pk}, derived from {@link + * #getSnapshotPathPrefix()} (default {@code "global"}). + * + *

The default {@link #getShardSize()} is {@value #DEFAULT_SHARD_SIZE} bytes, kept safely under + * the Cosmos DB 2 MB document-size limit; because the store forces a checkpoint whenever a + * diff would exceed the shard size, diff documents stay bounded too. + */ +public final class CosmosSessionStoreOptions { + + /** Default database name. */ + public static final String DEFAULT_DATABASE = "genkit"; + + /** Default container name. */ + public static final String DEFAULT_CONTAINER = "genkit-sessions"; + + /** Default number of turns between full checkpoints. */ + public static final int DEFAULT_CHECKPOINT_INTERVAL = 25; + + /** Default shard size in bytes for checkpoint state (1 MiB, under the 2 MB document cap). */ + public static final int DEFAULT_SHARD_SIZE = 1024 * 1024; + + /** Default subscription poll interval in milliseconds. */ + public static final long DEFAULT_POLL_INTERVAL_MS = 2000L; + + private final String databaseName; + private final String containerName; + private final int checkpointInterval; + private final int shardSize; + private final Function snapshotPathPrefix; + private final boolean createIfNotExists; + private final long pollIntervalMs; + + private CosmosSessionStoreOptions(Builder builder) { + this.databaseName = builder.databaseName; + this.containerName = builder.containerName; + this.checkpointInterval = builder.checkpointInterval; + this.shardSize = builder.shardSize; + this.snapshotPathPrefix = builder.snapshotPathPrefix; + this.createIfNotExists = builder.createIfNotExists; + this.pollIntervalMs = builder.pollIntervalMs; + } + + /** + * Returns the Cosmos DB database name (default {@value #DEFAULT_DATABASE}). + * + * @return the database name + */ + public String getDatabaseName() { + return databaseName; + } + + /** + * Returns the Cosmos DB container name (default {@value #DEFAULT_CONTAINER}). + * + * @return the container name + */ + public String getContainerName() { + return containerName; + } + + /** + * Returns the number of turns between full checkpoints (default {@value + * #DEFAULT_CHECKPOINT_INTERVAL}). + * + * @return the checkpoint interval + */ + public int getCheckpointInterval() { + return checkpointInterval; + } + + /** + * Returns the shard size in bytes for checkpoint state (default {@value #DEFAULT_SHARD_SIZE}). + * + * @return the shard size in bytes + */ + public int getShardSize() { + return shardSize; + } + + /** + * Returns the function that derives the per-tenant partition-key prefix from the per-request + * store options (default {@code o -> "global"}). + * + * @return the prefix function + */ + public Function getSnapshotPathPrefix() { + return snapshotPathPrefix; + } + + /** + * Returns whether the store should create the database and container on first use if they do not + * exist (default {@code false}). + * + * @return {@code true} if the database/container should be created when missing + */ + public boolean isCreateIfNotExists() { + return createIfNotExists; + } + + /** + * Returns the subscription poll interval in milliseconds (default {@value + * #DEFAULT_POLL_INTERVAL_MS}). + * + * @return the poll interval in milliseconds + */ + public long getPollIntervalMs() { + return pollIntervalMs; + } + + /** + * Returns default options. + * + * @return a {@code CosmosSessionStoreOptions} with all defaults + */ + public static CosmosSessionStoreOptions defaults() { + return builder().build(); + } + + /** + * Creates a new builder. + * + * @return a new builder + */ + public static Builder builder() { + return new Builder(); + } + + /** Builder for {@link CosmosSessionStoreOptions}. */ + public static final class Builder { + private String databaseName = DEFAULT_DATABASE; + private String containerName = DEFAULT_CONTAINER; + private int checkpointInterval = DEFAULT_CHECKPOINT_INTERVAL; + private int shardSize = DEFAULT_SHARD_SIZE; + private Function snapshotPathPrefix = o -> "global"; + private boolean createIfNotExists = false; + private long pollIntervalMs = DEFAULT_POLL_INTERVAL_MS; + + private Builder() {} + + /** + * Sets the Cosmos DB database name. + * + * @param databaseName the database name + * @return this builder + */ + public Builder databaseName(String databaseName) { + this.databaseName = databaseName; + return this; + } + + /** + * Sets the Cosmos DB container name. + * + * @param containerName the container name + * @return this builder + */ + public Builder containerName(String containerName) { + this.containerName = containerName; + return this; + } + + /** + * Sets the number of turns between full checkpoints. + * + * @param checkpointInterval the checkpoint interval (must be {@code >= 1}) + * @return this builder + */ + public Builder checkpointInterval(int checkpointInterval) { + this.checkpointInterval = checkpointInterval; + return this; + } + + /** + * Sets the shard size in bytes for checkpoint state (must stay under the 2 MB document cap). + * + * @param shardSize the shard size in bytes (must be {@code >= 1}) + * @return this builder + */ + public Builder shardSize(int shardSize) { + this.shardSize = shardSize; + return this; + } + + /** + * Sets the function that derives the per-tenant partition-key prefix. + * + * @param snapshotPathPrefix the prefix function + * @return this builder + */ + public Builder snapshotPathPrefix(Function snapshotPathPrefix) { + this.snapshotPathPrefix = snapshotPathPrefix; + return this; + } + + /** + * Sets whether to create the database and container on first use if they do not exist. + * + * @param createIfNotExists whether to create the database/container when missing + * @return this builder + */ + public Builder createIfNotExists(boolean createIfNotExists) { + this.createIfNotExists = createIfNotExists; + return this; + } + + /** + * Sets the subscription poll interval in milliseconds. + * + * @param pollIntervalMs the poll interval (must be {@code >= 1}) + * @return this builder + */ + public Builder pollIntervalMs(long pollIntervalMs) { + this.pollIntervalMs = pollIntervalMs; + return this; + } + + /** + * Builds a new {@code CosmosSessionStoreOptions}. + * + * @return a new options instance + */ + public CosmosSessionStoreOptions build() { + if (databaseName == null || databaseName.isBlank()) { + throw new IllegalArgumentException("databaseName must be non-empty"); + } + if (containerName == null || containerName.isBlank()) { + throw new IllegalArgumentException("containerName must be non-empty"); + } + if (checkpointInterval < 1) { + throw new IllegalArgumentException("checkpointInterval must be >= 1"); + } + if (shardSize < 1) { + throw new IllegalArgumentException("shardSize must be >= 1"); + } + if (snapshotPathPrefix == null) { + throw new IllegalArgumentException("snapshotPathPrefix must be non-null"); + } + if (pollIntervalMs < 1) { + throw new IllegalArgumentException("pollIntervalMs must be >= 1"); + } + return new CosmosSessionStoreOptions(this); + } + } +} diff --git a/plugins/azure-foundry/src/main/java/com/google/genkit/plugins/azurefoundry/session/package-info.java b/plugins/azure-foundry/src/main/java/com/google/genkit/plugins/azurefoundry/session/package-info.java new file mode 100644 index 000000000..7443ca427 --- /dev/null +++ b/plugins/azure-foundry/src/main/java/com/google/genkit/plugins/azurefoundry/session/package-info.java @@ -0,0 +1,27 @@ +/* + * 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. + * + * SPDX-License-Identifier: Apache-2.0 + */ + +/** + * Azure Cosmos DB-backed agent session persistence. + * + *

{@link com.google.genkit.plugins.azurefoundry.session.CosmosSessionStore} implements the + * Genkit {@code SessionStore} contract using the sharded checkpoint + RFC-6902 diff + pointer + * layout shared with the Firestore and DynamoDB backends. Construct it directly and pass it to an + * agent via {@code AgentConfig.store(...)}. + */ +package com.google.genkit.plugins.azurefoundry.session; diff --git a/plugins/azure-foundry/src/test/java/com/google/genkit/plugins/azurefoundry/session/CosmosSessionStoreOptionsTest.java b/plugins/azure-foundry/src/test/java/com/google/genkit/plugins/azurefoundry/session/CosmosSessionStoreOptionsTest.java new file mode 100644 index 000000000..7d3153e66 --- /dev/null +++ b/plugins/azure-foundry/src/test/java/com/google/genkit/plugins/azurefoundry/session/CosmosSessionStoreOptionsTest.java @@ -0,0 +1,79 @@ +/* + * 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. + * + * SPDX-License-Identifier: Apache-2.0 + */ + +package com.google.genkit.plugins.azurefoundry.session; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertThrows; + +import com.google.genkit.ai.agent.SessionStoreOptions; +import org.junit.jupiter.api.Test; + +/** Unit tests for {@link CosmosSessionStoreOptions}. */ +class CosmosSessionStoreOptionsTest { + + @Test + void defaultsAreSane() { + CosmosSessionStoreOptions o = CosmosSessionStoreOptions.defaults(); + assertEquals("genkit", o.getDatabaseName()); + assertEquals("genkit-sessions", o.getContainerName()); + assertEquals(25, o.getCheckpointInterval()); + assertEquals(1024 * 1024, o.getShardSize()); + assertEquals("global", o.getSnapshotPathPrefix().apply(SessionStoreOptions.empty())); + assertFalse(o.isCreateIfNotExists()); + assertEquals(2000L, o.getPollIntervalMs()); + } + + @Test + void customBuilder() { + CosmosSessionStoreOptions o = + CosmosSessionStoreOptions.builder() + .databaseName("my-db") + .containerName("my-sessions") + .checkpointInterval(10) + .shardSize(4096) + .snapshotPathPrefix(so -> "tenant-1") + .createIfNotExists(true) + .pollIntervalMs(500) + .build(); + assertEquals("my-db", o.getDatabaseName()); + assertEquals("my-sessions", o.getContainerName()); + assertEquals(10, o.getCheckpointInterval()); + assertEquals(4096, o.getShardSize()); + assertEquals("tenant-1", o.getSnapshotPathPrefix().apply(SessionStoreOptions.empty())); + assertEquals(true, o.isCreateIfNotExists()); + assertEquals(500L, o.getPollIntervalMs()); + } + + @Test + void builderValidates() { + assertThrows( + IllegalArgumentException.class, + () -> CosmosSessionStoreOptions.builder().databaseName("").build()); + assertThrows( + IllegalArgumentException.class, + () -> CosmosSessionStoreOptions.builder().containerName("").build()); + assertThrows( + IllegalArgumentException.class, + () -> CosmosSessionStoreOptions.builder().shardSize(0).build()); + assertThrows( + IllegalArgumentException.class, + () -> CosmosSessionStoreOptions.builder().pollIntervalMs(0).build()); + } +} diff --git a/plugins/azure-foundry/src/test/java/com/google/genkit/plugins/azurefoundry/session/CosmosSessionStoreTest.java b/plugins/azure-foundry/src/test/java/com/google/genkit/plugins/azurefoundry/session/CosmosSessionStoreTest.java new file mode 100644 index 000000000..999c5b77c --- /dev/null +++ b/plugins/azure-foundry/src/test/java/com/google/genkit/plugins/azurefoundry/session/CosmosSessionStoreTest.java @@ -0,0 +1,190 @@ +/* + * 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. + * + * SPDX-License-Identifier: Apache-2.0 + */ + +package com.google.genkit.plugins.azurefoundry.session; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertNull; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assumptions.assumeTrue; + +import com.azure.cosmos.CosmosClient; +import com.azure.cosmos.CosmosClientBuilder; +import com.google.genkit.ai.Message; +import com.google.genkit.ai.agent.GetSnapshotOptions; +import com.google.genkit.ai.agent.SessionSnapshot; +import com.google.genkit.ai.agent.SessionState; +import com.google.genkit.ai.agent.SessionStoreOptions; +import com.google.genkit.ai.agent.SnapshotStatus; +import com.google.genkit.core.GenkitException; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.UUID; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; + +/** + * Tests for {@link CosmosSessionStore}. + * + *

Integration tests are gated on the {@code COSMOS_ENDPOINT} and {@code COSMOS_KEY} environment + * variables (e.g. the local Cosmos DB emulator or an Azure account). When either is unset the tests + * are skipped via {@link org.junit.jupiter.api.Assumptions}. + */ +class CosmosSessionStoreTest { + + private static final String ENDPOINT = System.getenv("COSMOS_ENDPOINT"); + private static final String KEY = System.getenv("COSMOS_KEY"); + + private CosmosClient client; + private CosmosSessionStore> store; + + private static boolean configured() { + return ENDPOINT != null && !ENDPOINT.isEmpty() && KEY != null && !KEY.isEmpty(); + } + + @BeforeEach + void setUp() { + if (!configured()) { + return; // integration tests skip via assumeTrue + } + client = new CosmosClientBuilder().endpoint(ENDPOINT).key(KEY).gatewayMode().buildClient(); + String container = "genkit-sessions-test-" + UUID.randomUUID().toString().substring(0, 8); + store = + new CosmosSessionStore<>( + client, + CosmosSessionStoreOptions.builder() + .databaseName("genkit-test") + .containerName(container) + .checkpointInterval(3) + .createIfNotExists(true) + .build()); + } + + @AfterEach + void tearDown() { + if (client != null) { + client.close(); + } + } + + private static SessionSnapshot> snapshotWithState( + String sessionId, String parentId, Map custom) { + SessionState> state = + SessionState.>builder() + .sessionId(sessionId) + .messages(List.of(Message.user("hello"))) + .custom(custom) + .build(); + return SessionSnapshot.>builder() + .sessionId(sessionId) + .parentId(parentId) + .status(SnapshotStatus.COMPLETED) + .state(state) + .build(); + } + + @Test + void saveThenGetBySnapshotIdRoundTrips() { + assumeTrue(configured()); + String sessionId = "s-" + UUID.randomUUID(); + Map custom = new HashMap<>(); + custom.put("count", 1); + + String id = + store.saveSnapshot( + null, e -> snapshotWithState(sessionId, null, custom), SessionStoreOptions.empty()); + assertNotNull(id); + + SessionSnapshot> got = + store.getSnapshot(GetSnapshotOptions.builder().snapshotId(id).build()); + assertNotNull(got); + assertEquals(sessionId, got.getSessionId()); + assertEquals(SnapshotStatus.COMPLETED, got.getStatus()); + assertEquals(1, got.getState().getMessages().size()); + assertEquals(1, ((Number) got.getState().getCustom().get("count")).intValue()); + } + + @Test + void getBySessionIdReturnsLeaf() { + assumeTrue(configured()); + String sessionId = "s-" + UUID.randomUUID(); + + Map c1 = new HashMap<>(); + c1.put("count", 1); + String id1 = + store.saveSnapshot( + null, e -> snapshotWithState(sessionId, null, c1), SessionStoreOptions.empty()); + Map c2 = new HashMap<>(); + c2.put("count", 2); + String id2 = + store.saveSnapshot( + null, e -> snapshotWithState(sessionId, id1, c2), SessionStoreOptions.empty()); + + SessionSnapshot> latest = + store.getSnapshot(GetSnapshotOptions.builder().sessionId(sessionId).build()); + assertNotNull(latest); + assertEquals(id2, latest.getSnapshotId()); + assertEquals(2, ((Number) latest.getState().getCustom().get("count")).intValue()); + } + + @Test + void diffThenCheckpointReconstructs() { + assumeTrue(configured()); + // checkpointInterval is 3; save 5 turns across checkpoint boundaries and confirm the leaf + // reconstructs correctly (checkpoint shards + segment-path diffs). + String sessionId = "s-" + UUID.randomUUID(); + String parent = null; + String lastId = null; + for (int i = 1; i <= 5; i++) { + Map c = new HashMap<>(); + c.put("count", i); + final String p = parent; + lastId = + store.saveSnapshot( + null, e -> snapshotWithState(sessionId, p, c), SessionStoreOptions.empty()); + parent = lastId; + } + SessionSnapshot> got = + store.getSnapshot(GetSnapshotOptions.builder().snapshotId(lastId).build()); + assertNotNull(got); + assertEquals(5, ((Number) got.getState().getCustom().get("count")).intValue()); + } + + @Test + void rejectsEmptySessionId() { + assumeTrue(configured()); + GenkitException ex = + assertThrows( + GenkitException.class, + () -> + store.saveSnapshot( + null, + e -> snapshotWithState("", null, new HashMap<>()), + SessionStoreOptions.empty())); + assertEquals("INVALID_ARGUMENT", ex.getErrorCode()); + } + + @Test + void mutatorNullIsNoOp() { + assumeTrue(configured()); + assertNull(store.saveSnapshot(null, e -> null, SessionStoreOptions.empty())); + } +} diff --git a/plugins/firebase/src/main/java/com/google/genkit/plugins/firebase/session/FirestoreSessionStore.java b/plugins/firebase/src/main/java/com/google/genkit/plugins/firebase/session/FirestoreSessionStore.java new file mode 100644 index 000000000..9010b1472 --- /dev/null +++ b/plugins/firebase/src/main/java/com/google/genkit/plugins/firebase/session/FirestoreSessionStore.java @@ -0,0 +1,634 @@ +/* + * 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. + * + * SPDX-License-Identifier: Apache-2.0 + */ + +package com.google.genkit.plugins.firebase.session; + +import com.fasterxml.jackson.databind.JsonNode; +import com.fasterxml.jackson.databind.ObjectMapper; +import com.fasterxml.jackson.databind.node.NullNode; +import com.google.cloud.firestore.DocumentReference; +import com.google.cloud.firestore.DocumentSnapshot; +import com.google.cloud.firestore.Firestore; +import com.google.cloud.firestore.ListenerRegistration; +import com.google.cloud.firestore.Transaction; +import com.google.genkit.ai.agent.GetSnapshotOptions; +import com.google.genkit.ai.agent.RuntimeError; +import com.google.genkit.ai.agent.SessionSnapshot; +import com.google.genkit.ai.agent.SessionState; +import com.google.genkit.ai.agent.SessionStore; +import com.google.genkit.ai.agent.SessionStoreOptions; +import com.google.genkit.ai.agent.SnapshotMutator; +import com.google.genkit.ai.agent.SnapshotStatus; +import com.google.genkit.ai.agent.SnapshotSubscriber; +import com.google.genkit.ai.agent.internal.SnapshotSharding; +import com.google.genkit.core.GenkitException; +import com.google.genkit.core.JsonUtils; +import com.google.genkit.core.jsonpatch.JsonPatch; +import java.nio.charset.StandardCharsets; +import java.util.ArrayList; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.UUID; +import java.util.concurrent.ExecutionException; +import java.util.function.Consumer; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +/** + * Firestore-backed implementation of {@link SessionStore} and {@link SnapshotSubscriber}. + * + *

Persists session snapshots with a sharded checkpoint + diff + pointer layout, mirroring the + * upstream Go ({@code go/plugins/firebase/exp/firestore_session_store.go}) and JS ({@code + * js/plugins/google-cloud/src/session-store/firestore.ts}) implementations. + * + *

Storage layout

+ * + *

For a configured {@code collection} and per-tenant {@code prefix} (default {@code "global"}): + * + *

    + *
  • {@code //snapshots/} — one metadata document per snapshot. + * {@code kind} is {@code "checkpoint"} (full state stored out-of-band in shards) or {@code + * "diff"} (RFC-6902 {@code statePatch} from its parent). Each document carries {@code + * checkpointId} (nearest checkpoint ancestor), {@code checkpointShardCount}, {@code + * segmentPath} (ordered diff ids from the checkpoint exclusive → this doc inclusive) and the + * snapshot metadata fields. {@code statePatch} and {@code error} are stored as opaque JSON + * strings because Firestore disallows nested arrays. + *
  • {@code -shards//shards/_} — checkpoint state JSON + * (UTF-8) split into {@code shardSize}-byte chunks. + *
  • {@code -pointers//pointers/} — the current leaf pointer for + * a session ({@code currentSnapshotId}, {@code checkpointId}, {@code checkpointShardCount}, + * {@code segmentPath}, {@code currentCreatedAt}, {@code updatedAt}). Carries no state. + *
+ * + *

Concurrency

+ * + *

{@link #saveSnapshot} runs inside a Firestore transaction (read existing → apply mutator → + * write). The mutator is treated as pure and may be re-invoked on transaction retry. + * + * @param the type of custom session state + */ +public final class FirestoreSessionStore implements SessionStore, SnapshotSubscriber { + + private static final Logger logger = LoggerFactory.getLogger(FirestoreSessionStore.class); + private static final ObjectMapper MAPPER = JsonUtils.getObjectMapper(); + + static final String KIND_CHECKPOINT = "checkpoint"; + static final String KIND_DIFF = "diff"; + + private final Firestore db; + private final FirestoreSessionStoreOptions options; + + /** + * Creates a store with default options. + * + * @param db the Firestore client + */ + public FirestoreSessionStore(Firestore db) { + this(db, FirestoreSessionStoreOptions.defaults()); + } + + /** + * Creates a store. + * + * @param db the Firestore client + * @param options the store options + */ + public FirestoreSessionStore(Firestore db, FirestoreSessionStoreOptions options) { + if (db == null) { + throw new IllegalArgumentException("Firestore client must be non-null"); + } + if (options == null) { + throw new IllegalArgumentException("options must be non-null"); + } + this.db = db; + this.options = options; + } + + // ────────────────────────────────────────────────────────────────────────── + // Firestore path helpers + // ────────────────────────────────────────────────────────────────────────── + + private String prefix(SessionStoreOptions opts) { + String p = + options.getSnapshotPathPrefix().apply(opts != null ? opts : SessionStoreOptions.empty()); + return (p == null || p.isBlank()) ? "global" : p; + } + + private DocumentReference snapshotDoc(String prefix, String snapshotId) { + return db.collection(options.getCollection()) + .document(prefix) + .collection("snapshots") + .document(snapshotId); + } + + private DocumentReference shardDoc(String prefix, String checkpointId, int index) { + return db.collection(options.getCollection() + "-shards") + .document(prefix) + .collection("shards") + .document(checkpointId + "_" + index); + } + + private DocumentReference pointerDoc(String prefix, String sessionId) { + return db.collection(options.getCollection() + "-pointers") + .document(prefix) + .collection("pointers") + .document(sessionId); + } + + // ────────────────────────────────────────────────────────────────────────── + // SnapshotWriter + // ────────────────────────────────────────────────────────────────────────── + + /** + * {@inheritDoc} + * + *

Runs a Firestore transaction implementing the same identity/sessionId/status defaulting + * contract as the in-memory reference store: + * + *

    + *
  1. Read the existing snapshot doc; reconstruct its full state when present. + *
  2. Apply the (pure) mutator. If it returns {@code null}, no write occurs and {@code null} is + * returned. + *
  3. Mint a UUID id when none supplied; preserve {@code sessionId} from the existing row; + * reject empty {@code sessionId} with {@code INVALID_ARGUMENT}; default {@code null} status + * to {@code COMPLETED}. + *
  4. Decide checkpoint vs diff and write the snapshot doc (+ shards for a checkpoint). + *
  5. Advance the session pointer to the new leaf (never backward). + *
+ */ + @Override + public String saveSnapshot( + String snapshotId, SnapshotMutator mutator, SessionStoreOptions options) { + String prefix = prefix(options); + try { + return db.runTransaction( + (Transaction tx) -> { + // 1. Read existing snapshot + reconstruct state. + SessionSnapshot existing = null; + String existingSessionId = null; + if (snapshotId != null) { + SnapshotSharding.validateId(snapshotId); + DocumentSnapshot doc = tx.get(snapshotDoc(prefix, snapshotId)).get(); + if (doc.exists()) { + existing = readSnapshot(tx, prefix, doc); + existingSessionId = existing.getSessionId(); + } + } + + // 2. Apply mutator (pure). + SessionSnapshot result = mutator.apply(existing); + if (result == null) { + return null; + } + + // 3. Identity / sessionId / status defaulting (mirror InMemorySessionStore). + String finalId; + if (snapshotId != null) { + finalId = snapshotId; + } else if (result.getSnapshotId() != null && !result.getSnapshotId().isBlank()) { + finalId = result.getSnapshotId(); + } else { + finalId = UUID.randomUUID().toString(); + } + SnapshotSharding.validateId(finalId); + result.setSnapshotId(finalId); + + if (existingSessionId != null) { + result.setSessionId(existingSessionId); + } + if (result.getSessionId() == null && result.getState() != null) { + result.setSessionId(result.getState().getSessionId()); + } + if (result.getSessionId() == null || result.getSessionId().isBlank()) { + throw GenkitException.builder() + .message("snapshot requires sessionId") + .errorCode("INVALID_ARGUMENT") + .build(); + } + if (result.getStatus() == null) { + result.setStatus(SnapshotStatus.COMPLETED); + } + + // 4. Resolve parent metadata to decide checkpoint vs diff. + String parentId = result.getParentId(); + JsonNode newState = stateToJson(result.getState()); + + ParentInfo parent = null; + if (parentId != null && !parentId.isBlank()) { + DocumentSnapshot parentDoc = tx.get(snapshotDoc(prefix, parentId)).get(); + if (parentDoc.exists()) { + parent = loadParentInfo(tx, prefix, parentDoc); + } + } + + boolean parentExists = parent != null; + int depthFromCheckpoint = 0; + JsonNode statePatch = null; + int diffSizeBytes = 0; + if (parent != null) { + depthFromCheckpoint = parent.segmentPath.size() + 1; + statePatch = JsonPatch.diff(parent.state, newState); + diffSizeBytes = + MAPPER.writeValueAsString(statePatch).getBytes(StandardCharsets.UTF_8).length; + } + + boolean checkpoint = + SnapshotSharding.shouldCheckpoint( + parentExists, + depthFromCheckpoint, + FirestoreSessionStore.this.options.getCheckpointInterval(), + diffSizeBytes, + FirestoreSessionStore.this.options.getShardSize()); + + // Read the session pointer now — Firestore transactions require all reads to + // precede all writes, so this must happen before the tx.set calls below. + DocumentSnapshot currentPointer = + tx.get(pointerDoc(prefix, result.getSessionId())).get(); + + // 5. Write snapshot doc (+shards if checkpoint). + Map docData = baseDocData(result); + String checkpointId; + int checkpointShardCount; + List segmentPath; + + if (checkpoint) { + checkpointId = finalId; + String stateJson = MAPPER.writeValueAsString(newState); + List shards = + SnapshotSharding.shardString( + stateJson, FirestoreSessionStore.this.options.getShardSize()); + checkpointShardCount = shards.size(); + segmentPath = new ArrayList<>(); + for (int i = 0; i < shards.size(); i++) { + Map shardData = new HashMap<>(); + shardData.put("checkpointId", checkpointId); + shardData.put("index", i); + shardData.put("data", shards.get(i)); + tx.set(shardDoc(prefix, checkpointId, i), shardData); + } + docData.put("kind", KIND_CHECKPOINT); + } else { + // shouldCheckpoint() guarantees a non-null parent here (it returns true whenever + // there is no usable parent), but assert defensively for the analyzer. + ParentInfo p = parent; + if (p == null) { + throw new GenkitException("internal: diff path without parent"); + } + checkpointId = p.checkpointId; + checkpointShardCount = p.checkpointShardCount; + segmentPath = new ArrayList<>(p.segmentPath); + segmentPath.add(finalId); + docData.put("kind", KIND_DIFF); + docData.put("statePatch", MAPPER.writeValueAsString(statePatch)); + } + + docData.put("checkpointId", checkpointId); + docData.put("checkpointShardCount", checkpointShardCount); + docData.put("segmentPath", segmentPath); + tx.set(snapshotDoc(prefix, finalId), docData); + + // 6. Advance pointer (never backward). + advancePointer( + tx, + currentPointer, + prefix, + result.getSessionId(), + finalId, + result.getCreatedAt(), + result.getUpdatedAt(), + checkpointId, + checkpointShardCount, + segmentPath); + + return finalId; + }) + .get(); + } catch (ExecutionException e) { + Throwable cause = e.getCause(); + if (cause instanceof GenkitException) { + throw (GenkitException) cause; + } + throw new GenkitException("Failed to save snapshot: " + cause.getMessage(), cause); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new GenkitException("Interrupted while saving snapshot", e); + } + } + + /** + * Advances the session pointer to the new leaf unless the stored pointer is already newer. The + * pointer document must have been read earlier in the transaction (Firestore requires all reads + * before all writes) and is passed in as {@code current}. + */ + private void advancePointer( + Transaction tx, + DocumentSnapshot current, + String prefix, + String sessionId, + String snapshotId, + String createdAt, + String updatedAt, + String checkpointId, + int checkpointShardCount, + List segmentPath) { + DocumentReference ref = pointerDoc(prefix, sessionId); + if (current.exists()) { + String currentLeaf = current.getString("currentSnapshotId"); + String currentCreatedAt = current.getString("currentCreatedAt"); + // Don't move backward: only refresh when rewriting the same leaf, or when the new snapshot is + // at least as new as the stored leaf (by createdAt string, RFC-3339 lexically sortable). + boolean sameLeaf = snapshotId.equals(currentLeaf); + boolean newer = + createdAt == null + || currentCreatedAt == null + || createdAt.compareTo(currentCreatedAt) >= 0; + if (!sameLeaf && !newer) { + return; + } + } + Map pointerData = new HashMap<>(); + pointerData.put("currentSnapshotId", snapshotId); + pointerData.put("checkpointId", checkpointId); + pointerData.put("checkpointShardCount", checkpointShardCount); + pointerData.put("segmentPath", segmentPath); + if (createdAt != null) { + pointerData.put("currentCreatedAt", createdAt); + } + if (updatedAt != null) { + pointerData.put("updatedAt", updatedAt); + } + tx.set(ref, pointerData); + } + + // ────────────────────────────────────────────────────────────────────────── + // SnapshotReader + // ────────────────────────────────────────────────────────────────────────── + + /** + * {@inheritDoc} + * + *

By {@code snapshotId}: loads the snapshot doc and reconstructs its full state from the + * nearest checkpoint's shards plus the ordered {@code segmentPath} diffs. By {@code sessionId}: + * reads the session pointer, then loads and reconstructs the pointed snapshot. + */ + @Override + public SessionSnapshot getSnapshot(GetSnapshotOptions opts) { + if (opts == null) { + return null; + } + SessionStoreOptions storeOpts = SessionStoreOptions.empty(); + String prefix = prefix(storeOpts); + try { + if (opts.getSnapshotId() != null) { + SnapshotSharding.validateId(opts.getSnapshotId()); + return db.runTransaction( + (Transaction tx) -> { + DocumentSnapshot doc = tx.get(snapshotDoc(prefix, opts.getSnapshotId())).get(); + if (!doc.exists()) { + return null; + } + return readSnapshot(tx, prefix, doc); + }) + .get(); + } + if (opts.getSessionId() != null) { + SnapshotSharding.validateId(opts.getSessionId()); + return db.runTransaction( + (Transaction tx) -> { + DocumentSnapshot pointer = tx.get(pointerDoc(prefix, opts.getSessionId())).get(); + if (!pointer.exists()) { + return null; + } + String leafId = pointer.getString("currentSnapshotId"); + if (leafId == null) { + return null; + } + DocumentSnapshot doc = tx.get(snapshotDoc(prefix, leafId)).get(); + if (!doc.exists()) { + return null; + } + return readSnapshot(tx, prefix, doc); + }) + .get(); + } + return null; + } catch (ExecutionException e) { + Throwable cause = e.getCause(); + if (cause instanceof GenkitException) { + throw (GenkitException) cause; + } + throw new GenkitException("Failed to get snapshot: " + cause.getMessage(), cause); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new GenkitException("Interrupted while getting snapshot", e); + } + } + + // ────────────────────────────────────────────────────────────────────────── + // SnapshotSubscriber + // ────────────────────────────────────────────────────────────────────────── + + /** + * {@inheritDoc} + * + *

Registers a Firestore realtime listener on the snapshot document. On every event (including + * the initial snapshot when the document exists) the snapshot is reconstructed off-band (via a + * read transaction) and passed to {@code cb}. The returned {@link AutoCloseable} removes the + * listener. + */ + @Override + public AutoCloseable onSnapshotStateChange( + String snapshotId, Consumer> cb, SessionStoreOptions options) { + SnapshotSharding.validateId(snapshotId); + String prefix = prefix(options); + ListenerRegistration registration = + snapshotDoc(prefix, snapshotId) + .addSnapshotListener( + (doc, error) -> { + if (error != null) { + logger.warn( + "Snapshot listener error for {}: {}", snapshotId, error.getMessage()); + return; + } + if (doc == null || !doc.exists()) { + return; + } + try { + SessionSnapshot snap = + getSnapshot(GetSnapshotOptions.builder().snapshotId(snapshotId).build()); + if (snap != null) { + cb.accept(snap); + } + } catch (Exception e) { + logger.warn( + "Failed to reconstruct snapshot {} in listener: {}", + snapshotId, + e.getMessage()); + } + }); + return registration::remove; + } + + // ────────────────────────────────────────────────────────────────────────── + // Snapshot (de)serialization to/from Firestore documents + // ────────────────────────────────────────────────────────────────────────── + + /** Builds the non-state base document data for a snapshot. */ + private Map baseDocData(SessionSnapshot snap) throws Exception { + Map data = new HashMap<>(); + data.put("snapshotId", snap.getSnapshotId()); + data.put("sessionId", snap.getSessionId()); + if (snap.getParentId() != null) { + data.put("parentId", snap.getParentId()); + } + if (snap.getCreatedAt() != null) { + data.put("createdAt", snap.getCreatedAt()); + } + if (snap.getUpdatedAt() != null) { + data.put("updatedAt", snap.getUpdatedAt()); + } + if (snap.getHeartbeatAt() != null) { + data.put("heartbeatAt", snap.getHeartbeatAt()); + } + if (snap.getStatus() != null) { + data.put("status", snap.getStatus().getValue()); + } + if (snap.getFinishReason() != null) { + data.put("finishReason", snap.getFinishReason().getValue()); + } + if (snap.getError() != null) { + // Stored as opaque JSON string (Firestore nested-array restriction safety). + data.put("error", MAPPER.writeValueAsString(snap.getError())); + } + return data; + } + + /** Holds the reconstructed parent state and its checkpoint lineage. */ + private static final class ParentInfo { + JsonNode state; + String checkpointId; + int checkpointShardCount; + List segmentPath; + } + + /** Loads a parent snapshot's checkpoint lineage and reconstructs its state. */ + @SuppressWarnings("unchecked") + private ParentInfo loadParentInfo(Transaction tx, String prefix, DocumentSnapshot doc) + throws Exception { + ParentInfo info = new ParentInfo(); + info.checkpointId = doc.getString("checkpointId"); + Long shardCount = doc.getLong("checkpointShardCount"); + info.checkpointShardCount = shardCount != null ? shardCount.intValue() : 0; + Object segObj = doc.get("segmentPath"); + info.segmentPath = + segObj instanceof List ? new ArrayList<>((List) segObj) : new ArrayList<>(); + info.state = + reconstructFullState( + tx, prefix, info.checkpointId, info.checkpointShardCount, info.segmentPath); + return info; + } + + /** Reads and fully reconstructs a snapshot from a Firestore document. */ + @SuppressWarnings("unchecked") + private SessionSnapshot readSnapshot(Transaction tx, String prefix, DocumentSnapshot doc) + throws Exception { + String checkpointId = doc.getString("checkpointId"); + Long shardCount = doc.getLong("checkpointShardCount"); + int checkpointShardCount = shardCount != null ? shardCount.intValue() : 0; + Object segObj = doc.get("segmentPath"); + List segmentPath = + segObj instanceof List ? new ArrayList<>((List) segObj) : new ArrayList<>(); + + JsonNode state = + reconstructFullState(tx, prefix, checkpointId, checkpointShardCount, segmentPath); + return docToSnapshot(doc, state); + } + + /** + * Reconstructs full state: loads the checkpoint shards (concatenate, parse) then applies the + * {@code segmentPath} diffs in order. + */ + private JsonNode reconstructFullState( + Transaction tx, + String prefix, + String checkpointId, + int checkpointShardCount, + List segmentPath) + throws Exception { + if (checkpointId == null) { + return NullNode.getInstance(); + } + // Load shards. + List shardContents = new ArrayList<>(); + for (int i = 0; i < checkpointShardCount; i++) { + DocumentSnapshot shard = tx.get(shardDoc(prefix, checkpointId, i)).get(); + shardContents.add(shard.exists() ? (String) shard.get("data") : ""); + } + String checkpointJson = SnapshotSharding.reassembleShards(shardContents); + + // Load diffs along the segment path (each is an opaque JSON-string patch). + List diffs = new ArrayList<>(); + for (String diffId : segmentPath) { + DocumentSnapshot diffDoc = tx.get(snapshotDoc(prefix, diffId)).get(); + if (diffDoc.exists()) { + String patch = diffDoc.getString("statePatch"); + if (patch != null) { + diffs.add(patch); + } + } + } + return SnapshotSharding.reconstructState(checkpointJson, diffs); + } + + /** Builds a {@link SessionSnapshot} from the metadata document and reconstructed state. */ + @SuppressWarnings("unchecked") + private SessionSnapshot docToSnapshot(DocumentSnapshot doc, JsonNode state) throws Exception { + SessionSnapshot.Builder builder = SessionSnapshot.builder(); + builder.snapshotId(doc.getString("snapshotId")); + builder.sessionId(doc.getString("sessionId")); + builder.parentId(doc.getString("parentId")); + builder.createdAt(doc.getString("createdAt")); + builder.updatedAt(doc.getString("updatedAt")); + builder.heartbeatAt(doc.getString("heartbeatAt")); + String status = doc.getString("status"); + if (status != null) { + builder.status(SnapshotStatus.fromValueOrCompleted(status)); + } + String finishReason = doc.getString("finishReason"); + if (finishReason != null) { + builder.finishReason(com.google.genkit.ai.agent.AgentFinishReason.fromValue(finishReason)); + } + String errorJson = doc.getString("error"); + if (errorJson != null) { + builder.error(MAPPER.readValue(errorJson, RuntimeError.class)); + } + if (state != null && !state.isNull()) { + builder.state((SessionState) MAPPER.treeToValue(state, SessionState.class)); + } + return builder.build(); + } + + /** Serializes session state to a JSON node (null state → JSON null). */ + private JsonNode stateToJson(SessionState state) { + if (state == null) { + return NullNode.getInstance(); + } + return MAPPER.valueToTree(state); + } +} diff --git a/plugins/firebase/src/main/java/com/google/genkit/plugins/firebase/session/FirestoreSessionStoreOptions.java b/plugins/firebase/src/main/java/com/google/genkit/plugins/firebase/session/FirestoreSessionStoreOptions.java new file mode 100644 index 000000000..f77daac05 --- /dev/null +++ b/plugins/firebase/src/main/java/com/google/genkit/plugins/firebase/session/FirestoreSessionStoreOptions.java @@ -0,0 +1,192 @@ +/* + * 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. + * + * SPDX-License-Identifier: Apache-2.0 + */ + +package com.google.genkit.plugins.firebase.session; + +import com.google.genkit.ai.agent.SessionStoreOptions; +import java.util.function.Function; + +/** + * Configuration for {@link FirestoreSessionStore}. + * + *

Mirrors the upstream Go/JS Firestore session store options. The store derives three Firestore + * collection roots from {@link #getCollection()}: + * + *

    + *
  • {@code } — one document per snapshot (checkpoint or diff metadata). + *
  • {@code -shards} — sharded checkpoint state JSON. + *
  • {@code -pointers} — per-session pointer to the current leaf snapshot. + *
+ * + *

All three are namespaced under a per-tenant prefix derived from {@link + * #getSnapshotPathPrefix()} (default {@code "global"}). + */ +public final class FirestoreSessionStoreOptions { + + /** Default top-level collection name. */ + public static final String DEFAULT_COLLECTION = "genkit-sessions"; + + /** Default number of turns between full checkpoints. */ + public static final int DEFAULT_CHECKPOINT_INTERVAL = 25; + + /** Default shard size in bytes for checkpoint state (512 KiB). */ + public static final int DEFAULT_SHARD_SIZE = 512 * 1024; + + private final String collection; + private final int checkpointInterval; + private final int shardSize; + private final Function snapshotPathPrefix; + + private FirestoreSessionStoreOptions(Builder builder) { + this.collection = builder.collection; + this.checkpointInterval = builder.checkpointInterval; + this.shardSize = builder.shardSize; + this.snapshotPathPrefix = builder.snapshotPathPrefix; + } + + /** + * Returns the top-level collection name (default {@value #DEFAULT_COLLECTION}). + * + * @return the collection name + */ + public String getCollection() { + return collection; + } + + /** + * Returns the number of turns between full checkpoints (default {@value + * #DEFAULT_CHECKPOINT_INTERVAL}). + * + * @return the checkpoint interval + */ + public int getCheckpointInterval() { + return checkpointInterval; + } + + /** + * Returns the shard size in bytes for checkpoint state (default {@value #DEFAULT_SHARD_SIZE}). + * + * @return the shard size in bytes + */ + public int getShardSize() { + return shardSize; + } + + /** + * Returns the function that derives the per-tenant path prefix from the per-request store options + * (default {@code o -> "global"}). + * + * @return the prefix function + */ + public Function getSnapshotPathPrefix() { + return snapshotPathPrefix; + } + + /** + * Returns default options. + * + * @return a {@code FirestoreSessionStoreOptions} with all defaults + */ + public static FirestoreSessionStoreOptions defaults() { + return builder().build(); + } + + /** + * Creates a new builder. + * + * @return a new builder + */ + public static Builder builder() { + return new Builder(); + } + + /** Builder for {@link FirestoreSessionStoreOptions}. */ + public static final class Builder { + private String collection = DEFAULT_COLLECTION; + private int checkpointInterval = DEFAULT_CHECKPOINT_INTERVAL; + private int shardSize = DEFAULT_SHARD_SIZE; + private Function snapshotPathPrefix = o -> "global"; + + private Builder() {} + + /** + * Sets the top-level collection name. + * + * @param collection the collection name + * @return this builder + */ + public Builder collection(String collection) { + this.collection = collection; + return this; + } + + /** + * Sets the number of turns between full checkpoints. + * + * @param checkpointInterval the checkpoint interval (must be {@code >= 1}) + * @return this builder + */ + public Builder checkpointInterval(int checkpointInterval) { + this.checkpointInterval = checkpointInterval; + return this; + } + + /** + * Sets the shard size in bytes for checkpoint state. + * + * @param shardSize the shard size in bytes (must be {@code >= 1}) + * @return this builder + */ + public Builder shardSize(int shardSize) { + this.shardSize = shardSize; + return this; + } + + /** + * Sets the function that derives the per-tenant path prefix from the per-request store options. + * + * @param snapshotPathPrefix the prefix function + * @return this builder + */ + public Builder snapshotPathPrefix(Function snapshotPathPrefix) { + this.snapshotPathPrefix = snapshotPathPrefix; + return this; + } + + /** + * Builds a new {@code FirestoreSessionStoreOptions}. + * + * @return a new options instance + */ + public FirestoreSessionStoreOptions build() { + if (collection == null || collection.isBlank()) { + throw new IllegalArgumentException("collection must be non-empty"); + } + if (checkpointInterval < 1) { + throw new IllegalArgumentException("checkpointInterval must be >= 1"); + } + if (shardSize < 1) { + throw new IllegalArgumentException("shardSize must be >= 1"); + } + if (snapshotPathPrefix == null) { + throw new IllegalArgumentException("snapshotPathPrefix must be non-null"); + } + return new FirestoreSessionStoreOptions(this); + } + } +} diff --git a/plugins/firebase/src/main/java/com/google/genkit/plugins/firebase/session/package-info.java b/plugins/firebase/src/main/java/com/google/genkit/plugins/firebase/session/package-info.java new file mode 100644 index 000000000..39fc878ad --- /dev/null +++ b/plugins/firebase/src/main/java/com/google/genkit/plugins/firebase/session/package-info.java @@ -0,0 +1,29 @@ +/* + * 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. + * + * SPDX-License-Identifier: Apache-2.0 + */ + +/** + * Firestore-backed session store for Genkit agents. + * + *

{@link com.google.genkit.plugins.firebase.session.FirestoreSessionStore} implements {@link + * com.google.genkit.ai.agent.SessionStore} and {@link + * com.google.genkit.ai.agent.SnapshotSubscriber}, persisting session snapshots to Cloud Firestore + * using a sharded checkpoint + diff + pointer layout (mirroring the upstream Go/JS Firestore + * session stores). Configure it with {@link + * com.google.genkit.plugins.firebase.session.FirestoreSessionStoreOptions}. + */ +package com.google.genkit.plugins.firebase.session; diff --git a/plugins/firebase/src/test/java/com/google/genkit/plugins/firebase/session/FirestoreSessionStoreTest.java b/plugins/firebase/src/test/java/com/google/genkit/plugins/firebase/session/FirestoreSessionStoreTest.java new file mode 100644 index 000000000..fed116c12 --- /dev/null +++ b/plugins/firebase/src/test/java/com/google/genkit/plugins/firebase/session/FirestoreSessionStoreTest.java @@ -0,0 +1,292 @@ +/* + * 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. + * + * SPDX-License-Identifier: Apache-2.0 + */ + +package com.google.genkit.plugins.firebase.session; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertNull; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import com.google.cloud.firestore.Firestore; +import com.google.cloud.firestore.FirestoreOptions; +import com.google.genkit.ai.agent.GetSnapshotOptions; +import com.google.genkit.ai.agent.SessionSnapshot; +import com.google.genkit.ai.agent.SessionState; +import com.google.genkit.ai.agent.SessionStoreOptions; +import com.google.genkit.ai.agent.SnapshotStatus; +import com.google.genkit.core.GenkitException; +import java.util.HashMap; +import java.util.Map; +import java.util.UUID; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicReference; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.Assumptions; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; + +/** + * Tests for {@link FirestoreSessionStore}. + * + *

    + *
  • Options unit tests — validate {@link FirestoreSessionStoreOptions}. These always + * run. + *
  • Emulator-gated integration tests — gated on {@code FIRESTORE_EMULATOR_HOST}; skipped + * (via {@link Assumptions#assumeTrue}) when no emulator is configured. + *
+ * + *

The pure sharding / checkpoint / reconstruction helpers now live in {@code + * com.google.genkit.ai.agent.internal.SnapshotSharding} and are unit-tested by {@code + * SnapshotShardingTest} in the {@code ai} module. + */ +class FirestoreSessionStoreTest { + + // ────────────────────────────────────────────────────────────────────────── + // Options unit tests (no Firestore) + // ────────────────────────────────────────────────────────────────────────── + + @Test + void optionsDefaults() { + FirestoreSessionStoreOptions opts = FirestoreSessionStoreOptions.builder().build(); + assertEquals("genkit-sessions", opts.getCollection()); + assertEquals(25, opts.getCheckpointInterval()); + assertEquals(512 * 1024, opts.getShardSize()); + assertEquals("global", opts.getSnapshotPathPrefix().apply(SessionStoreOptions.empty())); + } + + @Test + void optionsCustomBuilder() { + FirestoreSessionStoreOptions opts = + FirestoreSessionStoreOptions.builder() + .collection("my-sessions") + .checkpointInterval(10) + .shardSize(1024) + .snapshotPathPrefix(o -> "tenant-1") + .build(); + assertEquals("my-sessions", opts.getCollection()); + assertEquals(10, opts.getCheckpointInterval()); + assertEquals(1024, opts.getShardSize()); + assertEquals("tenant-1", opts.getSnapshotPathPrefix().apply(SessionStoreOptions.empty())); + } + + // ────────────────────────────────────────────────────────────────────────── + // Emulator-gated integration tests + // ────────────────────────────────────────────────────────────────────────── + + private Firestore firestore; + private FirestoreSessionStore> store; + private String collection; + + @BeforeEach + void setUp() { + String emulatorHost = System.getenv("FIRESTORE_EMULATOR_HOST"); + if (emulatorHost == null || emulatorHost.isEmpty()) { + return; // emulator-gated tests will be skipped via assumeTrue + } + firestore = + FirestoreOptions.getDefaultInstance().toBuilder() + .setProjectId("genkit-test") + .setEmulatorHost(emulatorHost) + .build() + .getService(); + // Unique collection per run to isolate documents. + collection = "genkit-sessions-test-" + UUID.randomUUID().toString().substring(0, 8); + store = + new FirestoreSessionStore<>( + firestore, + FirestoreSessionStoreOptions.builder() + .collection(collection) + .checkpointInterval(3) + .build()); + } + + @AfterEach + void tearDown() throws Exception { + if (firestore != null) { + firestore.close(); + } + } + + private static SessionSnapshot> snapshotWithState( + String sessionId, String parentId, Map custom) { + SessionState> state = + SessionState.>builder().sessionId(sessionId).custom(custom).build(); + return SessionSnapshot.>builder() + .sessionId(sessionId) + .parentId(parentId) + .status(SnapshotStatus.COMPLETED) + .state(state) + .build(); + } + + @Test + void emulatorSaveGetRoundTrip() { + Assumptions.assumeTrue(System.getenv("FIRESTORE_EMULATOR_HOST") != null); + + String sessionId = "s-" + UUID.randomUUID(); + Map custom = new HashMap<>(); + custom.put("count", 1); + + String id = + store.saveSnapshot( + null, + existing -> snapshotWithState(sessionId, null, custom), + SessionStoreOptions.empty()); + assertNotNull(id); + + SessionSnapshot> got = + store.getSnapshot(GetSnapshotOptions.builder().snapshotId(id).build()); + assertNotNull(got); + assertEquals(sessionId, got.getSessionId()); + assertEquals(SnapshotStatus.COMPLETED, got.getStatus()); + assertNotNull(got.getState()); + assertEquals(1, ((Number) got.getState().getCustom().get("count")).intValue()); + } + + @Test + void emulatorRejectsEmptySessionId() { + Assumptions.assumeTrue(System.getenv("FIRESTORE_EMULATOR_HOST") != null); + GenkitException ex = + assertThrows( + GenkitException.class, + () -> + store.saveSnapshot( + null, + existing -> snapshotWithState("", null, new HashMap<>()), + SessionStoreOptions.empty())); + assertEquals("INVALID_ARGUMENT", ex.getErrorCode()); + } + + @Test + void emulatorMutatorNullIsNoOp() { + Assumptions.assumeTrue(System.getenv("FIRESTORE_EMULATOR_HOST") != null); + String result = store.saveSnapshot(null, existing -> null, SessionStoreOptions.empty()); + assertNull(result); + } + + @Test + void emulatorDiffChainReconstructs() { + Assumptions.assumeTrue(System.getenv("FIRESTORE_EMULATOR_HOST") != null); + + String sessionId = "s-" + UUID.randomUUID(); + + // Turn 1 (checkpoint - root) + Map c1 = new HashMap<>(); + c1.put("count", 1); + String id1 = + store.saveSnapshot( + null, e -> snapshotWithState(sessionId, null, c1), SessionStoreOptions.empty()); + + // Turn 2 (diff) + Map c2 = new HashMap<>(); + c2.put("count", 2); + String id2 = + store.saveSnapshot( + null, e -> snapshotWithState(sessionId, id1, c2), SessionStoreOptions.empty()); + + // Turn 3 (diff) + Map c3 = new HashMap<>(); + c3.put("count", 3); + c3.put("extra", "z"); + String id3 = + store.saveSnapshot( + null, e -> snapshotWithState(sessionId, id2, c3), SessionStoreOptions.empty()); + + SessionSnapshot> got = + store.getSnapshot(GetSnapshotOptions.builder().snapshotId(id3).build()); + assertNotNull(got); + assertEquals(3, ((Number) got.getState().getCustom().get("count")).intValue()); + assertEquals("z", got.getState().getCustom().get("extra")); + } + + @Test + void emulatorGetLatestViaPointer() { + Assumptions.assumeTrue(System.getenv("FIRESTORE_EMULATOR_HOST") != null); + + String sessionId = "s-" + UUID.randomUUID(); + Map c1 = new HashMap<>(); + c1.put("count", 1); + String id1 = + store.saveSnapshot( + null, e -> snapshotWithState(sessionId, null, c1), SessionStoreOptions.empty()); + Map c2 = new HashMap<>(); + c2.put("count", 2); + String id2 = + store.saveSnapshot( + null, e -> snapshotWithState(sessionId, id1, c2), SessionStoreOptions.empty()); + + SessionSnapshot> latest = + store.getSnapshot(GetSnapshotOptions.builder().sessionId(sessionId).build()); + assertNotNull(latest); + assertEquals(id2, latest.getSnapshotId()); + assertEquals(2, ((Number) latest.getState().getCustom().get("count")).intValue()); + } + + @Test + void emulatorCheckpointEveryInterval() { + Assumptions.assumeTrue(System.getenv("FIRESTORE_EMULATOR_HOST") != null); + + String sessionId = "s-" + UUID.randomUUID(); + String parent = null; + String lastId = null; + // checkpointInterval is 3; save 5 turns and confirm reconstruction stays correct. + for (int i = 1; i <= 5; i++) { + Map c = new HashMap<>(); + c.put("count", i); + final String p = parent; + lastId = + store.saveSnapshot( + null, e -> snapshotWithState(sessionId, p, c), SessionStoreOptions.empty()); + parent = lastId; + } + SessionSnapshot> got = + store.getSnapshot(GetSnapshotOptions.builder().snapshotId(lastId).build()); + assertNotNull(got); + assertEquals(5, ((Number) got.getState().getCustom().get("count")).intValue()); + } + + @Test + void emulatorOnSnapshotStateChangeFires() throws Exception { + Assumptions.assumeTrue(System.getenv("FIRESTORE_EMULATOR_HOST") != null); + + String sessionId = "s-" + UUID.randomUUID(); + Map c1 = new HashMap<>(); + c1.put("count", 1); + String id = + store.saveSnapshot( + null, e -> snapshotWithState(sessionId, null, c1), SessionStoreOptions.empty()); + + CountDownLatch latch = new CountDownLatch(1); + AtomicReference> seen = new AtomicReference<>(); + try (AutoCloseable sub = + store.onSnapshotStateChange( + id, + snap -> { + seen.set(snap); + latch.countDown(); + }, + SessionStoreOptions.empty())) { + assertTrue(latch.await(10, TimeUnit.SECONDS), "callback should fire at subscription time"); + assertNotNull(seen.get()); + assertEquals(id, seen.get().getSnapshotId()); + } + } +} diff --git a/plugins/jetty/src/main/java/com/google/genkit/plugins/jetty/JettyPlugin.java b/plugins/jetty/src/main/java/com/google/genkit/plugins/jetty/JettyPlugin.java index 178d5001c..a57c63377 100644 --- a/plugins/jetty/src/main/java/com/google/genkit/plugins/jetty/JettyPlugin.java +++ b/plugins/jetty/src/main/java/com/google/genkit/plugins/jetty/JettyPlugin.java @@ -25,6 +25,7 @@ import java.nio.ByteBuffer; import java.nio.charset.StandardCharsets; import java.util.*; +import java.util.function.Consumer; import org.eclipse.jetty.server.Handler; import org.eclipse.jetty.server.Request; import org.eclipse.jetty.server.Response; @@ -167,6 +168,9 @@ private void startServer() throws Exception { // Add flow endpoints addFlowHandlers(handlers); + // Add agent endpoints (bidi actions + companions) + addAgentHandlers(handlers); + // Add health endpoint ContextHandler healthHandler = new ContextHandler("/health"); healthHandler.setHandler(new HealthHandler()); @@ -228,6 +232,62 @@ private void addFlowHandlers(ContextHandlerCollection handlers) { } } + /** + * Adds HTTP handlers for all registered agents (bidi actions of type {@link ActionType#AGENT}). + * + *

For each agent named {@code } this mounts: + * + *

    + *
  • {@code POST /} -> one turn per request (non-streaming or SSE). + *
  • {@code POST //getSnapshot} -> companion {@code /agent-snapshot/} if present. + *
  • {@code POST //abort} -> companion {@code /agent-abort/} if present. + *
+ * + *

Companions may be absent for client-managed agents; only existing ones are mounted. + */ + private void addAgentHandlers(ContextHandlerCollection handlers) { + List> agents = registry.listActions(ActionType.AGENT); + + for (Action action : agents) { + if (!(action instanceof BidiAction)) { + logger.warn("Agent action {} is not a BidiAction; skipping", action.getName()); + continue; + } + String name = action.getName(); + String path = "/" + name; + + // Main one-turn-per-request endpoint. + ContextHandler agentHandler = new ContextHandler(path); + agentHandler.setAllowNullPathInContext(true); + agentHandler.setHandler(new AgentHandler((BidiAction) action)); + handlers.addHandler(agentHandler); + logger.info("Registered agent endpoint: {}", path); + + // Companion getSnapshot endpoint (if a companion action is registered). + Action snapshotAction = + registry.lookupAction(ActionType.AGENT_SNAPSHOT.keyFromName(name)); + if (snapshotAction != null) { + String snapshotPath = path + "/getSnapshot"; + ContextHandler snapshotHandler = new ContextHandler(snapshotPath); + snapshotHandler.setAllowNullPathInContext(true); + snapshotHandler.setHandler(new CompanionHandler(snapshotAction)); + handlers.addHandler(snapshotHandler); + logger.info("Registered agent companion endpoint: {}", snapshotPath); + } + + // Companion abort endpoint (if a companion action is registered). + Action abortAction = registry.lookupAction(ActionType.AGENT_ABORT.keyFromName(name)); + if (abortAction != null) { + String abortPath = path + "/abort"; + ContextHandler abortHandler = new ContextHandler(abortPath); + abortHandler.setAllowNullPathInContext(true); + abortHandler.setHandler(new CompanionHandler(abortAction)); + handlers.addHandler(abortHandler); + logger.info("Registered agent companion endpoint: {}", abortPath); + } + } + } + /** Handler for health check endpoint. */ private class HealthHandler extends Handler.Abstract { @Override @@ -317,4 +377,378 @@ public boolean handle(Request request, Response response, Callback callback) thr } } } + + /** + * Handler for agent endpoints. Serves one turn of a bidi agent per HTTP request. + * + *

Request body envelope is {@code {"data": , "init": }}. Streaming is + * selected when the {@code Accept} header contains {@code text/event-stream} or the query string + * contains {@code stream=true}; otherwise the final result is returned as a single JSON object. + */ + private class AgentHandler extends Handler.Abstract { + private final BidiAction action; + + @SuppressWarnings("unchecked") + AgentHandler(BidiAction action) { + this.action = (BidiAction) action; + } + + @Override + public boolean handle(Request request, Response response, Callback callback) throws Exception { + // Only accept POST requests. + if (!"POST".equals(request.getMethod())) { + writeMethodNotAllowed(response, callback); + return true; + } + + // Read and parse the request body. + JsonNode body; + try { + ByteArrayOutputStream baos = new ByteArrayOutputStream(); + Request.asInputStream(request).transferTo(baos); + String raw = baos.toString(StandardCharsets.UTF_8); + body = (raw == null || raw.isEmpty()) ? null : objectMapper.readTree(raw); + } catch (Exception e) { + logger.error("Error reading agent request body", e); + writeError(response, callback, 400, "INVALID_ARGUMENT", e); + return true; + } + + JsonNode data = body != null ? body.get("data") : null; + JsonNode init = body != null ? body.get("init") : null; + Map userContext = mergeHeadersIntoContext(parseContext(body), request); + + boolean streaming = isStreamingRequested(request); + if (streaming) { + handleStreaming(response, callback, data, init, userContext); + } else { + handleUnary(response, callback, data, init, userContext); + } + return true; + } + + /** Non-streaming: run the turn and write {@code {"result": }} or an error. */ + private void handleUnary( + Response response, + Callback callback, + JsonNode data, + JsonNode init, + Map userContext) { + try { + BufferedInputSource inputs = new BufferedInputSource<>(); + if (data != null && !data.isNull()) { + inputs.offer(data); + } + inputs.end(); + + ActionContext context = + ActionContext.builder().registry(registry).context(userContext).build(); + JsonNode result = action.runBidiJson(context, init, inputs, null); + + Map envelope = new HashMap<>(); + envelope.put("result", result); + + response.setStatus(200); + response.getHeaders().put("Content-Type", "application/json"); + byte[] bytes = objectMapper.writeValueAsBytes(envelope); + response.write(true, ByteBuffer.wrap(bytes), callback); + } catch (Exception e) { + logger.error("Error handling agent request", e); + writeError(response, callback, statusFromError(e), errorCodeFromError(e), e); + } + } + + /** + * Streaming: respond with {@code text/event-stream}, emitting one {@code data: {"message": + * }} frame per chunk, then a final {@code data: {"result": }} frame; errors emit + * a {@code data: {"error": {...}}} frame. + */ + private void handleStreaming( + Response response, + Callback callback, + JsonNode data, + JsonNode init, + Map userContext) { + response.setStatus(200); + response.getHeaders().put("Content-Type", "text/event-stream"); + response.getHeaders().put("Cache-Control", "no-cache"); + response.getHeaders().put("Connection", "keep-alive"); + response.getHeaders().put("X-Genkit-Stream-Id", UUID.randomUUID().toString()); + + Consumer streamCallback = + chunk -> { + Map frame = new HashMap<>(); + frame.put("message", chunk); + writeSseFrame(response, frame); + }; + + try { + BufferedInputSource inputs = new BufferedInputSource<>(); + if (data != null && !data.isNull()) { + inputs.offer(data); + } + inputs.end(); + + ActionContext context = + ActionContext.builder().registry(registry).context(userContext).build(); + JsonNode result = action.runBidiJson(context, init, inputs, streamCallback); + + Map resultFrame = new HashMap<>(); + resultFrame.put("result", result); + writeSseFrameLast(response, callback, resultFrame); + } catch (Exception e) { + logger.error("Error streaming agent request", e); + Map errorFrame = new HashMap<>(); + errorFrame.put("error", errorBody(errorCodeFromError(e), e)); + writeSseFrameLast(response, callback, errorFrame); + } + } + + /** Writes a non-final SSE frame, blocking until the write completes. */ + private void writeSseFrame(Response response, Object payload) { + try { + String data = "data: " + objectMapper.writeValueAsString(payload) + "\n\n"; + blockingWrite(response, false, ByteBuffer.wrap(data.getBytes(StandardCharsets.UTF_8))); + } catch (Exception e) { + throw new RuntimeException("Failed to write SSE frame", e); + } + } + + /** Writes the final SSE frame and completes the supplied request callback. */ + private void writeSseFrameLast(Response response, Callback callback, Object payload) { + try { + String data = "data: " + objectMapper.writeValueAsString(payload) + "\n\n"; + response.write(true, ByteBuffer.wrap(data.getBytes(StandardCharsets.UTF_8)), callback); + } catch (Exception e) { + callback.failed(e); + } + } + + /** Performs a Jetty async write and blocks the current thread until it completes. */ + private void blockingWrite(Response response, boolean last, ByteBuffer buffer) + throws Exception { + Callback.Completable completable = new Callback.Completable(); + response.write(last, buffer, completable); + completable.get(); + } + } + + /** + * Handler for agent companion endpoints (snapshot / abort). Unary POST over a plain {@link + * Action} looked up by key. Request body envelope is {@code {"data": }}; the response is + * {@code {"result": }} or an error. + */ + private class CompanionHandler extends Handler.Abstract { + private final Action action; + + @SuppressWarnings("unchecked") + CompanionHandler(Action action) { + this.action = (Action) action; + } + + @Override + public boolean handle(Request request, Response response, Callback callback) throws Exception { + if (!"POST".equals(request.getMethod())) { + writeMethodNotAllowed(response, callback); + return true; + } + + try { + ByteArrayOutputStream baos = new ByteArrayOutputStream(); + Request.asInputStream(request).transferTo(baos); + String raw = baos.toString(StandardCharsets.UTF_8); + JsonNode body = (raw == null || raw.isEmpty()) ? null : objectMapper.readTree(raw); + JsonNode data = body != null ? body.get("data") : null; + + ActionContext context = + ActionContext.builder() + .registry(registry) + .context(mergeHeadersIntoContext(parseContext(body), request)) + .build(); + JsonNode result = action.runJson(context, data, null); + + Map envelope = new HashMap<>(); + envelope.put("result", result); + + response.setStatus(200); + response.getHeaders().put("Content-Type", "application/json"); + byte[] bytes = objectMapper.writeValueAsBytes(envelope); + response.write(true, ByteBuffer.wrap(bytes), callback); + } catch (Exception e) { + logger.error("Error handling agent companion request", e); + writeError(response, callback, statusFromError(e), errorCodeFromError(e), e); + } + return true; + } + } + + // --------------------------------------------------------------------------- + // Shared helpers for agent / companion handlers + // --------------------------------------------------------------------------- + + /** + * Parses the optional {@code context} object from a request body envelope into a {@code + * Map}. Threaded into the run's ActionContext so tools/flows can read it (e.g. + * {@code {"auth": {"user": "alice"}}}). + * + * @param body the parsed request body (may be null) + * @return the parsed context map, or null if absent/blank + */ + private Map parseContext(JsonNode body) { + if (body == null || !body.has("context") || body.get("context").isNull()) { + return null; + } + JsonNode contextNode = body.get("context"); + if (!contextNode.isObject()) { + return null; + } + return objectMapper.convertValue( + contextNode, new com.fasterxml.jackson.core.type.TypeReference>() {}); + } + + /** + * HTTP framing/transport headers that are excluded from the {@code "headers"} sub-map threaded + * into the run's {@link ActionContext}. These describe the HTTP message itself (body encoding, + * connection lifecycle, routing) rather than application-level data a tool/flow would care about. + * Everything else — including custom headers like {@code Authorization} or {@code X-*} — is + * included; when in doubt we err on the side of including a header rather than filtering it. + */ + private static final Set EXCLUDED_HEADERS = + Set.of("content-type", "content-length", "accept", "accept-encoding", "connection", "host"); + + /** + * Merges incoming HTTP request headers into the request-scoped user context returned by {@link + * #parseContext(JsonNode)}. + * + *

Design: headers are exposed to tools/flows the same way the JSON-body {@code context} + * field already is — via {@link com.google.genkit.core.ActionContext#getContext()} — by nesting + * them under a reserved {@code "headers"} key as a {@code Map} (single-valued, to + * match the shape of {@code RemoteAgentOptions.headers()} on the client). This is additive and + * cannot silently collide with arbitrary body-context keys other than a literal top-level {@code + * "headers"} key. Precedence: if the body-supplied {@code context} already defines a + * {@code "headers"} entry, that body value wins and incoming HTTP headers are dropped for that + * key (body {@code context} always takes precedence over transport-level data); otherwise the + * HTTP headers populate {@code context.get("headers")}. + * + * @param bodyContext the context map parsed from the request body (may be null) + * @param request the incoming Jetty request, whose headers are folded in + * @return a merged context map (never null if headers are present; may be null if both the body + * context and the header set are empty) + */ + private Map mergeHeadersIntoContext( + Map bodyContext, Request request) { + Map headers = parseHeaders(request); + if (headers.isEmpty()) { + return bodyContext; + } + Map merged = bodyContext != null ? new HashMap<>(bodyContext) : new HashMap<>(); + merged.putIfAbsent("headers", headers); + return merged; + } + + /** + * Extracts incoming HTTP request headers (excluding standard framing headers, see {@link + * #EXCLUDED_HEADERS}) into a {@code Map}. If a header name repeats, the last value + * wins (mirroring {@code HttpFields#get}). + * + * @param request the incoming Jetty request + * @return a mutable map of header name to value (never null; may be empty) + */ + private static Map parseHeaders(Request request) { + Map result = new HashMap<>(); + for (org.eclipse.jetty.http.HttpField field : request.getHeaders()) { + String name = field.getName(); + if (name == null || EXCLUDED_HEADERS.contains(name.toLowerCase(Locale.ROOT))) { + continue; + } + result.put(name, field.getValue()); + } + return result; + } + + /** Returns true if the request asks for SSE streaming. */ + private static boolean isStreamingRequested(Request request) { + String accept = request.getHeaders().get("Accept"); + if (accept != null && accept.contains("text/event-stream")) { + return true; + } + String query = request.getHttpURI().getQuery(); + return query != null && query.contains("stream=true"); + } + + /** Writes a 405 Method Not Allowed JSON response. */ + private void writeMethodNotAllowed(Response response, Callback callback) { + response.setStatus(405); + response.getHeaders().put("Content-Type", "application/json"); + String error = "{\"error\":\"Method not allowed\"}"; + response.write(true, ByteBuffer.wrap(error.getBytes(StandardCharsets.UTF_8)), callback); + } + + /** Writes an error response with the given HTTP status and {@code {"error": {...}}} body. */ + private void writeError( + Response response, Callback callback, int httpStatus, String code, Throwable e) { + response.setStatus(httpStatus); + response.getHeaders().put("Content-Type", "application/json"); + Map envelope = new HashMap<>(); + envelope.put("error", errorBody(code, e)); + try { + byte[] bytes = objectMapper.writeValueAsBytes(envelope); + response.write(true, ByteBuffer.wrap(bytes), callback); + } catch (Exception writeError) { + callback.failed(writeError); + } + } + + /** Builds a structured error body: {@code {status, message, details: {stack}}}. */ + private static Map errorBody(String code, Throwable e) { + String message = e.getMessage() != null ? e.getMessage() : "Unknown error"; + java.io.StringWriter sw = new java.io.StringWriter(); + e.printStackTrace(new java.io.PrintWriter(sw)); + Map body = new HashMap<>(); + body.put("status", code); + body.put("message", message); + body.put("details", Map.of("stack", sw.toString())); + return body; + } + + /** Derives a status string (mirroring the error's status field) from a thrown error. */ + private static String errorCodeFromError(Throwable e) { + if (e instanceof GenkitException) { + String c = ((GenkitException) e).getErrorCode(); + if (c != null && !c.isEmpty()) { + return c; + } + } + return "INTERNAL"; + } + + /** Maps a thrown error to an HTTP status code derived from its status field. */ + private static int statusFromError(Throwable e) { + String code = errorCodeFromError(e); + switch (code) { + case "NOT_FOUND": + return 404; + case "INVALID_ARGUMENT": + case "FAILED_PRECONDITION": + case "OUT_OF_RANGE": + return 400; + case "UNAUTHENTICATED": + return 401; + case "PERMISSION_DENIED": + return 403; + case "ALREADY_EXISTS": + case "ABORTED": + return 409; + case "RESOURCE_EXHAUSTED": + return 429; + case "UNIMPLEMENTED": + return 501; + case "UNAVAILABLE": + return 503; + case "DEADLINE_EXCEEDED": + return 504; + default: + return 500; + } + } } diff --git a/plugins/jetty/src/test/java/com/google/genkit/plugins/jetty/AgentAbortHttpTest.java b/plugins/jetty/src/test/java/com/google/genkit/plugins/jetty/AgentAbortHttpTest.java new file mode 100644 index 000000000..a18c38497 --- /dev/null +++ b/plugins/jetty/src/test/java/com/google/genkit/plugins/jetty/AgentAbortHttpTest.java @@ -0,0 +1,244 @@ +/* + * 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. + * + * SPDX-License-Identifier: Apache-2.0 + */ + +package com.google.genkit.plugins.jetty; + +import static org.junit.jupiter.api.Assertions.*; + +import com.fasterxml.jackson.databind.JsonNode; +import com.fasterxml.jackson.databind.ObjectMapper; +import com.google.genkit.ai.Message; +import com.google.genkit.ai.agent.AgentFinishReason; +import com.google.genkit.ai.agent.AgentResult; +import com.google.genkit.ai.agent.CustomAgentConfig; +import com.google.genkit.ai.agent.InMemorySessionStore; +import com.google.genkit.ai.agent.internal.AgentActions; +import com.google.genkit.core.DefaultRegistry; +import java.io.IOException; +import java.net.HttpURLConnection; +import java.net.URI; +import java.net.URL; +import java.net.http.HttpClient; +import java.net.http.HttpRequest; +import java.net.http.HttpResponse; +import java.util.Map; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.TimeUnit; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.Test; + +/** + * Integration test verifying the abort companion endpoint over the HTTP wire using JettyPlugin. + * + *

Cross-referenced with {@code Agent.abort(String)} / {@code AgentActions.buildAbortAction}: the + * abort companion only flips a stored snapshot's status from {@code PENDING} to {@code ABORTED}; it + * never interrupts a running foreground {@code AgentFn} call (there is no cancellation hook + * threaded into a synchronous {@code AgentFn.run}). This test therefore targets what is actually + * implemented: aborting a snapshot that is genuinely in the {@code PENDING} window opened by a + * detached turn (see {@code DetachController}), which synchronously persists the {@code PENDING} + * row before the HTTP response returns and whose finalizer never overwrites an {@code ABORTED} + * status (race-safe by design) — then confirms via {@code getSnapshot} that the status is {@code + * ABORTED} and stays that way even after the background turn body completes. + */ +class AgentAbortHttpTest { + + private static final ObjectMapper MAPPER = new ObjectMapper(); + + private JettyPlugin plugin; + + @AfterEach + void tearDown() throws Exception { + if (plugin != null) { + plugin.stop(); + } + } + + /** + * Starts a JettyPlugin on the given port with a server-managed custom agent named {@code + * abortAgent} whose AgentFn blocks on {@code latch} until released, giving the test a + * deterministic window in which the snapshot is guaranteed to still be {@code PENDING}. + */ + private void startWithBlockingAgent(int port, CountDownLatch releaseLatch) throws Exception { + DefaultRegistry registry = new DefaultRegistry(); + + AgentActions.defineCustomAgent( + registry, + CustomAgentConfig.>builder() + .name("abortAgent") + .store(new InMemorySessionStore<>()) + .build(), + (sess, fnCtx) -> { + // Block the detached background turn body until the test releases it, so the + // PENDING window is deterministic and long enough to reliably call abort(). + releaseLatch.await(10, TimeUnit.SECONDS); + return AgentResult.builder() + .message(Message.model("done")) + .finishReason(AgentFinishReason.STOP) + .build(); + }); + + JettyPluginOptions options = JettyPluginOptions.builder().port(port).build(); + plugin = new JettyPlugin(options); + plugin.init(registry); + + Thread serverThread = + new Thread( + () -> { + try { + plugin.start(); + } catch (Exception e) { + // Server stopped. + } + }); + serverThread.setDaemon(true); + serverThread.start(); + + waitForServer(port); + } + + /** + * Integration test: detach a turn (leaving its snapshot PENDING while the AgentFn blocks) -> POST + * /abort for that snapshot -> assert the abort response reports status ABORTED -> release the + * blocked turn so it finalizes in the background -> poll getSnapshot and assert the status + * remains ABORTED (proving the finalizer's abort-aware guard truly won the race, not just that + * the abort call itself returned ABORTED). + */ + @Test + void testAbortOverHttp() throws Exception { + int port = findAvailablePort(); + CountDownLatch releaseLatch = new CountDownLatch(1); + startWithBlockingAgent(port, releaseLatch); + + HttpClient client = HttpClient.newHttpClient(); + + // Step 1: POST a detach turn. The AgentFn blocks on releaseLatch, so the snapshot remains + // PENDING until we count down the latch below. + String detachBody = + "{\"data\":{\"detach\":true,\"message\":{\"role\":\"user\"," + + "\"content\":[{\"text\":\"go\"}]}},\"init\":{}}"; + + HttpRequest detachRequest = + HttpRequest.newBuilder() + .uri(URI.create("http://localhost:" + port + "/abortAgent")) + .header("Content-Type", "application/json") + .POST(HttpRequest.BodyPublishers.ofString(detachBody)) + .build(); + + HttpResponse detachResponse = + client.send(detachRequest, HttpResponse.BodyHandlers.ofString()); + assertEquals(200, detachResponse.statusCode(), "detach POST body: " + detachResponse.body()); + + JsonNode detachRoot = MAPPER.readTree(detachResponse.body()); + JsonNode result = detachRoot.path("result"); + assertEquals( + "detached", + result.path("finishReason").asText(""), + "expected finishReason=detached: " + detachResponse.body()); + String snapshotId = result.path("snapshotId").asText(""); + assertFalse(snapshotId.isEmpty(), "expected non-empty snapshotId"); + + // Step 2: POST /abort for the still-PENDING snapshot. + String abortBody = "{\"data\":{\"snapshotId\":\"" + snapshotId + "\"}}"; + HttpRequest abortRequest = + HttpRequest.newBuilder() + .uri(URI.create("http://localhost:" + port + "/abortAgent/abort")) + .header("Content-Type", "application/json") + .POST(HttpRequest.BodyPublishers.ofString(abortBody)) + .build(); + + HttpResponse abortResponse = + client.send(abortRequest, HttpResponse.BodyHandlers.ofString()); + assertEquals(200, abortResponse.statusCode(), "abort POST body: " + abortResponse.body()); + + JsonNode abortRoot = MAPPER.readTree(abortResponse.body()); + assertTrue(abortRoot.has("result"), "expected result envelope: " + abortResponse.body()); + JsonNode abortResult = abortRoot.path("result"); + assertEquals(snapshotId, abortResult.path("snapshotId").asText("")); + assertEquals( + "aborted", + abortResult.path("status").asText(""), + "expected abort response status=aborted: " + abortResponse.body()); + + // Step 3: getSnapshot immediately (still PENDING-turned-ABORTED, before the background turn + // body has been released) should already reflect ABORTED. + String getSnapshotBody = "{\"data\":{\"snapshotId\":\"" + snapshotId + "\"}}"; + HttpRequest snapshotRequest = + HttpRequest.newBuilder() + .uri(URI.create("http://localhost:" + port + "/abortAgent/getSnapshot")) + .header("Content-Type", "application/json") + .POST(HttpRequest.BodyPublishers.ofString(getSnapshotBody)) + .build(); + + HttpResponse snapResponse1 = + client.send(snapshotRequest, HttpResponse.BodyHandlers.ofString()); + assertEquals(200, snapResponse1.statusCode(), "getSnapshot body: " + snapResponse1.body()); + JsonNode snapResult1 = MAPPER.readTree(snapResponse1.body()).path("result"); + assertEquals( + "aborted", + snapResult1.path("status").asText(""), + "expected snapshot status=aborted before finalize: " + snapResult1); + + // Step 4: release the blocked background turn so it runs to completion and attempts to + // finalize the (now ABORTED) snapshot to COMPLETED. + releaseLatch.countDown(); + + // Step 5: poll getSnapshot for up to 3s and assert the status NEVER reverts from "aborted" + // (proving DetachController's finalizer really never overwrites an ABORTED row). + long deadline = System.currentTimeMillis() + 3000; + String lastStatus = "aborted"; + while (System.currentTimeMillis() < deadline) { + HttpResponse snapResponse = + client.send(snapshotRequest, HttpResponse.BodyHandlers.ofString()); + assertEquals(200, snapResponse.statusCode(), "getSnapshot body: " + snapResponse.body()); + JsonNode snapResult = MAPPER.readTree(snapResponse.body()).path("result"); + lastStatus = snapResult.path("status").asText(""); + assertEquals( + "aborted", + lastStatus, + "snapshot status must remain aborted (finalizer must not overwrite it): " + snapResult); + Thread.sleep(50); + } + assertEquals("aborted", lastStatus); + } + + private static int findAvailablePort() throws IOException { + try (java.net.ServerSocket socket = new java.net.ServerSocket(0)) { + return socket.getLocalPort(); + } + } + + private static void waitForServer(int port) throws Exception { + int maxRetries = 50; + for (int i = 0; i < maxRetries; i++) { + try { + HttpURLConnection conn = + (HttpURLConnection) new URL("http://localhost:" + port + "/health").openConnection(); + conn.setRequestMethod("GET"); + conn.setConnectTimeout(200); + conn.setReadTimeout(200); + if (conn.getResponseCode() == 200) { + return; + } + } catch (IOException e) { + // Server not ready yet. + } + Thread.sleep(100); + } + fail("Server did not start within timeout"); + } +} diff --git a/plugins/jetty/src/test/java/com/google/genkit/plugins/jetty/AgentDetachHttpTest.java b/plugins/jetty/src/test/java/com/google/genkit/plugins/jetty/AgentDetachHttpTest.java new file mode 100644 index 000000000..dcbf04b3e --- /dev/null +++ b/plugins/jetty/src/test/java/com/google/genkit/plugins/jetty/AgentDetachHttpTest.java @@ -0,0 +1,221 @@ +/* + * 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. + * + * SPDX-License-Identifier: Apache-2.0 + */ + +package com.google.genkit.plugins.jetty; + +import static org.junit.jupiter.api.Assertions.*; + +import com.fasterxml.jackson.databind.JsonNode; +import com.fasterxml.jackson.databind.ObjectMapper; +import com.google.genkit.ai.Message; +import com.google.genkit.ai.agent.AgentFinishReason; +import com.google.genkit.ai.agent.AgentResult; +import com.google.genkit.ai.agent.CustomAgentConfig; +import com.google.genkit.ai.agent.InMemorySessionStore; +import com.google.genkit.ai.agent.internal.AgentActions; +import com.google.genkit.core.DefaultRegistry; +import java.io.IOException; +import java.net.HttpURLConnection; +import java.net.URI; +import java.net.URL; +import java.net.http.HttpClient; +import java.net.http.HttpRequest; +import java.net.http.HttpResponse; +import java.util.Map; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.Test; + +/** + * Integration test verifying the detach feature works over the HTTP wire using JettyPlugin. + * + *

Proves that a DETACH turn works through the Jetty agent HTTP endpoint: POST a turn with {@code + * detach:true} → server returns {@code finishReason: "detached"} + a pending {@code snapshotId} + * immediately → the background work finalizes → polling the {@code getSnapshot} companion shows + * {@code status: "completed"} with the accumulated state. + */ +class AgentDetachHttpTest { + + private static final ObjectMapper MAPPER = new ObjectMapper(); + + private JettyPlugin plugin; + + @AfterEach + void tearDown() throws Exception { + if (plugin != null) { + plugin.stop(); + } + } + + /** + * Starts a JettyPlugin on the given port with a server-managed custom agent named {@code + * detachAgent} that returns immediately (so background work finalizes quickly). + */ + private void startWithDetachAgent(int port) throws Exception { + DefaultRegistry registry = new DefaultRegistry(); + + // Register a server-managed custom agent. The AgentFn returns immediately with a STOP + // finish reason — detach background work will finalize quickly. + AgentActions.defineCustomAgent( + registry, + CustomAgentConfig.>builder() + .name("detachAgent") + .store(new InMemorySessionStore<>()) + .build(), + (sess, fnCtx) -> + AgentResult.builder() + .message(Message.model("done")) + .finishReason(AgentFinishReason.STOP) + .build()); + + JettyPluginOptions options = JettyPluginOptions.builder().port(port).build(); + plugin = new JettyPlugin(options); + plugin.init(registry); + + Thread serverThread = + new Thread( + () -> { + try { + plugin.start(); + } catch (Exception e) { + // Server stopped. + } + }); + serverThread.setDaemon(true); + serverThread.start(); + + waitForServer(port); + } + + /** + * Integration test: POST with detach:true → assert finishReason=="detached" + snapshotId → poll + * getSnapshot until status=="completed" → assert messages non-empty. + */ + @Test + void testDetachOverHttp() throws Exception { + int port = findAvailablePort(); + startWithDetachAgent(port); + + HttpClient client = HttpClient.newHttpClient(); + + // Step 1: POST a detach turn. + String detachBody = + "{\"data\":{\"detach\":true,\"message\":{\"role\":\"user\"," + + "\"content\":[{\"text\":\"go\"}]}},\"init\":{}}"; + + HttpRequest detachRequest = + HttpRequest.newBuilder() + .uri(URI.create("http://localhost:" + port + "/detachAgent")) + .header("Content-Type", "application/json") + .POST(HttpRequest.BodyPublishers.ofString(detachBody)) + .build(); + + HttpResponse detachResponse = + client.send(detachRequest, HttpResponse.BodyHandlers.ofString()); + assertEquals(200, detachResponse.statusCode(), "detach POST body: " + detachResponse.body()); + + JsonNode detachRoot = MAPPER.readTree(detachResponse.body()); + assertTrue(detachRoot.has("result"), "expected result envelope, got: " + detachResponse.body()); + + JsonNode result = detachRoot.path("result"); + + // Assert: finishReason is "detached". + String finishReason = result.path("finishReason").asText(""); + assertEquals( + "detached", + finishReason, + "expected finishReason=detached, got: " + finishReason + " body: " + detachResponse.body()); + + // Assert: snapshotId is present and non-empty. + String snapshotId = result.path("snapshotId").asText(""); + assertFalse( + snapshotId.isEmpty(), + "expected non-empty snapshotId, got: " + snapshotId + " body: " + detachResponse.body()); + + // Step 2: Poll getSnapshot until status becomes "completed" (or timeout at 5s). + String getSnapshotBody = "{\"data\":{\"snapshotId\":\"" + snapshotId + "\"}}"; + HttpRequest snapshotRequest = + HttpRequest.newBuilder() + .uri(URI.create("http://localhost:" + port + "/detachAgent/getSnapshot")) + .header("Content-Type", "application/json") + .POST(HttpRequest.BodyPublishers.ofString(getSnapshotBody)) + .build(); + + JsonNode snapResult = null; + String snapStatus = ""; + long deadline = System.currentTimeMillis() + 5000; + while (System.currentTimeMillis() < deadline) { + HttpResponse snapResponse = + client.send(snapshotRequest, HttpResponse.BodyHandlers.ofString()); + assertEquals(200, snapResponse.statusCode(), "getSnapshot body: " + snapResponse.body()); + + JsonNode snapRoot = MAPPER.readTree(snapResponse.body()); + assertTrue( + snapRoot.has("result"), + "expected result envelope from getSnapshot, got: " + snapResponse.body()); + + snapResult = snapRoot.path("result"); + snapStatus = snapResult.path("status").asText(""); + + if ("completed".equals(snapStatus) + || "failed".equals(snapStatus) + || "aborted".equals(snapStatus)) { + break; + } + Thread.sleep(100); + } + + // Assert: snapshot finalized to "completed". + assertNotNull(snapResult, "snapResult should not be null"); + assertEquals( + "completed", + snapStatus, + "expected snapshot status=completed, got: " + snapStatus + " snap: " + snapResult); + + // Assert: state.messages is non-empty (background turn ran and accumulated messages). + JsonNode messages = snapResult.path("state").path("messages"); + assertFalse( + messages.isMissingNode() || messages.isEmpty(), + "expected non-empty state.messages in completed snapshot, got: " + snapResult); + } + + private static int findAvailablePort() throws IOException { + try (java.net.ServerSocket socket = new java.net.ServerSocket(0)) { + return socket.getLocalPort(); + } + } + + private static void waitForServer(int port) throws Exception { + int maxRetries = 50; + for (int i = 0; i < maxRetries; i++) { + try { + HttpURLConnection conn = + (HttpURLConnection) new URL("http://localhost:" + port + "/health").openConnection(); + conn.setRequestMethod("GET"); + conn.setConnectTimeout(200); + conn.setReadTimeout(200); + if (conn.getResponseCode() == 200) { + return; + } + } catch (IOException e) { + // Server not ready yet. + } + Thread.sleep(100); + } + fail("Server did not start within timeout"); + } +} diff --git a/plugins/jetty/src/test/java/com/google/genkit/plugins/jetty/AgentHeaderPropagationTest.java b/plugins/jetty/src/test/java/com/google/genkit/plugins/jetty/AgentHeaderPropagationTest.java new file mode 100644 index 000000000..d150cba07 --- /dev/null +++ b/plugins/jetty/src/test/java/com/google/genkit/plugins/jetty/AgentHeaderPropagationTest.java @@ -0,0 +1,205 @@ +/* + * 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. + * + * SPDX-License-Identifier: Apache-2.0 + */ + +package com.google.genkit.plugins.jetty; + +import static org.junit.jupiter.api.Assertions.*; + +import com.google.genkit.ai.Message; +import com.google.genkit.ai.agent.AgentChat; +import com.google.genkit.ai.agent.AgentFinishReason; +import com.google.genkit.ai.agent.AgentResponse; +import com.google.genkit.ai.agent.AgentResult; +import com.google.genkit.ai.agent.CustomAgentConfig; +import com.google.genkit.ai.agent.InMemorySessionStore; +import com.google.genkit.ai.agent.internal.AgentActions; +import com.google.genkit.client.RemoteAgent; +import com.google.genkit.client.RemoteAgentOptions; +import com.google.genkit.core.ActionContext; +import com.google.genkit.core.DefaultRegistry; +import java.io.IOException; +import java.net.HttpURLConnection; +import java.net.URI; +import java.net.URL; +import java.util.HashMap; +import java.util.Map; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.Test; + +/** + * End-to-end proof that custom HTTP request headers set via {@link RemoteAgentOptions#headers()} on + * the client actually reach a served {@code AgentFn} through {@code JettyPlugin}'s {@code + * AgentHandler} / {@code CompanionHandler}. + * + *

This is the full client-to-server round trip that was previously broken: {@code + * HttpAgentTransport} already sent {@code RemoteAgentOptions.headers()} as real HTTP headers, but + * {@code JettyPlugin} never read incoming HTTP headers on the server side, so a tool/flow had no + * way to observe them. These tests start a real Jetty server, send a real HTTP request (via {@link + * RemoteAgent}/{@link RemoteAgentOptions#headers()}) carrying a custom header, and assert a + * no-model custom {@code AgentFn} can read that header's value out of {@code + * AgentFnContext#context()} and reflect it back in its response. + */ +class AgentHeaderPropagationTest { + + private JettyPlugin plugin; + + @AfterEach + void tearDown() throws Exception { + if (plugin != null) { + plugin.stop(); + } + } + + /** + * Starts a Jetty server on {@code port} hosting a server-managed, no-model custom agent named + * {@code headerEchoAgent}. The agent's {@code AgentFn} reads {@code + * fnCtx.context().getContext().get("headers")} (the reserved key that {@code JettyPlugin} merges + * incoming HTTP headers into) and echoes the requested header's value back in its reply text, so + * the test can assert on it without any model call. + */ + private void startHeaderEchoAgent(int port) throws Exception { + DefaultRegistry registry = new DefaultRegistry(); + + AgentActions.defineCustomAgent( + registry, + CustomAgentConfig.>builder() + .name("headerEchoAgent") + .store(new InMemorySessionStore<>()) + .build(), + (sess, fnCtx) -> { + String headerValue = "(missing)"; + ActionContext actionContext = fnCtx.context(); + if (actionContext != null && actionContext.getContext() != null) { + Object headersObj = actionContext.getContext().get("headers"); + if (headersObj instanceof Map headers) { + Object v = headers.get("X-Custom-Auth"); + if (v != null) { + headerValue = String.valueOf(v); + } + } + } + return AgentResult.builder() + .message(Message.model("header=" + headerValue)) + .finishReason(AgentFinishReason.STOP) + .build(); + }); + + JettyPluginOptions options = JettyPluginOptions.builder().port(port).build(); + plugin = new JettyPlugin(options); + plugin.init(registry); + + Thread serverThread = + new Thread( + () -> { + try { + plugin.start(); + } catch (Exception e) { + // Server stopped. + } + }); + serverThread.setDaemon(true); + serverThread.start(); + + waitForServer(port); + } + + /** + * Full round trip: a {@link RemoteAgent} configured with {@link + * RemoteAgentOptions.Builder#headers(Map)} sends a real HTTP POST to a live Jetty server; the + * server-side {@code AgentFn} reads the header back out of {@code ActionContext.getContext()} and + * reflects it in its reply. Proves the client-to-server header pipe end-to-end, not just that a + * map got populated somewhere. + */ + @Test + void testCustomHeaderReachesAgentFnOverHttp() throws Exception { + int port = findAvailablePort(); + startHeaderEchoAgent(port); + + Map headers = new HashMap<>(); + headers.put("X-Custom-Auth", "secret-token-123"); + + AgentChat> chat = + RemoteAgent.chat( + RemoteAgentOptions.builder() + .url("http://localhost:" + port + "/headerEchoAgent") + .headers(headers) + .build()); + + AgentResponse> resp = chat.send("hello"); + + assertNotNull(resp.text(), "expected non-null text"); + assertEquals( + "header=secret-token-123", + resp.text(), + "expected the AgentFn to read the custom HTTP header via ActionContext.getContext()"); + } + + /** + * Without any custom header configured, the AgentFn should observe the reserved {@code "headers"} + * key as either absent or not containing {@code X-Custom-Auth} — i.e. the header plumbing does + * not fabricate values, and the (missing) fallback path is exercised. + */ + @Test + void testNoCustomHeaderMeansMissingInAgentFn() throws Exception { + int port = findAvailablePort(); + startHeaderEchoAgent(port); + + AgentChat> chat = + RemoteAgent.chat( + RemoteAgentOptions.builder() + .url("http://localhost:" + port + "/headerEchoAgent") + .build()); + + AgentResponse> resp = chat.send("hello"); + + assertEquals( + "header=(missing)", + resp.text(), + "expected no X-Custom-Auth header to be observed when the client sets none"); + } + + // --------------------------------------------------------------------------- + // Helpers + // --------------------------------------------------------------------------- + + private static int findAvailablePort() throws IOException { + try (java.net.ServerSocket socket = new java.net.ServerSocket(0)) { + return socket.getLocalPort(); + } + } + + private static void waitForServer(int port) throws Exception { + int maxRetries = 50; + for (int i = 0; i < maxRetries; i++) { + try { + URL healthUrl = new URI("http://localhost:" + port + "/health").toURL(); + HttpURLConnection conn = (HttpURLConnection) healthUrl.openConnection(); + conn.setRequestMethod("GET"); + conn.setConnectTimeout(200); + conn.setReadTimeout(200); + if (conn.getResponseCode() == 200) { + return; + } + } catch (IOException e) { + // Server not ready yet. + } + Thread.sleep(100); + } + fail("Server did not start within timeout"); + } +} diff --git a/plugins/jetty/src/test/java/com/google/genkit/plugins/jetty/AgentHttpServingTest.java b/plugins/jetty/src/test/java/com/google/genkit/plugins/jetty/AgentHttpServingTest.java new file mode 100644 index 000000000..8b5d95027 --- /dev/null +++ b/plugins/jetty/src/test/java/com/google/genkit/plugins/jetty/AgentHttpServingTest.java @@ -0,0 +1,246 @@ +/* + * 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. + * + * SPDX-License-Identifier: Apache-2.0 + */ + +package com.google.genkit.plugins.jetty; + +import static org.junit.jupiter.api.Assertions.*; + +import com.fasterxml.jackson.databind.JsonNode; +import com.fasterxml.jackson.databind.ObjectMapper; +import com.fasterxml.jackson.databind.node.ObjectNode; +import com.google.genkit.core.ActionDef; +import com.google.genkit.core.ActionType; +import com.google.genkit.core.BidiActionImpl; +import com.google.genkit.core.DefaultRegistry; +import java.io.IOException; +import java.net.HttpURLConnection; +import java.net.URI; +import java.net.URL; +import java.net.http.HttpClient; +import java.net.http.HttpRequest; +import java.net.http.HttpResponse; +import java.util.Optional; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.Test; + +/** + * Integration tests for agent HTTP serving in the Jetty plugin. + * + *

Verifies the one-turn-per-request transport for {@code ActionType.AGENT} bidi actions: a + * non-streaming {@code {data, init}} request, an SSE streaming request, and a companion {@code + * getSnapshot} endpoint. + */ +class AgentHttpServingTest { + + private static final ObjectMapper MAPPER = new ObjectMapper(); + + private JettyPlugin plugin; + + @AfterEach + void tearDown() throws Exception { + if (plugin != null) { + plugin.stop(); + } + } + + /** + * Builds a registry containing a server-managed bidi agent named {@code chatAgent} plus a + * companion {@code agent-snapshot} action, then starts a JettyPlugin on the given port. + */ + private void startWithChatAgent(int port) throws Exception { + DefaultRegistry registry = new DefaultRegistry(); + + // A minimal server-managed bidi agent: reads one input, emits one stream chunk, returns a + // final output containing snapshotId + an echoed message. + BidiActionImpl agent = + BidiActionImpl.builder() + .name("chatAgent") + .inputClass(JsonNode.class) + .outputClass(JsonNode.class) + .streamClass(JsonNode.class) + .initClass(JsonNode.class) + .handler( + (ctx, init, inputs, cb) -> { + Optional first = inputs.next(); + JsonNode data = first.orElse(MAPPER.nullNode()); + // The AgentInput envelope carries the turn message under "message". + JsonNode message = data.path("message"); + // Emit one streamed chunk. + if (cb != null) { + ObjectNode chunk = MAPPER.createObjectNode(); + chunk.put("text", "thinking..."); + cb.accept(chunk); + } + // Final output echoes the message back. + ObjectNode result = MAPPER.createObjectNode(); + result.put("snapshotId", "s1"); + result.set("message", message); + return result; + }) + .build(); + agent.register(registry); + + // Companion agent-snapshot action looked up at key "/agent-snapshot/chatAgent". + ActionDef snapshot = + new ActionDef<>( + "chatAgent", + ActionType.AGENT_SNAPSHOT, + null, + null, + JsonNode.class, + JsonNode.class, + (ctx, input, cb) -> { + ObjectNode snap = MAPPER.createObjectNode(); + snap.put("snapshotId", input != null ? input.path("snapshotId").asText("?") : "?"); + snap.put("status", "captured"); + return snap; + }); + snapshot.register(registry); + + JettyPluginOptions options = JettyPluginOptions.builder().port(port).build(); + plugin = new JettyPlugin(options); + plugin.init(registry); + + Thread serverThread = + new Thread( + () -> { + try { + plugin.start(); + } catch (Exception e) { + // Server stopped. + } + }); + serverThread.setDaemon(true); + serverThread.start(); + + waitForServer(port); + } + + @Test + void testAgentNonStreaming() throws Exception { + int port = findAvailablePort(); + startWithChatAgent(port); + + String body = + "{\"data\":{\"message\":{\"role\":\"user\",\"content\":[{\"text\":\"hi\"}]}},\"init\":{}}"; + + HttpClient client = HttpClient.newHttpClient(); + HttpRequest request = + HttpRequest.newBuilder() + .uri(URI.create("http://localhost:" + port + "/chatAgent")) + .header("Content-Type", "application/json") + .POST(HttpRequest.BodyPublishers.ofString(body)) + .build(); + + HttpResponse response = client.send(request, HttpResponse.BodyHandlers.ofString()); + assertEquals(200, response.statusCode(), "body: " + response.body()); + + JsonNode root = MAPPER.readTree(response.body()); + assertTrue(root.has("result"), "expected result envelope: " + response.body()); + assertEquals("s1", root.path("result").path("snapshotId").asText()); + // Echoed message should be present. + assertEquals( + "hi", root.path("result").path("message").path("content").path(0).path("text").asText()); + } + + @Test + void testAgentStreamingSse() throws Exception { + int port = findAvailablePort(); + startWithChatAgent(port); + + String body = + "{\"data\":{\"message\":{\"role\":\"user\",\"content\":[{\"text\":\"hi\"}]}},\"init\":{}}"; + + HttpClient client = HttpClient.newHttpClient(); + HttpRequest request = + HttpRequest.newBuilder() + .uri(URI.create("http://localhost:" + port + "/chatAgent")) + .header("Content-Type", "application/json") + .header("Accept", "text/event-stream") + .POST(HttpRequest.BodyPublishers.ofString(body)) + .build(); + + HttpResponse response = client.send(request, HttpResponse.BodyHandlers.ofString()); + assertEquals(200, response.statusCode(), "body: " + response.body()); + + String contentType = response.headers().firstValue("Content-Type").orElse(""); + assertTrue(contentType.contains("text/event-stream"), "content-type: " + contentType); + assertTrue( + response.headers().firstValue("X-Genkit-Stream-Id").isPresent(), + "expected X-Genkit-Stream-Id header"); + + String text = response.body(); + // A message frame then a result frame. + assertTrue(text.contains("\"message\""), "expected a message frame: " + text); + assertTrue(text.contains("thinking..."), "expected streamed chunk text: " + text); + assertTrue(text.contains("\"result\""), "expected a result frame: " + text); + assertTrue(text.contains("\"snapshotId\":\"s1\""), "expected snapshotId in result: " + text); + // Frames are SSE-formatted. + assertTrue(text.contains("data: "), "expected SSE data prefix: " + text); + } + + @Test + void testAgentSnapshotCompanion() throws Exception { + int port = findAvailablePort(); + startWithChatAgent(port); + + String body = "{\"data\":{\"snapshotId\":\"s1\"}}"; + + HttpClient client = HttpClient.newHttpClient(); + HttpRequest request = + HttpRequest.newBuilder() + .uri(URI.create("http://localhost:" + port + "/chatAgent/getSnapshot")) + .header("Content-Type", "application/json") + .POST(HttpRequest.BodyPublishers.ofString(body)) + .build(); + + HttpResponse response = client.send(request, HttpResponse.BodyHandlers.ofString()); + assertEquals(200, response.statusCode(), "body: " + response.body()); + + JsonNode root = MAPPER.readTree(response.body()); + assertTrue(root.has("result"), "expected result envelope: " + response.body()); + assertEquals("s1", root.path("result").path("snapshotId").asText()); + assertEquals("captured", root.path("result").path("status").asText()); + } + + private static int findAvailablePort() throws IOException { + try (java.net.ServerSocket socket = new java.net.ServerSocket(0)) { + return socket.getLocalPort(); + } + } + + private static void waitForServer(int port) throws Exception { + int maxRetries = 50; + for (int i = 0; i < maxRetries; i++) { + try { + HttpURLConnection conn = + (HttpURLConnection) new URL("http://localhost:" + port + "/health").openConnection(); + conn.setRequestMethod("GET"); + conn.setConnectTimeout(200); + conn.setReadTimeout(200); + if (conn.getResponseCode() == 200) { + return; + } + } catch (IOException e) { + // Server not ready yet. + } + Thread.sleep(100); + } + fail("Server did not start within timeout"); + } +} diff --git a/plugins/jetty/src/test/java/com/google/genkit/plugins/jetty/RemoteAgentTest.java b/plugins/jetty/src/test/java/com/google/genkit/plugins/jetty/RemoteAgentTest.java new file mode 100644 index 000000000..d1fb7c5e1 --- /dev/null +++ b/plugins/jetty/src/test/java/com/google/genkit/plugins/jetty/RemoteAgentTest.java @@ -0,0 +1,487 @@ +/* + * 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. + * + * SPDX-License-Identifier: Apache-2.0 + */ + +package com.google.genkit.plugins.jetty; + +import static org.junit.jupiter.api.Assertions.*; + +import com.fasterxml.jackson.databind.JsonNode; +import com.fasterxml.jackson.databind.ObjectMapper; +import com.fasterxml.jackson.databind.node.ObjectNode; +import com.google.genkit.ai.Message; +import com.google.genkit.ai.agent.AgentChat; +import com.google.genkit.ai.agent.AgentFinishReason; +import com.google.genkit.ai.agent.AgentResponse; +import com.google.genkit.ai.agent.AgentResult; +import com.google.genkit.ai.agent.CustomAgentConfig; +import com.google.genkit.ai.agent.GetSnapshotRequest; +import com.google.genkit.ai.agent.InMemorySessionStore; +import com.google.genkit.ai.agent.SessionSnapshot; +import com.google.genkit.ai.agent.internal.AgentActions; +import com.google.genkit.client.HttpAgentTransport; +import com.google.genkit.client.RemoteAgent; +import com.google.genkit.client.RemoteAgentOptions; +import com.google.genkit.core.ActionDef; +import com.google.genkit.core.ActionType; +import com.google.genkit.core.BidiActionImpl; +import com.google.genkit.core.DefaultRegistry; +import java.io.IOException; +import java.net.HttpURLConnection; +import java.net.URI; +import java.net.URL; +import java.util.ArrayList; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.Optional; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicInteger; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.Test; + +/** + * Integration tests for {@link RemoteAgent} / {@link HttpAgentTransport} against a live Jetty + * server running a server-managed echo agent. + * + *

Placed in the {@code plugins/jetty} module because that module depends on {@code genkit} + * (which contains {@code RemoteAgent}) and can also start a {@link JettyPlugin} — giving us both + * halves of the test in one module without a circular dependency. + */ +class RemoteAgentTest { + + private static final ObjectMapper MAPPER = new ObjectMapper(); + + private JettyPlugin plugin; + + @AfterEach + void tearDown() throws Exception { + if (plugin != null) { + plugin.stop(); + } + } + + // --------------------------------------------------------------------------- + // Server setup + // --------------------------------------------------------------------------- + + /** + * Starts a Jetty server on {@code port} hosting a server-managed echo agent named {@code + * echoAgent}. + * + *

The agent: + * + *

    + *
  • Emits one SSE chunk with {@code {"text":"streaming..."}}. + *
  • Returns a final output whose {@code snapshotId} is {@code "snap-"} (turn counter + * increments on each invocation), {@code sessionId} is {@code "session-1"}, and {@code + * message} echoes the user message content back. + *
+ * + *

A companion {@code agent-snapshot} action is registered so {@code getSnapshot} works. + */ + private void startEchoAgent(int port) throws Exception { + DefaultRegistry registry = new DefaultRegistry(); + + AtomicInteger turnCounter = new AtomicInteger(0); + + BidiActionImpl agent = + BidiActionImpl.builder() + .name("echoAgent") + .inputClass(JsonNode.class) + .outputClass(JsonNode.class) + .streamClass(JsonNode.class) + .initClass(JsonNode.class) + .handler( + (ctx, init, inputs, cb) -> { + Optional first = inputs.next(); + JsonNode data = first.orElse(MAPPER.nullNode()); + JsonNode messageNode = data.path("message"); + + int turn = turnCounter.incrementAndGet(); + + // Emit one streamed chunk. + if (cb != null) { + ObjectNode chunk = MAPPER.createObjectNode(); + chunk.put("text", "streaming..."); + cb.accept(chunk); + } + + // Build final output. + ObjectNode result = MAPPER.createObjectNode(); + result.put("snapshotId", "snap-" + turn); + result.put("sessionId", "session-1"); + result.put("finishReason", AgentFinishReason.STOP.getValue()); + // Echo message back. + ObjectNode message = MAPPER.createObjectNode(); + message.put("role", "model"); + ObjectNode part = MAPPER.createObjectNode(); + // Echo the first text part of the user message. + String userText = + messageNode.path("content").path(0).path("text").asText("(empty)"); + part.put("text", "echo: " + userText); + message.putArray("content").add(part); + result.set("message", message); + return result; + }) + .build(); + agent.register(registry); + + // Companion snapshot action. + ActionDef snapshot = + new ActionDef<>( + "echoAgent", + ActionType.AGENT_SNAPSHOT, + null, + null, + JsonNode.class, + JsonNode.class, + (ctx, input, cb) -> { + ObjectNode snap = MAPPER.createObjectNode(); + String sid = + (input != null && input.has("snapshotId")) + ? input.get("snapshotId").asText("?") + : "?"; + snap.put("snapshotId", sid); + snap.put("sessionId", "session-1"); + snap.put("status", "completed"); + return snap; + }); + snapshot.register(registry); + + JettyPluginOptions options = JettyPluginOptions.builder().port(port).build(); + plugin = new JettyPlugin(options); + plugin.init(registry); + + Thread serverThread = + new Thread( + () -> { + try { + plugin.start(); + } catch (Exception e) { + // Server stopped. + } + }); + serverThread.setDaemon(true); + serverThread.start(); + + waitForServer(port); + } + + // --------------------------------------------------------------------------- + // Tests + // --------------------------------------------------------------------------- + + @Test + void testSendFirstTurn() throws Exception { + int port = findAvailablePort(); + startEchoAgent(port); + + AgentChat> chat = + RemoteAgent.chat( + RemoteAgentOptions.builder().url("http://localhost:" + port + "/echoAgent").build()); + + AgentResponse> resp = chat.send("hello"); + + // Response text must be non-empty (echoed user message). + assertNotNull(resp.text(), "expected non-null text"); + assertFalse(resp.text().isEmpty(), "expected non-empty text, got: " + resp.text()); + assertTrue(resp.text().contains("echo: hello"), "expected echo in text, got: " + resp.text()); + + // snapshotId must be set after the first turn. + assertNotNull(chat.snapshotId(), "expected snapshotId to be set after first turn"); + assertEquals("snap-1", chat.snapshotId()); + } + + @Test + void testSendSecondTurnResumes() throws Exception { + int port = findAvailablePort(); + startEchoAgent(port); + + AgentChat> chat = + RemoteAgent.chat( + RemoteAgentOptions.builder().url("http://localhost:" + port + "/echoAgent").build()); + + chat.send("first"); + AgentResponse> resp2 = chat.send("second"); + + // Second turn should have incremented the counter. + assertEquals("snap-2", chat.snapshotId()); + assertTrue( + resp2.text().contains("echo: second"), "expected echo of second turn: " + resp2.text()); + } + + @Test + void testStreamingChunksDelivered() throws Exception { + int port = findAvailablePort(); + startEchoAgent(port); + + AgentChat> chat = + RemoteAgent.chat( + RemoteAgentOptions.builder().url("http://localhost:" + port + "/echoAgent").build()); + + List chunks = new ArrayList<>(); + chat.sendStream( + "streaming test", + chunk -> { + // Chunks may contain a modelChunk; collect whatever the server sends. + chunks.add(chunk.toString()); + }); + + // The server emits one streaming chunk; the onChunk callback must have been called. + // (AgentStreamChunk is not a plain text — it wraps modelChunk; we just verify it fired.) + assertFalse(chunks.isEmpty(), "expected at least one SSE chunk"); + } + + @Test + void testGetSnapshot() throws Exception { + int port = findAvailablePort(); + startEchoAgent(port); + + AgentChat> chat = + RemoteAgent.chat( + RemoteAgentOptions.builder().url("http://localhost:" + port + "/echoAgent").build()); + + chat.send("hi"); + String snapId = chat.snapshotId(); + assertNotNull(snapId, "expected snapshotId after send"); + + // Retrieve snapshot via the transport directly. + HttpAgentTransport> transport = + new HttpAgentTransport<>( + RemoteAgentOptions.builder().url("http://localhost:" + port + "/echoAgent").build()); + + SessionSnapshot> snap = + transport.getSnapshot(GetSnapshotRequest.builder().snapshotId(snapId).build()); + + assertNotNull(snap, "expected non-null snapshot"); + assertEquals(snapId, snap.getSnapshotId(), "snapshot ID should match requested ID"); + } + + /** + * Starts a Jetty server on {@code port} hosting a server-managed custom agent named {@code + * blockingAgent} whose {@code AgentFn} blocks on {@code releaseLatch} until released. Used to + * open a deterministic {@code PENDING} window for {@link #testAbortOverHttp()}. + */ + private void startBlockingAgent(int port, CountDownLatch releaseLatch) throws Exception { + DefaultRegistry registry = new DefaultRegistry(); + + AgentActions.defineCustomAgent( + registry, + CustomAgentConfig.>builder() + .name("blockingAgent") + .store(new InMemorySessionStore<>()) + .build(), + (sess, fnCtx) -> { + releaseLatch.await(10, TimeUnit.SECONDS); + return AgentResult.builder() + .message(Message.model("done")) + .finishReason(AgentFinishReason.STOP) + .build(); + }); + + JettyPluginOptions options = JettyPluginOptions.builder().port(port).build(); + plugin = new JettyPlugin(options); + plugin.init(registry); + + Thread serverThread = + new Thread( + () -> { + try { + plugin.start(); + } catch (Exception e) { + // Server stopped. + } + }); + serverThread.setDaemon(true); + serverThread.start(); + + waitForServer(port); + } + + /** + * Starts a Jetty server on {@code port} hosting a client-managed (no store) custom agent named + * {@code counterAgent} that increments an integer counter in its custom state each turn and + * echoes back the running total in its reply text. + */ + private void startClientManagedCounterAgent(int port) throws Exception { + DefaultRegistry registry = new DefaultRegistry(); + + AgentActions.defineCustomAgent( + registry, + CustomAgentConfig.>builder() + .name("counterAgent") + // No store() call => client-managed: state round-trips via AgentInit/AgentOutput. + .build(), + (sess, fnCtx) -> { + Map custom = sess.getCustom(); + int count = 0; + if (custom != null && custom.get("count") instanceof Number) { + count = ((Number) custom.get("count")).intValue(); + } + count++; + Map updated = new HashMap<>(); + updated.put("count", count); + sess.updateCustom(c -> updated); + return AgentResult.builder() + .message(Message.model("count=" + count)) + .finishReason(AgentFinishReason.STOP) + .build(); + }); + + JettyPluginOptions options = JettyPluginOptions.builder().port(port).build(); + plugin = new JettyPlugin(options); + plugin.init(registry); + + Thread serverThread = + new Thread( + () -> { + try { + plugin.start(); + } catch (Exception e) { + // Server stopped. + } + }); + serverThread.setDaemon(true); + serverThread.start(); + + waitForServer(port); + } + + /** + * Verifies {@link AgentChat#abort()} (backed by {@link HttpAgentTransport#abort}) reaches the + * server's {@code /abort} companion endpoint and that a subsequent {@code getSnapshot} (via the + * transport directly, mirroring how {@code AgentChat.loadChat} would read it back) reflects the + * flipped status. + * + *

Scoped to what is actually implemented: aborting a snapshot that is genuinely {@code + * PENDING} (opened by a detached turn whose {@code AgentFn} we block deterministically) — not + * interrupting a live foreground call, which {@code Agent.abort()} / the abort companion action + * do not support (no cancellation hook is threaded into a running {@code AgentFn}). + */ + @Test + void testAbortOverHttp() throws Exception { + int port = findAvailablePort(); + CountDownLatch releaseLatch = new CountDownLatch(1); + startBlockingAgent(port, releaseLatch); + + RemoteAgentOptions opts = + RemoteAgentOptions.builder().url("http://localhost:" + port + "/blockingAgent").build(); + AgentChat> chat = RemoteAgent.chat(opts); + + // Detach so the snapshot is written PENDING and the background AgentFn blocks on the latch. + chat.sendStream( + com.google.genkit.ai.agent.AgentInput.builder() + .message(Message.user("go")) + .detach(true) + .build(), + c -> {}); + + String snapshotId = chat.snapshotId(); + assertNotNull(snapshotId, "expected snapshotId after detach"); + + // Abort via the client — reaches the server's /abort companion endpoint. + com.google.genkit.ai.agent.SnapshotStatus status = chat.abort(); + assertEquals( + com.google.genkit.ai.agent.SnapshotStatus.ABORTED, + status, + "expected client-side abort() to report ABORTED"); + + // Read the snapshot back directly via the transport (as AgentChat.loadChat would internally). + HttpAgentTransport> transport = new HttpAgentTransport<>(opts); + SessionSnapshot> snap = + transport.getSnapshot(GetSnapshotRequest.builder().snapshotId(snapshotId).build()); + assertNotNull(snap, "expected non-null snapshot"); + assertEquals( + com.google.genkit.ai.agent.SnapshotStatus.ABORTED, + snap.getStatus(), + "expected stored snapshot status to be ABORTED"); + + // Release the blocked background turn; its finalizer must not revert the ABORTED status. + releaseLatch.countDown(); + Thread.sleep(300); + SessionSnapshot> snapAfter = + transport.getSnapshot(GetSnapshotRequest.builder().snapshotId(snapshotId).build()); + assertEquals( + com.google.genkit.ai.agent.SnapshotStatus.ABORTED, + snapAfter.getStatus(), + "expected status to remain ABORTED after the background turn finalizes"); + } + + /** + * Verifies a client-managed agent (no {@code SessionStore}) genuinely round-trips its session + * state over real HTTP: {@link AgentChat} carries the full {@link + * com.google.genkit.ai.agent.SessionState} (including custom state) in {@code AgentInit} on each + * turn, and the served agent's reply on turn 2 reflects state seeded by turn 1 — proving + * serialization/deserialization across the wire, not just in-process object sharing. + */ + @Test + void testClientManagedRemoteAgent() throws Exception { + int port = findAvailablePort(); + startClientManagedCounterAgent(port); + + AgentChat> chat = + RemoteAgent.chat( + RemoteAgentOptions.builder() + .url("http://localhost:" + port + "/counterAgent") + .serverManaged(false) + .build()); + + AgentResponse> resp1 = chat.send("first"); + assertEquals("count=1", resp1.text(), "expected first turn to start the counter at 1"); + + AgentResponse> resp2 = chat.send("second"); + assertEquals( + "count=2", + resp2.text(), + "expected second turn's reply to reflect state seeded from turn 1 over the wire"); + + AgentResponse> resp3 = chat.send("third"); + assertEquals("count=3", resp3.text(), "expected third turn to continue incrementing"); + } + + // --------------------------------------------------------------------------- + // Helpers + // --------------------------------------------------------------------------- + + private static int findAvailablePort() throws IOException { + try (java.net.ServerSocket socket = new java.net.ServerSocket(0)) { + return socket.getLocalPort(); + } + } + + private static void waitForServer(int port) throws Exception { + int maxRetries = 50; + for (int i = 0; i < maxRetries; i++) { + try { + URL healthUrl = new URI("http://localhost:" + port + "/health").toURL(); + HttpURLConnection conn = (HttpURLConnection) healthUrl.openConnection(); + conn.setRequestMethod("GET"); + conn.setConnectTimeout(200); + conn.setReadTimeout(200); + if (conn.getResponseCode() == 200) { + return; + } + } catch (IOException e) { + // Server not ready yet. + } + Thread.sleep(100); + } + fail("Server did not start within timeout"); + } +} diff --git a/plugins/middleware/pom.xml b/plugins/middleware/pom.xml new file mode 100644 index 000000000..ff695804d --- /dev/null +++ b/plugins/middleware/pom.xml @@ -0,0 +1,92 @@ + + + + 4.0.0 + + + com.google.genkit + genkit-parent + 1.0.0-SNAPSHOT + ../../pom.xml + + + genkit-plugin-middleware + jar + Genkit Middleware Plugin + Generation middleware for Genkit - sub-agent delegation (agents) and artifact tools (artifacts) + + + false + + + + + + com.google.genkit + genkit-core + ${project.version} + + + com.google.genkit + genkit-ai + ${project.version} + + + com.google.genkit + genkit + ${project.version} + + + + + com.fasterxml.jackson.core + jackson-databind + + + + + org.slf4j + slf4j-api + + + + + org.junit.jupiter + junit-jupiter + test + + + org.mockito + mockito-core + test + + + + + + + org.apache.maven.plugins + maven-compiler-plugin + + + + diff --git a/plugins/middleware/src/main/java/com/google/genkit/plugins/middleware/Agents.java b/plugins/middleware/src/main/java/com/google/genkit/plugins/middleware/Agents.java new file mode 100644 index 000000000..3fdec1fcc --- /dev/null +++ b/plugins/middleware/src/main/java/com/google/genkit/plugins/middleware/Agents.java @@ -0,0 +1,244 @@ +/* + * 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. + * + * SPDX-License-Identifier: Apache-2.0 + */ + +package com.google.genkit.plugins.middleware; + +import com.fasterxml.jackson.annotation.JsonInclude; +import com.fasterxml.jackson.annotation.JsonProperty; +import com.google.genkit.ai.Tool; +import com.google.genkit.plugins.middleware.internal.Delegation; +import java.util.ArrayList; +import java.util.List; +import java.util.concurrent.atomic.AtomicInteger; + +/** + * Factory for sub-agent delegation tools. + * + *

Each configured sub-agent is exposed to the delegating model as a {@code _} tool + * (default prefix {@code delegate_to}, so {@code delegate_to_}). When the model calls that + * tool with {@code {"task": "..."}}, the named sub-agent (resolved from the registry at {@code + * /agent/}) is run for a single turn and its assistant text is returned as {@code response}. + * Any artifacts the sub-agent produces are namespaced ({@code /}) and merged + * into the active parent session via {@link + * com.google.genkit.ai.agent.AgentSessionContext#currentArtifactStore()}. + * + *

Pragmatic scope decision (tools-factory, not a generate hook)

+ * + *

This repo's generate pipeline has no first-class "wrap generate / inject system prompt" + * middleware seam that an agent's synthesized {@code AgentFn} runs through, but {@code AgentConfig} + * already exposes {@code tools(List)} and {@code system(String)}. {@code Agents} is therefore + * implemented as a tool factory: {@link #delegationTools(AgentsOptions)} returns the + * delegation tools and {@link #systemPromptFragment(AgentsOptions)} returns a {@code } + * system-prompt fragment. Callers wire them in via {@code AgentConfig.tools(...)} and {@code + * AgentConfig.system(...)}. This achieves model-driven sub-agent delegation without modifying the + * generate pipeline. + * + *

History forwarding

+ * + *

Each (parent session, sub-agent) pair is delegated to through its own stable sub-session (see + * {@code Delegation} for the exact key derivation and mechanics). Repeated delegation calls to the + * SAME sub-agent within the SAME parent session see up to {@code historyLength} trailing messages + * of that sub-session's own accumulated history; {@code historyLength <= 0} (the default) means + * task-only, no history forwarded. For client-managed sub-agents this is fully honored via an + * internal history ledger forwarded as {@code AgentInit.state}. For server-managed sub-agents + * (configured with a {@code SessionStore}), history is forwarded by resuming the sub-agent's own + * store-backed session via {@code AgentInit.sessionId} — the store's full accumulated history is + * used and {@code historyLength} trimming is not applied in that case (no seam exists to truncate a + * resolved session before the sub-agent's turn runs without mutating its own persisted snapshot). + * + *

Structured failure/interrupt propagation

+ * + *

A sub-agent turn that finishes {@code INTERRUPTED} causes the delegation tool to throw {@link + * com.google.genkit.ai.ToolInterruptException}, carrying the sub-agent's interrupted tool + * name/input as metadata, so the parent's own generate loop pauses too (nested interrupt + * visibility). Resuming the parent's interrupt surfaces that information; it does not, by itself, + * re-invoke the sub-agent with a resume — full nested-resume is out of scope for this fix. A + * sub-agent turn that finishes {@code FAILED} (or whose {@code runBidiJson} call itself throws) + * causes the delegation tool to throw {@link com.google.genkit.core.GenkitException} carrying the + * sub-agent's error message, so the failure propagates like any other tool exception. + * + *

Other limitations (v1)

+ * + *
    + *
  • {@code maxDelegations} is enforced per {@code delegationTools(...)} invocation, shared + * across all delegation tools returned by that call (an in-memory counter). + *
+ */ +public final class Agents { + + private Agents() {} + + /** + * Builds the delegation tools (one per configured sub-agent). + * + *

The returned tools share a single delegation counter so that {@code + * options.getMaxDelegations()} is enforced across all of them for the lifetime of this call's + * result. + * + * @param options the options (must not be null; at least one agent) + * @return the delegation tools + */ + public static List> delegationTools(AgentsOptions options) { + if (options == null) { + throw new IllegalArgumentException("options must not be null"); + } + boolean includeContent = options.getArtifactStrategy() == ArtifactStrategy.INLINE; + int maxDelegations = options.getMaxDelegations(); + AtomicInteger counter = new AtomicInteger(0); + + List> tools = new ArrayList<>(); + for (String agentName : options.getAgents()) { + tools.add(delegationTool(agentName, options, includeContent, maxDelegations, counter)); + } + return tools; + } + + /** + * Builds the tool name for a sub-agent given the configured prefix. An empty prefix yields the + * bare agent name. + * + * @param prefix the tool prefix + * @param agentName the sub-agent name + * @return the delegation tool name + */ + public static String toolName(String prefix, String agentName) { + if (prefix == null || prefix.isEmpty()) { + return agentName; + } + return prefix + "_" + agentName; + } + + /** + * Builds a {@code } system-prompt fragment listing each delegation tool and the agent + * it delegates to. + * + * @param options the options (must not be null) + * @return the system-prompt fragment + */ + public static String systemPromptFragment(AgentsOptions options) { + if (options == null) { + throw new IllegalArgumentException("options must not be null"); + } + StringBuilder sb = new StringBuilder(); + sb.append("\n"); + sb.append( + "You can delegate a self-contained task to a specialized sub-agent by calling one of the" + + " tools below with a 'task' describing what it should do. Use the sub-agent's text" + + " response to continue your own work.\n"); + for (String agentName : options.getAgents()) { + String tool = toolName(options.getToolPrefix(), agentName); + sb.append(" - ") + .append(tool) + .append(": delegates to the '") + .append(agentName) + .append("' agent.\n"); + } + sb.append(""); + return sb.toString(); + } + + // ── Internal ────────────────────────────────────────────────────────────────── + + private static Tool delegationTool( + String agentName, + AgentsOptions options, + boolean includeContent, + int maxDelegations, + AtomicInteger counter) { + String tool = toolName(options.getToolPrefix(), agentName); + return Tool.builder() + .name(tool) + .description( + "Delegate a self-contained task to the '" + + agentName + + "' sub-agent and return its text response.") + .inputClass(DelegateInput.class) + .outputClass(DelegateOutput.class) + .handler( + (ctx, in) -> { + if (maxDelegations > 0 && counter.incrementAndGet() > maxDelegations) { + return DelegateOutput.text( + "Delegation limit reached (" + + maxDelegations + + "). Answer using the information already gathered."); + } + String task = in != null ? in.task : null; + Delegation.Result result = + Delegation.run(ctx, agentName, task, includeContent, options.getHistoryLength()); + + DelegateOutput out = new DelegateOutput(); + out.response = result.response; + if (result.artifacts != null && !result.artifacts.isEmpty()) { + out.artifacts = new ArrayList<>(); + for (Delegation.NamedArtifact na : result.artifacts) { + ArtifactRef ref = new ArtifactRef(); + ref.name = na.name; + ref.content = na.content; + out.artifacts.add(ref); + } + } + return out; + }) + .build(); + } + + // ── Tool I/O types ────────────────────────────────────────────────────────────── + + /** Input for a delegation tool. */ + @JsonInclude(JsonInclude.Include.NON_NULL) + public static final class DelegateInput { + @JsonProperty("task") + public String task; + + /** Default constructor for JSON deserialization. */ + public DelegateInput() {} + } + + /** Output for a delegation tool. */ + @JsonInclude(JsonInclude.Include.NON_NULL) + public static final class DelegateOutput { + @JsonProperty("response") + public String response; + + @JsonProperty("artifacts") + public List artifacts; + + /** Default constructor for JSON deserialization. */ + public DelegateOutput() {} + + static DelegateOutput text(String response) { + DelegateOutput o = new DelegateOutput(); + o.response = response; + return o; + } + } + + /** A namespaced reference to an artifact produced by a sub-agent. */ + @JsonInclude(JsonInclude.Include.NON_NULL) + public static final class ArtifactRef { + @JsonProperty("name") + public String name; + + /** Content; populated only under {@link ArtifactStrategy#INLINE}. */ + @JsonProperty("content") + public String content; + + /** Default constructor for JSON deserialization. */ + public ArtifactRef() {} + } +} diff --git a/plugins/middleware/src/main/java/com/google/genkit/plugins/middleware/AgentsOptions.java b/plugins/middleware/src/main/java/com/google/genkit/plugins/middleware/AgentsOptions.java new file mode 100644 index 000000000..36e04ece4 --- /dev/null +++ b/plugins/middleware/src/main/java/com/google/genkit/plugins/middleware/AgentsOptions.java @@ -0,0 +1,227 @@ +/* + * 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. + * + * SPDX-License-Identifier: Apache-2.0 + */ + +package com.google.genkit.plugins.middleware; + +import com.google.genkit.ai.agent.AgentRef; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.List; + +/** + * Options controlling the sub-agent delegation tools produced by {@link Agents}. + * + *

The only required field is {@link #getAgents()} (at least one sub-agent). The remaining fields + * have the following defaults: + * + *

    + *
  • {@code toolPrefix} — {@code "delegate_to"}; an empty string means the bare agent name + * is used as the tool name. + *
  • {@code maxDelegations} — {@code 0} (unlimited). + *
  • {@code historyLength} — {@code 0} (task only; see {@link Agents} for exactly how a + * non-zero value is forwarded to client-managed vs. server-managed sub-agents). + *
  • {@code artifactStrategy} — {@link ArtifactStrategy#INLINE}. + *
+ */ +public final class AgentsOptions { + + private final List agents; + private final String toolPrefix; + private final int maxDelegations; + private final int historyLength; + private final ArtifactStrategy artifactStrategy; + + private AgentsOptions(Builder builder) { + if (builder.agents == null || builder.agents.isEmpty()) { + throw new IllegalArgumentException("at least one agent is required"); + } + this.agents = new ArrayList<>(builder.agents); + this.toolPrefix = builder.toolPrefix != null ? builder.toolPrefix : "delegate_to"; + this.maxDelegations = builder.maxDelegations; + this.historyLength = builder.historyLength; + this.artifactStrategy = + builder.artifactStrategy != null ? builder.artifactStrategy : ArtifactStrategy.INLINE; + } + + /** + * Creates a builder for {@link AgentsOptions}. + * + * @return a new builder + */ + public static Builder builder() { + return new Builder(); + } + + /** + * Returns the configured sub-agent names. + * + * @return the sub-agent names (never empty) + */ + public List getAgents() { + return agents; + } + + /** + * Returns the tool-name prefix. An empty string means the bare agent name is the tool name. + * + * @return the tool prefix + */ + public String getToolPrefix() { + return toolPrefix; + } + + /** + * Returns the maximum number of delegations per parent generate invocation. {@code 0} means + * unlimited. + * + * @return the max delegations + */ + public int getMaxDelegations() { + return maxDelegations; + } + + /** + * Returns the number of trailing prior messages (from this sub-agent's own accumulated + * conversation with the delegating parent session, not the parent's own history) to forward on + * each delegation call. {@code 0} means task only, no history forwarded. + * + *

See {@link Agents} for exactly how this is applied for client-managed vs. server-managed + * sub-agents. + * + * @return the history length + */ + public int getHistoryLength() { + return historyLength; + } + + /** + * Returns the artifact strategy. + * + * @return the artifact strategy + */ + public ArtifactStrategy getArtifactStrategy() { + return artifactStrategy; + } + + /** Builder for {@link AgentsOptions}. */ + public static final class Builder { + private List agents; + private String toolPrefix; + private int maxDelegations; + private int historyLength; + private ArtifactStrategy artifactStrategy; + + private Builder() {} + + /** + * Sets the sub-agent names (required, at least one). + * + * @param agents the sub-agent names + * @return this builder + */ + public Builder agents(List agents) { + this.agents = agents != null ? new ArrayList<>(agents) : null; + return this; + } + + /** + * Sets the sub-agent names (required, at least one). + * + * @param agents the sub-agent names + * @return this builder + */ + public Builder agents(String... agents) { + this.agents = agents != null ? new ArrayList<>(Arrays.asList(agents)) : null; + return this; + } + + /** + * Adds sub-agents from {@link AgentRef}s (uses {@link AgentRef#getName()}). + * + * @param refs the agent refs + * @return this builder + */ + public Builder agentRefs(AgentRef... refs) { + List names = new ArrayList<>(); + if (refs != null) { + for (AgentRef ref : refs) { + if (ref != null) { + names.add(ref.getName()); + } + } + } + this.agents = names; + return this; + } + + /** + * Sets the tool-name prefix (default {@code "delegate_to"}). An empty string uses the bare + * agent name. + * + * @param toolPrefix the tool prefix + * @return this builder + */ + public Builder toolPrefix(String toolPrefix) { + this.toolPrefix = toolPrefix; + return this; + } + + /** + * Sets the maximum delegations per parent invocation ({@code 0} = unlimited). + * + * @param maxDelegations the max delegations + * @return this builder + */ + public Builder maxDelegations(int maxDelegations) { + this.maxDelegations = maxDelegations; + return this; + } + + /** + * Sets the trailing parent-history length to forward ({@code 0} = task only). + * + * @param historyLength the history length + * @return this builder + */ + public Builder historyLength(int historyLength) { + this.historyLength = historyLength; + return this; + } + + /** + * Sets the artifact strategy (default {@link ArtifactStrategy#INLINE}). + * + * @param artifactStrategy the artifact strategy + * @return this builder + */ + public Builder artifactStrategy(ArtifactStrategy artifactStrategy) { + this.artifactStrategy = artifactStrategy; + return this; + } + + /** + * Builds the {@link AgentsOptions}. + * + * @return a new {@link AgentsOptions} + * @throws IllegalArgumentException if no agents are configured + */ + public AgentsOptions build() { + return new AgentsOptions(this); + } + } +} diff --git a/plugins/middleware/src/main/java/com/google/genkit/plugins/middleware/ArtifactStrategy.java b/plugins/middleware/src/main/java/com/google/genkit/plugins/middleware/ArtifactStrategy.java new file mode 100644 index 000000000..c591c539d --- /dev/null +++ b/plugins/middleware/src/main/java/com/google/genkit/plugins/middleware/ArtifactStrategy.java @@ -0,0 +1,37 @@ +/* + * 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. + * + * SPDX-License-Identifier: Apache-2.0 + */ + +package com.google.genkit.plugins.middleware; + +/** + * Strategy controlling how artifacts produced by a sub-agent are surfaced back to the delegating + * model. + */ +public enum ArtifactStrategy { + /** + * Include each sub-agent artifact's text content inline in the delegation tool result (in + * addition to merging it into the parent session). This is the default. + */ + INLINE, + + /** + * Merge sub-agent artifacts into the parent session and return only their (namespaced) names in + * the delegation tool result, not their content. + */ + SESSION +} diff --git a/plugins/middleware/src/main/java/com/google/genkit/plugins/middleware/Artifacts.java b/plugins/middleware/src/main/java/com/google/genkit/plugins/middleware/Artifacts.java new file mode 100644 index 000000000..a0f7bcdfc --- /dev/null +++ b/plugins/middleware/src/main/java/com/google/genkit/plugins/middleware/Artifacts.java @@ -0,0 +1,222 @@ +/* + * 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. + * + * SPDX-License-Identifier: Apache-2.0 + */ + +package com.google.genkit.plugins.middleware; + +import com.fasterxml.jackson.annotation.JsonInclude; +import com.fasterxml.jackson.annotation.JsonProperty; +import com.google.genkit.ai.Part; +import com.google.genkit.ai.Tool; +import com.google.genkit.ai.agent.AgentSessionContext; +import com.google.genkit.ai.agent.Artifact; +import com.google.genkit.ai.agent.ArtifactStore; +import java.util.ArrayList; +import java.util.Collections; +import java.util.List; + +/** + * Factory for artifact tools that let a model read and write named artifacts on the active agent + * session. + * + *

Both tools operate on {@link AgentSessionContext#currentArtifactStore()}. When no agent + * session is bound to the current thread, the tools degrade gracefully: {@code read_artifact} + * reports {@code found = false} and {@code write_artifact} reports a {@code "no active session"} + * status without throwing. + * + *

Usage: + * + * 

{@code
+ * AgentConfig.builder()
+ *     .name("writer")
+ *     .tools(Artifacts.tools(ArtifactsOptions.defaults()))
+ *     .build();
+ * }
+ */ +public final class Artifacts { + + /** Name of the read tool. */ + public static final String READ_TOOL = "read_artifact"; + + /** Name of the write tool. */ + public static final String WRITE_TOOL = "write_artifact"; + + private Artifacts() {} + + /** + * Builds the artifact tools using default options ({@code read_artifact} + {@code + * write_artifact}). + * + * @return the artifact tools + */ + public static List> tools() { + return tools(ArtifactsOptions.defaults()); + } + + /** + * Builds the artifact tools. + * + *

Always includes {@code read_artifact}. Includes {@code write_artifact} unless {@code + * options.isReadonly()} is {@code true}. + * + * @param options the options (must not be null) + * @return the artifact tools + */ + public static List> tools(ArtifactsOptions options) { + if (options == null) { + throw new IllegalArgumentException("options must not be null"); + } + List> result = new ArrayList<>(); + result.add(readTool()); + if (!options.isReadonly()) { + result.add(writeTool()); + } + return result; + } + + // ── Tool builders ───────────────────────────────────────────────────────────── + + private static Tool readTool() { + return Tool.builder() + .name(READ_TOOL) + .description( + "Read a named artifact from the current session. Returns its text content and whether " + + "it was found.") + .inputClass(ReadInput.class) + .outputClass(ReadOutput.class) + .handler( + (ctx, in) -> { + String name = in != null ? in.name : null; + ArtifactStore store = AgentSessionContext.currentArtifactStore(); + if (store == null || name == null) { + return new ReadOutput(name, null, false); + } + for (Artifact a : store.getArtifacts()) { + if (name.equals(a.getName())) { + return new ReadOutput(name, textOf(a), true); + } + } + return new ReadOutput(name, null, false); + }) + .build(); + } + + private static Tool writeTool() { + return Tool.builder() + .name(WRITE_TOOL) + .description( + "Write (create or replace) a named artifact in the current session with the given text " + + "content.") + .inputClass(WriteInput.class) + .outputClass(WriteOutput.class) + .handler( + (ctx, in) -> { + String name = in != null ? in.name : null; + String content = in != null ? in.content : null; + ArtifactStore store = AgentSessionContext.currentArtifactStore(); + if (store == null) { + return new WriteOutput("no active session"); + } + if (name == null || name.isEmpty()) { + return new WriteOutput("error: artifact name is required"); + } + Artifact artifact = + Artifact.builder() + .name(name) + .parts(Collections.singletonList(Part.text(content != null ? content : ""))) + .build(); + // addArtifacts deduplicates by name, replacing any existing artifact with this name. + store.addArtifacts(artifact); + return new WriteOutput("ok"); + }) + .build(); + } + + /** Concatenates the text of all text parts of an artifact. */ + static String textOf(Artifact artifact) { + StringBuilder sb = new StringBuilder(); + if (artifact.getParts() != null) { + for (Part p : artifact.getParts()) { + if (p.getText() != null) { + sb.append(p.getText()); + } + } + } + return sb.toString(); + } + + // ── Tool I/O types ────────────────────────────────────────────────────────────── + + /** Input for {@code read_artifact}. */ + @JsonInclude(JsonInclude.Include.NON_NULL) + public static final class ReadInput { + @JsonProperty("name") + public String name; + + /** Default constructor for JSON deserialization. */ + public ReadInput() {} + } + + /** Output for {@code read_artifact}. */ + @JsonInclude(JsonInclude.Include.NON_NULL) + public static final class ReadOutput { + @JsonProperty("name") + public String name; + + @JsonProperty("content") + public String content; + + @JsonProperty("found") + public boolean found; + + /** Default constructor for JSON deserialization. */ + public ReadOutput() {} + + ReadOutput(String name, String content, boolean found) { + this.name = name; + this.content = content; + this.found = found; + } + } + + /** Input for {@code write_artifact}. */ + @JsonInclude(JsonInclude.Include.NON_NULL) + public static final class WriteInput { + @JsonProperty("name") + public String name; + + @JsonProperty("content") + public String content; + + /** Default constructor for JSON deserialization. */ + public WriteInput() {} + } + + /** Output for {@code write_artifact}. */ + @JsonInclude(JsonInclude.Include.NON_NULL) + public static final class WriteOutput { + @JsonProperty("status") + public String status; + + /** Default constructor for JSON deserialization. */ + public WriteOutput() {} + + WriteOutput(String status) { + this.status = status; + } + } +} diff --git a/plugins/middleware/src/main/java/com/google/genkit/plugins/middleware/ArtifactsOptions.java b/plugins/middleware/src/main/java/com/google/genkit/plugins/middleware/ArtifactsOptions.java new file mode 100644 index 000000000..8a3982d54 --- /dev/null +++ b/plugins/middleware/src/main/java/com/google/genkit/plugins/middleware/ArtifactsOptions.java @@ -0,0 +1,89 @@ +/* + * 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. + * + * SPDX-License-Identifier: Apache-2.0 + */ + +package com.google.genkit.plugins.middleware; + +/** + * Options controlling the artifact tools produced by {@link Artifacts}. + * + *

The only knob is {@link #isReadonly()}: when {@code true}, only the {@code read_artifact} tool + * is produced; when {@code false} (the default) a {@code write_artifact} tool is produced as well. + */ +public final class ArtifactsOptions { + + private final boolean readonly; + + private ArtifactsOptions(Builder builder) { + this.readonly = builder.readonly; + } + + /** + * Returns a default options instance ({@code readonly = false}). + * + * @return default options + */ + public static ArtifactsOptions defaults() { + return builder().build(); + } + + /** + * Creates a builder for {@link ArtifactsOptions}. + * + * @return a new builder + */ + public static Builder builder() { + return new Builder(); + } + + /** + * Returns whether artifact writing is disabled. + * + * @return {@code true} if only the read tool should be produced + */ + public boolean isReadonly() { + return readonly; + } + + /** Builder for {@link ArtifactsOptions}. */ + public static final class Builder { + private boolean readonly; + + private Builder() {} + + /** + * Sets whether artifact writing is disabled. When {@code true}, {@link Artifacts#tools} omits + * the {@code write_artifact} tool. + * + * @param readonly {@code true} to produce only the read tool + * @return this builder + */ + public Builder readonly(boolean readonly) { + this.readonly = readonly; + return this; + } + + /** + * Builds the {@link ArtifactsOptions}. + * + * @return a new {@link ArtifactsOptions} + */ + public ArtifactsOptions build() { + return new ArtifactsOptions(this); + } + } +} diff --git a/plugins/middleware/src/main/java/com/google/genkit/plugins/middleware/internal/Delegation.java b/plugins/middleware/src/main/java/com/google/genkit/plugins/middleware/internal/Delegation.java new file mode 100644 index 000000000..7a2a42b6d --- /dev/null +++ b/plugins/middleware/src/main/java/com/google/genkit/plugins/middleware/internal/Delegation.java @@ -0,0 +1,356 @@ +/* + * 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. + * + * SPDX-License-Identifier: Apache-2.0 + */ + +package com.google.genkit.plugins.middleware.internal; + +import com.fasterxml.jackson.databind.JsonNode; +import com.google.genkit.ai.Message; +import com.google.genkit.ai.Part; +import com.google.genkit.ai.ToolInterruptException; +import com.google.genkit.ai.agent.Agent; +import com.google.genkit.ai.agent.AgentFinishReason; +import com.google.genkit.ai.agent.AgentInit; +import com.google.genkit.ai.agent.AgentInput; +import com.google.genkit.ai.agent.AgentOutput; +import com.google.genkit.ai.agent.AgentSessionContext; +import com.google.genkit.ai.agent.Artifact; +import com.google.genkit.ai.agent.ArtifactStore; +import com.google.genkit.ai.agent.Session; +import com.google.genkit.ai.agent.SessionState; +import com.google.genkit.core.Action; +import com.google.genkit.core.ActionContext; +import com.google.genkit.core.ActionType; +import com.google.genkit.core.BidiAction; +import com.google.genkit.core.BufferedInputSource; +import com.google.genkit.core.GenkitException; +import com.google.genkit.core.JsonUtils; +import com.google.genkit.core.Registry; +import java.util.ArrayList; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.UUID; +import java.util.concurrent.ConcurrentHashMap; + +/** + * Internal helper that runs a single one-shot turn against a sub-agent and packages the result. + * + *

This isolates the registry lookup, the {@code runBidiJson} one-shot invocation, output + * parsing, artifact namespacing/merging, and cross-call history bookkeeping from the public {@code + * Agents} façade. + * + *

History forwarding (per parent-session, per-sub-agent)

+ * + *

Each (parent session id, sub-agent name) pair gets its own stable sub-session key ({@link + * #subSessionKey}). For a client-managed sub-agent (no {@code SessionStore}), this class + * keeps a small in-memory ledger ({@link #CLIENT_HISTORY}) of that sub-session's accumulated + * messages and forwards the trailing {@code historyLength} of them via {@code + * AgentInit.state.messages} on every call, then updates the ledger from the returned {@code + * AgentOutput.state} afterward. This leans on the existing client-managed hydration path in {@code + * SessionResolver.resolve} (state → {@code new Session<>(state)}) rather than inventing new + * persistence. + * + *

For a server-managed sub-agent (has a {@code SessionStore}), the sub-session key is + * forwarded as {@code AgentInit.sessionId}; the sub-agent's own store then naturally accumulates + * history via {@code SessionResolver.resolveBySessionId} across repeated delegation calls. Note: + * {@code historyLength} trimming is not applied in the server-managed case — the store + * resolves and hands the agent its full accumulated history, and there is no seam in {@code + * defineCustomAgent} to truncate what a resolved session sees before {@code AgentFn} runs without + * mutating the sub-agent's own persisted snapshot (out of scope for this fix; see {@code + * Agents}/{@code AgentsOptions} javadoc). + */ +public final class Delegation { + + private Delegation() {} + + /** + * Ledger of accumulated message history per sub-session key, for client-managed sub-agents only. + * Server-managed sub-agents accumulate history in their own {@code SessionStore} instead. + */ + private static final Map> CLIENT_HISTORY = new ConcurrentHashMap<>(); + + /** Result of a single delegation run, ready to be mapped onto the delegation tool output. */ + public static final class Result { + public final String response; + public final List artifacts; + + Result(String response, List artifacts) { + this.response = response; + this.artifacts = artifacts; + } + } + + /** A sub-agent artifact after namespacing, ready for the tool output. */ + public static final class NamedArtifact { + public final String name; + public final String content; + + NamedArtifact(String name, String content) { + this.name = name; + this.content = content; + } + } + + /** + * Runs the sub-agent named {@code agentName} for a single turn with {@code task} as the user + * message, merges any produced artifacts (namespaced by a fresh invocation id) into the active + * parent session, and returns the sub-agent's text plus the (namespaced) artifacts. + * + *

Up to {@code historyLength} prior messages from this (parent session, sub-agent) pair's own + * accumulated conversation are forwarded to the sub-agent (see class Javadoc for the + * client-managed vs. server-managed mechanics). A {@code historyLength <= 0} means no history is + * forwarded (task-only), matching the pre-fix behavior. + * + * @param ctx the calling tool's action context (provides the registry) + * @param agentName the bare sub-agent name (looked up at {@code /agent/}) + * @param task the task text to send as the sub-agent's user message + * @param includeArtifactContent whether to include artifact content in the returned artifacts + * ({@link com.google.genkit.plugins.middleware.ArtifactStrategy#INLINE}) or names only + * ({@link com.google.genkit.plugins.middleware.ArtifactStrategy#SESSION}) + * @param historyLength the number of trailing prior messages to forward to the sub-agent ({@code + * <= 0} means task-only, no history forwarded) + * @return the delegation result + * @throws ToolInterruptException if the sub-agent's turn finished {@code INTERRUPTED}; carries + * the sub-agent's interrupted tool name/input as metadata so the parent's own generate loop + * pauses too + * @throws GenkitException if the sub-agent's turn finished {@code FAILED}, or the sub-agent's + * {@code runBidiJson} call itself threw + */ + public static Result run( + ActionContext ctx, + String agentName, + String task, + boolean includeArtifactContent, + int historyLength) { + Registry registry = ctx.getRegistry(); + Action action = registry.lookupAction(ActionType.AGENT.keyFromName(agentName)); + if (action == null) { + return new Result( + "Error: sub-agent '" + agentName + "' is not registered.", new ArrayList<>()); + } + if (!(action instanceof BidiAction)) { + return new Result( + "Error: '" + agentName + "' is not a bidirectional agent action.", new ArrayList<>()); + } + BidiAction agent = (BidiAction) action; + + boolean serverManaged = isServerManaged(agent); + String subSessionKey = subSessionKey(agentName); + + // Build a one-shot input source carrying the task as a single user turn. + AgentInput input = AgentInput.builder().message(Message.user(task != null ? task : "")).build(); + BufferedInputSource inputs = new BufferedInputSource<>(); + inputs.offer(JsonUtils.toJsonNode(input)); + inputs.end(); + + AgentInit init = new AgentInit<>(); + List forwardedHistory = null; + if (serverManaged) { + // Server-managed: resume the sub-agent's own persistent sub-session by sessionId so its + // store naturally accumulates history across repeated delegation calls (see class Javadoc). + init.setSessionId(subSessionKey); + } else if (historyLength > 0) { + // Client-managed: hydrate from our own ledger, trimmed to historyLength. + forwardedHistory = trimTrailing(CLIENT_HISTORY.get(subSessionKey), historyLength); + if (!forwardedHistory.isEmpty()) { + SessionState state = new SessionState<>(); + state.setMessages(forwardedHistory); + init.setState(state); + } + } + JsonNode initJson = JsonUtils.toJsonNode(init); + + JsonNode outJson; + try { + outJson = agent.runBidiJson(ctx, initJson, inputs, chunk -> {}); + } catch (Exception e) { + throw new GenkitException("Sub-agent '" + agentName + "' failed: " + e.getMessage(), e); + } + + AgentOutput output = JsonUtils.fromJsonNode(outJson, AgentOutput.class); + + // Update the client-managed history ledger with this turn's full message history (task + + // response), trimmed to historyLength, ready for the next delegation call to this sub-agent. + if (!serverManaged && historyLength > 0) { + List updated = + output != null && output.getState() != null ? output.getState().getMessages() : null; + if (updated == null) { + // Fall back to reconstructing task + response if the sub-agent didn't return state. + updated = new ArrayList<>(); + if (forwardedHistory != null) { + updated.addAll(forwardedHistory); + } + updated.add(input.getMessage()); + if (output != null && output.getMessage() != null) { + updated.add(output.getMessage()); + } + } + CLIENT_HISTORY.put(subSessionKey, trimTrailing(updated, historyLength)); + } + + throwIfInterruptedOrFailed(agentName, output); + + String response = extractResponse(output); + + // Namespace + merge artifacts into the parent session. + List namespaced = new ArrayList<>(); + if (output != null && output.getArtifacts() != null && !output.getArtifacts().isEmpty()) { + String invocationId = UUID.randomUUID().toString(); + ArtifactStore parentStore = AgentSessionContext.currentArtifactStore(); + for (Artifact a : output.getArtifacts()) { + String baseName = a.getName() != null ? a.getName() : "artifact"; + String namespacedName = invocationId + "/" + baseName; + Artifact merged = + Artifact.builder() + .name(namespacedName) + .parts(a.getParts()) + .metadata(a.getMetadata()) + .build(); + if (parentStore != null) { + parentStore.addArtifacts(merged); + } + namespaced.add( + new NamedArtifact(namespacedName, includeArtifactContent ? textOf(a) : null)); + } + } + + return new Result(response, namespaced); + } + + /** + * Derives the stable sub-session key for a (parent session, sub-agent) pair: the currently bound + * parent session id (or a fixed fallback if no {@link AgentSessionContext} is bound) combined + * with the sub-agent's name. Repeated delegation to the SAME sub-agent within the SAME parent + * session resumes the same sub-session; a different parent session or a different sub-agent name + * gets an independent one. + */ + private static String subSessionKey(String agentName) { + Session parent = AgentSessionContext.current(); + String parentSessionId = parent != null ? parent.sessionId() : "no-parent-session"; + return parentSessionId + "::" + agentName; + } + + /** + * Returns {@code true} if the resolved agent action reports server-managed state. {@link Agent} + * is the only implementation produced by {@code AgentActions.defineCustomAgent} and exposes + * {@link Agent#serverManaged()} directly. + */ + private static boolean isServerManaged(BidiAction agent) { + return agent instanceof Agent && ((Agent) agent).serverManaged(); + } + + /** Returns a new list containing at most the last {@code n} elements of {@code list}. */ + private static List trimTrailing(List list, int n) { + if (list == null || list.isEmpty() || n <= 0) { + return new ArrayList<>(); + } + if (list.size() <= n) { + return new ArrayList<>(list); + } + return new ArrayList<>(list.subList(list.size() - n, list.size())); + } + + /** + * Throws a structured exception for {@code INTERRUPTED} and {@code FAILED} sub-agent turns so + * they propagate through the parent's own generate/tool-calling loop instead of collapsing to + * plain text. + * + *

INTERRUPTED: throws {@link ToolInterruptException} carrying the sub-agent's + * interrupted tool name/input (extracted from the final message's tool-request part) as metadata, + * mirroring the convention used by {@code Tool.run}'s catch/rethrow of the same exception type. + * This makes the PARENT's generate loop stop with {@code FinishReason.INTERRUPTED} too, exactly + * as if the parent had called an interrupting tool directly. Scope decision: this fix surfaces + * the interrupt structurally to the parent; it does not implement an end-to-end nested-resume + * path (resuming the parent's interrupt does not automatically re-invoke the sub-agent with a + * resume) — see {@code Agents}/{@code AgentsOptions} javadoc and the fix-middleware report for + * the full rationale. + * + *

FAILED: throws {@link GenkitException} carrying the sub-agent's actual error message, + * matching how {@code Tool.run} wraps a handler's checked failure and how {@code Genkit.generate} + * treats any other tool exception (folded into a {@code ToolResponse} error, not propagated as a + * structured interrupt). + */ + private static void throwIfInterruptedOrFailed(String agentName, AgentOutput output) { + if (output == null) { + return; + } + AgentFinishReason reason = output.getFinishReason(); + if (reason == AgentFinishReason.INTERRUPTED) { + Map metadata = new HashMap<>(); + Part toolRequestPart = findToolRequestPart(output.getMessage()); + if (toolRequestPart != null && toolRequestPart.getToolRequest() != null) { + metadata.put("name", toolRequestPart.getToolRequest().getName()); + metadata.put("input", toolRequestPart.getToolRequest().getInput()); + } + metadata.put("subAgent", agentName); + String text = output.getMessage() != null ? output.getMessage().getText() : null; + String message = + "Sub-agent '" + + agentName + + "' was interrupted before completing" + + (text != null && !text.isEmpty() ? ": " + text : "."); + throw new ToolInterruptException(message, metadata); + } + if (reason == AgentFinishReason.FAILED) { + String err = + output.getError() != null ? String.valueOf(output.getError().getMessage()) : null; + String text = output.getMessage() != null ? output.getMessage().getText() : null; + String detail = err != null ? err : text; + throw new GenkitException( + "Sub-agent '" + + agentName + + "' failed" + + (detail != null && !detail.isEmpty() ? ": " + detail : ".")); + } + } + + /** Finds the first tool-request part in a message's content, or {@code null} if none. */ + private static Part findToolRequestPart(Message message) { + if (message == null || message.getContent() == null) { + return null; + } + for (Part p : message.getContent()) { + if (p.getToolRequest() != null) { + return p; + } + } + return null; + } + + /** Builds the textual response for a normal (STOP/etc.) sub-agent turn. */ + private static String extractResponse(AgentOutput output) { + if (output == null) { + return "Sub-agent produced no output."; + } + String text = output.getMessage() != null ? output.getMessage().getText() : null; + return text != null ? text : ""; + } + + /** Concatenates the text of all text parts of an artifact. */ + private static String textOf(Artifact artifact) { + StringBuilder sb = new StringBuilder(); + if (artifact.getParts() != null) { + for (Part p : artifact.getParts()) { + if (p.getText() != null) { + sb.append(p.getText()); + } + } + } + return sb.toString(); + } +} diff --git a/plugins/middleware/src/main/java/com/google/genkit/plugins/middleware/package-info.java b/plugins/middleware/src/main/java/com/google/genkit/plugins/middleware/package-info.java new file mode 100644 index 000000000..5e8d28ba2 --- /dev/null +++ b/plugins/middleware/src/main/java/com/google/genkit/plugins/middleware/package-info.java @@ -0,0 +1,40 @@ +/* + * 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. + * + * SPDX-License-Identifier: Apache-2.0 + */ + +/** + * Middleware plugin for Genkit providing higher-level generation building blocks. + * + *

This plugin provides: + * + *

    + *
  • {@link com.google.genkit.plugins.middleware.Agents} — sub-agent delegation, where + * each configured sub-agent is exposed to the model as a {@code delegate_to_} tool that + * runs the sub-agent for a single turn and returns its text (plus optional artifacts). + *
  • {@link com.google.genkit.plugins.middleware.Artifacts} — {@code read_artifact} and + * {@code write_artifact} tools operating on the active agent session's artifact store. + *
+ * + *

Both are implemented as tool factories: they return {@code List>} (and, + * for {@code agents()}, a system-prompt fragment) that callers wire into an agent via {@code + * AgentConfig.tools(...)} and {@code AgentConfig.system(...)}. This avoids modifying the generate + * pipeline while still letting a model delegate to sub-agents and read/write artifacts. + * + * @see com.google.genkit.plugins.middleware.Agents + * @see com.google.genkit.plugins.middleware.Artifacts + */ +package com.google.genkit.plugins.middleware; diff --git a/plugins/middleware/src/test/java/com/google/genkit/plugins/middleware/AgentsTest.java b/plugins/middleware/src/test/java/com/google/genkit/plugins/middleware/AgentsTest.java new file mode 100644 index 000000000..e13b6cfbe --- /dev/null +++ b/plugins/middleware/src/test/java/com/google/genkit/plugins/middleware/AgentsTest.java @@ -0,0 +1,551 @@ +/* + * 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. + * + * SPDX-License-Identifier: Apache-2.0 + */ + +package com.google.genkit.plugins.middleware; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import com.google.genkit.ai.Message; +import com.google.genkit.ai.Part; +import com.google.genkit.ai.agent.AgentFinishReason; +import com.google.genkit.ai.agent.AgentFn; +import com.google.genkit.ai.agent.AgentResult; +import com.google.genkit.ai.agent.AgentSessionContext; +import com.google.genkit.ai.agent.Artifact; +import com.google.genkit.ai.agent.CustomAgentConfig; +import com.google.genkit.ai.agent.InMemorySessionStore; +import com.google.genkit.ai.agent.Session; +import com.google.genkit.ai.agent.SessionState; +import com.google.genkit.ai.agent.internal.AgentActions; +import com.google.genkit.core.ActionContext; +import com.google.genkit.core.DefaultRegistry; +import com.google.genkit.core.Registry; +import java.util.Collections; +import java.util.List; +import java.util.Map; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; + +/** TDD tests for {@link Agents} sub-agent delegation (Stage 6). */ +class AgentsTest { + + private Registry registry; + private ActionContext ctx; + + @BeforeEach + void setUp() { + registry = new DefaultRegistry(); + ctx = new ActionContext(registry); + } + + // ── helpers ────────────────────────────────────────────────────────────────── + + /** Registers a client-managed sub-agent whose turn returns a fixed assistant message. */ + private void defineFixedAgent(String name, String responseText) { + AgentFn> fn = + (sess, fnCtx) -> + AgentResult.builder() + .message(Message.model(responseText)) + .finishReason(AgentFinishReason.STOP) + .build(); + CustomAgentConfig> config = + CustomAgentConfig.>builder().name(name).build(); + AgentActions.defineCustomAgent(registry, config, fn); + } + + /** Registers a sub-agent that also produces a named artifact. */ + private void defineArtifactAgent(String name, String responseText, String artifactContent) { + AgentFn> fn = + (sess, fnCtx) -> + AgentResult.builder() + .message(Message.model(responseText)) + .artifacts( + Collections.singletonList( + Artifact.builder() + .name("report") + .parts(Collections.singletonList(Part.text(artifactContent))) + .build())) + .finishReason(AgentFinishReason.STOP) + .build(); + CustomAgentConfig> config = + CustomAgentConfig.>builder().name(name).build(); + AgentActions.defineCustomAgent(registry, config, fn); + } + + @SuppressWarnings("unchecked") + private static com.google.genkit.ai.Tool tool( + List> tools, String name) { + return (com.google.genkit.ai.Tool) + tools.stream().filter(t -> name.equals(t.getName())).findFirst().orElseThrow(); + } + + private static Session> newSession() { + return new Session<>(SessionState.>builder().build()); + } + + // ── tests ────────────────────────────────────────────────────────────────────── + + @Test + void delegationToolNamesUsePrefix() { + defineFixedAgent("researcher", "result"); + List> tools = + Agents.delegationTools(AgentsOptions.builder().agents("researcher").build()); + assertEquals(1, tools.size()); + assertEquals("delegate_to_researcher", tools.get(0).getName()); + } + + @Test + void emptyPrefixUsesBareName() { + defineFixedAgent("researcher", "result"); + List> tools = + Agents.delegationTools(AgentsOptions.builder().agents("researcher").toolPrefix("").build()); + assertEquals("researcher", tools.get(0).getName()); + } + + @Test + void delegateRunsSubAgentAndReturnsText() { + defineFixedAgent("researcher", "I found X in the archives."); + List> tools = + Agents.delegationTools(AgentsOptions.builder().agents("researcher").build()); + com.google.genkit.ai.Tool delegate = + tool(tools, "delegate_to_researcher"); + + Agents.DelegateInput in = new Agents.DelegateInput(); + in.task = "find X"; + Agents.DelegateOutput out = delegate.run(ctx, in); + + assertEquals("I found X in the archives.", out.response); + } + + @Test + void maxDelegationsCapReturnsLimitMessage() { + defineFixedAgent("researcher", "ok"); + List> tools = + Agents.delegationTools( + AgentsOptions.builder().agents("researcher").maxDelegations(1).build()); + com.google.genkit.ai.Tool delegate = + tool(tools, "delegate_to_researcher"); + + Agents.DelegateInput in = new Agents.DelegateInput(); + in.task = "t"; + + Agents.DelegateOutput first = delegate.run(ctx, in); + assertEquals("ok", first.response); + + Agents.DelegateOutput second = delegate.run(ctx, in); + assertTrue( + second.response.toLowerCase().contains("delegation limit reached"), + "expected limit message, got: " + second.response); + } + + @Test + void unknownAgentReturnsErrorText() { + // No agent registered under this name. + List> tools = + Agents.delegationTools(AgentsOptions.builder().agents("ghost").build()); + com.google.genkit.ai.Tool delegate = + tool(tools, "delegate_to_ghost"); + + Agents.DelegateInput in = new Agents.DelegateInput(); + in.task = "t"; + Agents.DelegateOutput out = delegate.run(ctx, in); + assertTrue(out.response.toLowerCase().contains("not registered"), out.response); + } + + @Test + void inlineArtifactsAreNamespacedMergedAndContentIncluded() { + defineArtifactAgent("researcher", "done", "the body of the report"); + List> tools = + Agents.delegationTools( + AgentsOptions.builder() + .agents("researcher") + .artifactStrategy(ArtifactStrategy.INLINE) + .build()); + com.google.genkit.ai.Tool delegate = + tool(tools, "delegate_to_researcher"); + + Session> parent = newSession(); + AgentSessionContext.run( + parent, + () -> { + Agents.DelegateInput in = new Agents.DelegateInput(); + in.task = "write report"; + Agents.DelegateOutput out = delegate.run(ctx, in); + + assertEquals("done", out.response); + assertNotNull(out.artifacts); + assertEquals(1, out.artifacts.size()); + // namespaced as /report + assertTrue(out.artifacts.get(0).name.endsWith("/report"), out.artifacts.get(0).name); + // INLINE includes content + assertEquals("the body of the report", out.artifacts.get(0).content); + }); + + // merged into the parent session, namespaced + assertEquals(1, parent.getArtifacts().size()); + assertTrue(parent.getArtifacts().get(0).getName().endsWith("/report")); + } + + @Test + void sessionStrategyMergesButOmitsContent() { + defineArtifactAgent("researcher", "done", "secret body"); + List> tools = + Agents.delegationTools( + AgentsOptions.builder() + .agents("researcher") + .artifactStrategy(ArtifactStrategy.SESSION) + .build()); + com.google.genkit.ai.Tool delegate = + tool(tools, "delegate_to_researcher"); + + Session> parent = newSession(); + AgentSessionContext.run( + parent, + () -> { + Agents.DelegateInput in = new Agents.DelegateInput(); + in.task = "write report"; + Agents.DelegateOutput out = delegate.run(ctx, in); + assertEquals(1, out.artifacts.size()); + // SESSION: name present, content omitted + assertTrue(out.artifacts.get(0).name.endsWith("/report")); + org.junit.jupiter.api.Assertions.assertNull(out.artifacts.get(0).content); + }); + + assertEquals(1, parent.getArtifacts().size()); + } + + @Test + void systemPromptFragmentListsTools() { + String fragment = + Agents.systemPromptFragment(AgentsOptions.builder().agents("researcher", "writer").build()); + assertTrue(fragment.contains("")); + assertTrue(fragment.contains("delegate_to_researcher")); + assertTrue(fragment.contains("delegate_to_writer")); + } + + @Test + void historyLengthOptionForwardsPriorDelegationContext() { + // Fixed behavior (see Agents/AgentsOptions javadoc and internal.Delegation): repeated + // delegation to the SAME sub-agent within the SAME parent session now resumes that + // (parent-session, sub-agent) pair's own accumulated sub-session, trimmed to historyLength. + // The sub-agent is client-managed (no store), so history is forwarded via the internal + // client-history ledger + AgentInit.state. Prove real forwarding by having the sub-agent + // report exactly how many messages it can see on each call: call 1 sees only its own task (1); + // call 2 sees call 1's task + response + its own task (3); call 3 sees all six messages so far + // (call 1's task+response, call 2's task+response, its own task) = 5, still under the + // historyLength=10 cap, so nothing is trimmed yet. + AgentFn> fn = + (sess, fnCtx) -> + AgentResult.builder() + .message(Message.model("saw " + sess.getMessages().size() + " message(s)")) + .finishReason(AgentFinishReason.STOP) + .build(); + CustomAgentConfig> config = + CustomAgentConfig.>builder().name("researcher").build(); + AgentActions.defineCustomAgent(registry, config, fn); + + List> tools = + Agents.delegationTools( + AgentsOptions.builder().agents("researcher").historyLength(10).build()); + com.google.genkit.ai.Tool delegate = + tool(tools, "delegate_to_researcher"); + + // Bind a parent session so the derived sub-session key is stable and isolated from other + // tests running against the same "researcher" agent name. + Session> parent = newSession(); + AgentSessionContext.run( + parent, + () -> { + Agents.DelegateInput in = new Agents.DelegateInput(); + + in.task = "task 1"; + assertEquals( + "saw 1 message(s)", + delegate.run(ctx, in).response, + "first delegation call has no prior history: sees only its own task"); + + in.task = "task 2"; + assertEquals( + "saw 3 message(s)", + delegate.run(ctx, in).response, + "second call should see call 1's task+response plus its own task"); + + in.task = "task 3"; + assertEquals( + "saw 5 message(s)", + delegate.run(ctx, in).response, + "third call should see calls 1 and 2's task+response plus its own task"); + }); + } + + @Test + void historyLengthOptionTrimsToConfiguredLength() { + // historyLength caps how much prior context is forwarded: with historyLength=2, the sub-agent + // should only ever see its own task plus at most 2 prior messages, never the full backlog. + AgentFn> fn = + (sess, fnCtx) -> + AgentResult.builder() + .message(Message.model("saw " + sess.getMessages().size() + " message(s)")) + .finishReason(AgentFinishReason.STOP) + .build(); + CustomAgentConfig> config = + CustomAgentConfig.>builder().name("trimmer").build(); + AgentActions.defineCustomAgent(registry, config, fn); + + List> tools = + Agents.delegationTools(AgentsOptions.builder().agents("trimmer").historyLength(2).build()); + com.google.genkit.ai.Tool delegate = + tool(tools, "delegate_to_trimmer"); + + Session> parent = newSession(); + AgentSessionContext.run( + parent, + () -> { + Agents.DelegateInput in = new Agents.DelegateInput(); + + in.task = "task 1"; + assertEquals("saw 1 message(s)", delegate.run(ctx, in).response); + + in.task = "task 2"; + // Prior history (task 1 + response = 2 messages) trimmed to historyLength=2, plus the + // new task message = 3. + assertEquals("saw 3 message(s)", delegate.run(ctx, in).response); + + in.task = "task 3"; + // Accumulated history is now 4 messages (task1, resp1, task2, resp2); trimmed to the + // trailing 2, plus the new task message = 3 (never grows past the cap + 1). + assertEquals( + "saw 3 message(s)", + delegate.run(ctx, in).response, + "history forwarded to the sub-agent must stay capped at historyLength, not grow" + + " unbounded"); + }); + } + + @Test + void historyIsIsolatedPerParentSessionAndPerSubAgent() { + // A different parent session delegating to the SAME sub-agent name must get an independent + // sub-session (no bleed-over from another parent's conversation). + AgentFn> fn = + (sess, fnCtx) -> + AgentResult.builder() + .message(Message.model("saw " + sess.getMessages().size() + " message(s)")) + .finishReason(AgentFinishReason.STOP) + .build(); + CustomAgentConfig> config = + CustomAgentConfig.>builder().name("isolated").build(); + AgentActions.defineCustomAgent(registry, config, fn); + + List> tools = + Agents.delegationTools( + AgentsOptions.builder().agents("isolated").historyLength(10).build()); + com.google.genkit.ai.Tool delegate = + tool(tools, "delegate_to_isolated"); + + Session> parentA = newSession(); + AgentSessionContext.run( + parentA, + () -> { + Agents.DelegateInput in = new Agents.DelegateInput(); + in.task = "A task 1"; + assertEquals("saw 1 message(s)", delegate.run(ctx, in).response); + in.task = "A task 2"; + assertEquals("saw 3 message(s)", delegate.run(ctx, in).response); + }); + + Session> parentB = newSession(); + AgentSessionContext.run( + parentB, + () -> { + Agents.DelegateInput in = new Agents.DelegateInput(); + in.task = "B task 1"; + assertEquals( + "saw 1 message(s)", + delegate.run(ctx, in).response, + "a different parent session must not see parent A's history for the same sub-agent"); + }); + } + + @Test + void interruptedSubAgentThrowsStructuredToolInterrupt() { + // Fixed behavior (see Agents/internal.Delegation javadoc): finishReason==INTERRUPTED now + // causes the delegation tool to throw ToolInterruptException (the same exception type/ + // convention Tool.run re-throws for human-in-the-loop), carrying the sub-agent's interrupted + // tool name/input as metadata, instead of collapsing to plain text. This lets the PARENT's own + // generate/tool-calling loop pause too (Genkit.java's tool-execution loop catches + // ToolInterruptException specifically and surfaces FinishReason.INTERRUPTED). + AgentFn> fn = + (sess, fnCtx) -> { + com.google.genkit.ai.ToolRequest toolRequest = + new com.google.genkit.ai.ToolRequest( + "confirmAction", Map.of("action", "do the risky thing")); + com.google.genkit.ai.Part toolRequestPart = new com.google.genkit.ai.Part(); + toolRequestPart.setToolRequest(toolRequest); + Message interrupted = + new Message( + com.google.genkit.ai.Role.MODEL, + List.of( + com.google.genkit.ai.Part.text("need approval to continue"), + toolRequestPart)); + return AgentResult.builder() + .message(interrupted) + .finishReason(AgentFinishReason.INTERRUPTED) + .build(); + }; + CustomAgentConfig> config = + CustomAgentConfig.>builder().name("approver").build(); + AgentActions.defineCustomAgent(registry, config, fn); + + List> tools = + Agents.delegationTools(AgentsOptions.builder().agents("approver").build()); + com.google.genkit.ai.Tool delegate = + tool(tools, "delegate_to_approver"); + + Agents.DelegateInput in = new Agents.DelegateInput(); + in.task = "do the risky thing"; + + com.google.genkit.ai.ToolInterruptException thrown = + org.junit.jupiter.api.Assertions.assertThrows( + com.google.genkit.ai.ToolInterruptException.class, () -> delegate.run(ctx, in)); + + assertTrue( + thrown.getMessage().toLowerCase().contains("interrupted"), + "expected a descriptive interrupted message, got: " + thrown.getMessage()); + assertTrue( + thrown.getMessage().contains("need approval to continue"), + "expected the sub-agent's message text folded into the exception message, got: " + + thrown.getMessage()); + assertEquals( + "confirmAction", + thrown.getMetadata().get("name"), + "expected the sub-agent's interrupted tool name to be carried as metadata"); + assertEquals( + Map.of("action", "do the risky thing"), + thrown.getMetadata().get("input"), + "expected the sub-agent's interrupted tool input to be carried as metadata"); + } + + @Test + void interruptedSubAgentWithNoToolRequestStillThrowsInterrupt() { + // Even if the sub-agent's INTERRUPTED message carries no tool-request part (e.g. a + // hand-rolled AgentFn that sets finishReason without populating toolRequests), the delegation + // tool must still throw ToolInterruptException rather than falling back to descriptive text. + AgentFn> fn = + (sess, fnCtx) -> + AgentResult.builder() + .message(Message.model("need approval to continue")) + .finishReason(AgentFinishReason.INTERRUPTED) + .build(); + CustomAgentConfig> config = + CustomAgentConfig.>builder().name("bareapprover").build(); + AgentActions.defineCustomAgent(registry, config, fn); + + List> tools = + Agents.delegationTools(AgentsOptions.builder().agents("bareapprover").build()); + com.google.genkit.ai.Tool delegate = + tool(tools, "delegate_to_bareapprover"); + + Agents.DelegateInput in = new Agents.DelegateInput(); + in.task = "do the risky thing"; + + com.google.genkit.ai.ToolInterruptException thrown = + org.junit.jupiter.api.Assertions.assertThrows( + com.google.genkit.ai.ToolInterruptException.class, () -> delegate.run(ctx, in)); + assertTrue(thrown.getMessage().contains("need approval to continue")); + } + + @Test + void failingSubAgentThrowsGenkitException() { + // Fixed behavior (see Agents/internal.Delegation javadoc): finishReason==FAILED now causes the + // delegation tool to throw a GenkitException carrying the sub-agent's actual error message, + // instead of collapsing to plain text. This propagates like any other tool exception in + // Genkit's normal tool-calling loop (Genkit.java's generic catch folds it into a ToolResponse + // error rather than a structured interrupt). + AgentFn> fn = + (sess, fnCtx) -> { + throw new RuntimeException("boom: downstream service unavailable"); + }; + CustomAgentConfig> config = + CustomAgentConfig.>builder().name("flaky").build(); + AgentActions.defineCustomAgent(registry, config, fn); + + List> tools = + Agents.delegationTools(AgentsOptions.builder().agents("flaky").build()); + com.google.genkit.ai.Tool delegate = + tool(tools, "delegate_to_flaky"); + + Agents.DelegateInput in = new Agents.DelegateInput(); + in.task = "do something"; + + com.google.genkit.core.GenkitException thrown = + org.junit.jupiter.api.Assertions.assertThrows( + com.google.genkit.core.GenkitException.class, () -> delegate.run(ctx, in)); + + assertTrue( + thrown.getMessage().toLowerCase().contains("failed"), + "expected descriptive failure text in the exception message, got: " + thrown.getMessage()); + assertTrue( + thrown.getMessage().contains("boom: downstream service unavailable"), + "expected the underlying error message folded into the exception message, got: " + + thrown.getMessage()); + } + + @Test + void delegationToolsProducesOneToolPerConfiguredAgentWithDescriptiveText() { + defineFixedAgent("researcher", "r"); + defineFixedAgent("writer", "w"); + List> tools = + Agents.delegationTools(AgentsOptions.builder().agents("researcher", "writer").build()); + + assertEquals(2, tools.size()); + com.google.genkit.ai.Tool researcherTool = + tool(tools, "delegate_to_researcher"); + com.google.genkit.ai.Tool writerTool = + tool(tools, "delegate_to_writer"); + + assertTrue(researcherTool.getDesc().getDescription().contains("researcher")); + assertTrue(writerTool.getDesc().getDescription().contains("writer")); + } + + @Test + void serverManagedSubAgentAlsoWorks() { + // sub-agent with a session store (server-managed) – delegation should still get its text. + InMemorySessionStore> store = new InMemorySessionStore<>(); + AgentFn> fn = + (sess, fnCtx) -> + AgentResult.builder() + .message(Message.model("server says hi")) + .finishReason(AgentFinishReason.STOP) + .build(); + CustomAgentConfig> config = + CustomAgentConfig.>builder().name("svragent").store(store).build(); + AgentActions.defineCustomAgent(registry, config, fn); + + List> tools = + Agents.delegationTools(AgentsOptions.builder().agents("svragent").build()); + com.google.genkit.ai.Tool delegate = + tool(tools, "delegate_to_svragent"); + + Agents.DelegateInput in = new Agents.DelegateInput(); + in.task = "ping"; + Agents.DelegateOutput out = delegate.run(ctx, in); + assertEquals("server says hi", out.response); + } +} diff --git a/plugins/middleware/src/test/java/com/google/genkit/plugins/middleware/ArtifactsTest.java b/plugins/middleware/src/test/java/com/google/genkit/plugins/middleware/ArtifactsTest.java new file mode 100644 index 000000000..078d03ace --- /dev/null +++ b/plugins/middleware/src/test/java/com/google/genkit/plugins/middleware/ArtifactsTest.java @@ -0,0 +1,182 @@ +/* + * 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. + * + * SPDX-License-Identifier: Apache-2.0 + */ + +package com.google.genkit.plugins.middleware; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertNull; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import com.google.genkit.ai.Tool; +import com.google.genkit.ai.agent.AgentSessionContext; +import com.google.genkit.ai.agent.Session; +import com.google.genkit.ai.agent.SessionState; +import com.google.genkit.core.ActionContext; +import com.google.genkit.core.DefaultRegistry; +import com.google.genkit.core.Registry; +import java.util.List; +import java.util.Map; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; + +/** TDD tests for {@link Artifacts} (Stage 6). */ +class ArtifactsTest { + + private Registry registry; + private ActionContext ctx; + + @BeforeEach + void setUp() { + registry = new DefaultRegistry(); + ctx = new ActionContext(registry); + } + + private static Session> newSession() { + return new Session<>(SessionState.>builder().build()); + } + + @SuppressWarnings("unchecked") + private static Tool readTool(List> tools) { + return (Tool) + tools.stream() + .filter(t -> Artifacts.READ_TOOL.equals(t.getName())) + .findFirst() + .orElseThrow(); + } + + @SuppressWarnings("unchecked") + private static Tool writeTool( + List> tools) { + return (Tool) + tools.stream() + .filter(t -> Artifacts.WRITE_TOOL.equals(t.getName())) + .findFirst() + .orElseThrow(); + } + + @Test + void defaultProducesReadAndWriteTools() { + List> tools = Artifacts.tools(ArtifactsOptions.defaults()); + assertEquals(2, tools.size()); + assertTrue(tools.stream().anyMatch(t -> Artifacts.READ_TOOL.equals(t.getName()))); + assertTrue(tools.stream().anyMatch(t -> Artifacts.WRITE_TOOL.equals(t.getName()))); + } + + @Test + void readonlyOmitsWriteTool() { + List> tools = Artifacts.tools(ArtifactsOptions.builder().readonly(true).build()); + assertEquals(1, tools.size()); + assertEquals(Artifacts.READ_TOOL, tools.get(0).getName()); + assertFalse(tools.stream().anyMatch(t -> Artifacts.WRITE_TOOL.equals(t.getName()))); + } + + @Test + void writeThenReadRoundTrips() { + List> tools = Artifacts.tools(ArtifactsOptions.defaults()); + Tool write = writeTool(tools); + Tool read = readTool(tools); + + Session> session = newSession(); + AgentSessionContext.run( + session, + () -> { + Artifacts.WriteInput wi = new Artifacts.WriteInput(); + wi.name = "notes"; + wi.content = "hello world"; + Artifacts.WriteOutput wo = write.run(ctx, wi); + assertEquals("ok", wo.status); + + Artifacts.ReadInput ri = new Artifacts.ReadInput(); + ri.name = "notes"; + Artifacts.ReadOutput ro = read.run(ctx, ri); + assertTrue(ro.found); + assertEquals("hello world", ro.content); + assertEquals("notes", ro.name); + }); + + // Persisted on the underlying store, deduplicated by name. + assertEquals(1, session.getArtifacts().size()); + assertEquals("notes", session.getArtifacts().get(0).getName()); + } + + @Test + void writeDeduplicatesByName() { + List> tools = Artifacts.tools(ArtifactsOptions.defaults()); + Tool write = writeTool(tools); + Tool read = readTool(tools); + + Session> session = newSession(); + AgentSessionContext.run( + session, + () -> { + Artifacts.WriteInput w1 = new Artifacts.WriteInput(); + w1.name = "doc"; + w1.content = "v1"; + write.run(ctx, w1); + + Artifacts.WriteInput w2 = new Artifacts.WriteInput(); + w2.name = "doc"; + w2.content = "v2"; + write.run(ctx, w2); + + Artifacts.ReadInput ri = new Artifacts.ReadInput(); + ri.name = "doc"; + Artifacts.ReadOutput ro = read.run(ctx, ri); + assertEquals("v2", ro.content); + }); + + assertEquals(1, session.getArtifacts().size()); + } + + @Test + void readMissingReportsNotFound() { + List> tools = Artifacts.tools(ArtifactsOptions.defaults()); + Tool read = readTool(tools); + + Session> session = newSession(); + AgentSessionContext.run( + session, + () -> { + Artifacts.ReadInput ri = new Artifacts.ReadInput(); + ri.name = "absent"; + Artifacts.ReadOutput ro = read.run(ctx, ri); + assertFalse(ro.found); + assertNull(ro.content); + }); + } + + @Test + void degradesWhenNoActiveSession() { + List> tools = Artifacts.tools(ArtifactsOptions.defaults()); + Tool read = readTool(tools); + Tool write = writeTool(tools); + + // No AgentSessionContext bound on this thread. + Artifacts.ReadInput ri = new Artifacts.ReadInput(); + ri.name = "x"; + Artifacts.ReadOutput ro = read.run(ctx, ri); + assertFalse(ro.found); + + Artifacts.WriteInput wi = new Artifacts.WriteInput(); + wi.name = "x"; + wi.content = "y"; + Artifacts.WriteOutput wo = write.run(ctx, wi); + assertEquals("no active session", wo.status); + } +} diff --git a/plugins/spring/src/main/java/com/google/genkit/plugins/spring/GenkitAgentController.java b/plugins/spring/src/main/java/com/google/genkit/plugins/spring/GenkitAgentController.java new file mode 100644 index 000000000..19ee10265 --- /dev/null +++ b/plugins/spring/src/main/java/com/google/genkit/plugins/spring/GenkitAgentController.java @@ -0,0 +1,569 @@ +/* + * 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. + * + * SPDX-License-Identifier: Apache-2.0 + */ + +package com.google.genkit.plugins.spring; + +import com.fasterxml.jackson.core.type.TypeReference; +import com.fasterxml.jackson.databind.JsonNode; +import com.fasterxml.jackson.databind.ObjectMapper; +import com.google.genkit.core.Action; +import com.google.genkit.core.ActionContext; +import com.google.genkit.core.ActionType; +import com.google.genkit.core.BidiAction; +import com.google.genkit.core.BufferedInputSource; +import com.google.genkit.core.GenkitException; +import com.google.genkit.core.Registry; +import jakarta.servlet.http.HttpServletRequest; +import java.io.PrintWriter; +import java.io.StringWriter; +import java.util.Enumeration; +import java.util.HashMap; +import java.util.List; +import java.util.Locale; +import java.util.Map; +import java.util.Set; +import java.util.UUID; +import java.util.function.Consumer; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; +import org.springframework.http.HttpStatus; +import org.springframework.http.MediaType; +import org.springframework.http.ResponseEntity; +import org.springframework.web.bind.annotation.PathVariable; +import org.springframework.web.bind.annotation.PostMapping; +import org.springframework.web.bind.annotation.RequestBody; +import org.springframework.web.bind.annotation.RestController; +import org.springframework.web.servlet.mvc.method.annotation.ResponseBodyEmitter; + +/** + * REST controller that exposes Genkit agents ({@code ActionType.AGENT} bidi actions) as HTTP + * endpoints, at parity with the {@code plugins/jetty} module's {@code AgentHandler} / {@code + * CompanionHandler}. + * + *

For each registered agent named {@code } this mounts, at the root path (NOT under {@code + * /api/...}, so that {@code RemoteAgent}/{@code HttpAgentTransport} clients that derive companion + * URLs by simple string concatenation on the base agent URL keep working regardless of which server + * plugin they talk to): + * + *

    + *
  • {@code POST /} — one turn per request (non-streaming or SSE). + *
  • {@code POST //getSnapshot} — companion {@code agent-snapshot} action, if registered. + *
  • {@code POST //abort} — companion {@code agent-abort} action, if registered. + *
+ * + *

Actions are looked up by name at request time (mirroring {@link + * GenkitFlowController#findFlowByName}) rather than pre-registered per-agent handlers, since Spring + * MVC controllers use static {@code @RequestMapping}-family annotations. + */ +@RestController +public class GenkitAgentController { + + private static final Logger logger = LoggerFactory.getLogger(GenkitAgentController.class); + + private final ObjectMapper objectMapper; + + /** + * Creates a new GenkitAgentController. + * + * @param objectMapper the ObjectMapper for JSON serialization + */ + public GenkitAgentController(ObjectMapper objectMapper) { + this.objectMapper = objectMapper; + logRegisteredEndpoints(); + } + + private Registry getRegistry() { + SpringPlugin plugin = SpringPlugin.getInstance(); + return plugin != null ? plugin.getRegistry() : null; + } + + /** Logs all registered agent endpoints (and their companions, if present). */ + private void logRegisteredEndpoints() { + Registry registry = getRegistry(); + if (registry == null) { + return; + } + List> agents = registry.listActions(ActionType.AGENT); + for (Action action : agents) { + String name = action.getName(); + logger.info("Registered agent endpoint: /{}", name); + if (registry.lookupAction(ActionType.AGENT_SNAPSHOT.keyFromName(name)) != null) { + logger.info("Registered agent companion endpoint: /{}/getSnapshot", name); + } + if (registry.lookupAction(ActionType.AGENT_ABORT.keyFromName(name)) != null) { + logger.info("Registered agent companion endpoint: /{}/abort", name); + } + } + } + + /** + * Finds an agent action by name. + * + * @param agentName the name of the agent + * @return the bidi action, or null if not found or not a {@link BidiAction} + */ + private BidiAction findAgentByName(String agentName) { + Registry registry = getRegistry(); + if (registry == null) { + return null; + } + List> agents = registry.listActions(ActionType.AGENT); + for (Action action : agents) { + if (action.getName().equals(agentName) && action instanceof BidiAction) { + @SuppressWarnings("unchecked") + BidiAction bidi = + (BidiAction) action; + return bidi; + } + } + return null; + } + + /** + * Finds a companion action (getSnapshot/abort) by its {@link ActionType} and agent name. + * + * @param type the companion action type ({@code AGENT_SNAPSHOT} or {@code AGENT_ABORT}) + * @param agentName the agent name + * @return the companion action, or null if not registered + */ + @SuppressWarnings("unchecked") + private Action findCompanion(ActionType type, String agentName) { + Registry registry = getRegistry(); + if (registry == null) { + return null; + } + Action action = registry.lookupAction(type.keyFromName(agentName)); + return (Action) action; + } + + // --------------------------------------------------------------------------- + // Main turn endpoint + // --------------------------------------------------------------------------- + + /** + * Handles one turn of an agent conversation. Non-streaming requests receive a single {@code + * {"result": }} JSON body; requests that ask for streaming (via {@code Accept: + * text/event-stream} or {@code ?stream=true}) receive an SSE response instead — see {@link + * #streamAgent}. + * + * @param agentName the name of the agent + * @param body the request envelope: {@code {"data": , "init": , "context": + * }} + * @param request the servlet request, used to detect streaming + * @return the JSON result envelope, or an error envelope with a mapped HTTP status + */ + @PostMapping(value = "/{agentName}", consumes = MediaType.APPLICATION_JSON_VALUE) + public Object runAgent( + @PathVariable String agentName, + @RequestBody(required = false) JsonNode body, + HttpServletRequest request) { + if (isStreamingRequested(request)) { + return streamAgent(agentName, body, request); + } + + Registry registry = getRegistry(); + if (registry == null) { + return ResponseEntity.status(HttpStatus.SERVICE_UNAVAILABLE) + .body(Map.of("error", "Registry not initialized")); + } + + BidiAction action = findAgentByName(agentName); + if (action == null) { + return ResponseEntity.status(HttpStatus.NOT_FOUND) + .body(Map.of("error", "Agent not found: " + agentName)); + } + + JsonNode data = body != null ? body.get("data") : null; + JsonNode init = body != null ? body.get("init") : null; + Map userContext = mergeHeadersIntoContext(parseContext(body), request); + + try { + BufferedInputSource inputs = new BufferedInputSource<>(); + if (data != null && !data.isNull()) { + inputs.offer(data); + } + inputs.end(); + + ActionContext context = + ActionContext.builder().registry(registry).context(userContext).build(); + JsonNode result = action.runBidiJson(context, init, inputs, null); + + Map envelope = new HashMap<>(); + envelope.put("result", result); + return ResponseEntity.ok(envelope); + } catch (Exception e) { + logger.error("Error handling agent request: {}", agentName, e); + return errorResponse(e); + } + } + + /** + * Handles one streaming turn of an agent conversation over SSE. Emits one {@code data: + * {"message": }} frame per streamed chunk, then a final {@code data: {"result": }} + * frame (or {@code data: {"error": {...}}} on failure). The turn runs on a background thread so + * the servlet container's async dispatch mechanism can flush frames as they are produced. + * + * @param agentName the name of the agent + * @param body the request envelope: {@code {"data": , "init": , "context": + * }} + * @param request the servlet request, whose headers are merged into the run's user context + * @return a response wrapping a {@link ResponseBodyEmitter} that streams the turn's chunks and + * final result, with the {@code X-Genkit-Stream-Id} header set (matching Jetty's {@code + * AgentHandler}) + */ + private ResponseEntity streamAgent( + String agentName, JsonNode body, HttpServletRequest request) { + // A plain ResponseBodyEmitter (rather than SseEmitter) is used because SseEmitter's + // SseEventBuilder hardcodes a "data:" prefix with no trailing space, whereas Jetty's + // AgentHandler (and the HttpAgentTransport client's SSE_DATA_PREFIX = "data: " parser) require + // a space after the colon. Sending fully pre-formatted "data: \n\n" strings directly as + // TEXT_PLAIN chunks reproduces Jetty's exact byte-level framing. + ResponseBodyEmitter emitter = new ResponseBodyEmitter(0L); + + Registry registry = getRegistry(); + if (registry == null) { + completeWithEnvelopeError(emitter, "SERVICE_UNAVAILABLE", "Registry not initialized"); + return sseResponse(emitter); + } + + BidiAction action = findAgentByName(agentName); + if (action == null) { + completeWithEnvelopeError(emitter, "NOT_FOUND", "Agent not found: " + agentName); + return sseResponse(emitter); + } + + JsonNode data = body != null ? body.get("data") : null; + JsonNode init = body != null ? body.get("init") : null; + Map userContext = mergeHeadersIntoContext(parseContext(body), request); + + Thread worker = + new Thread( + () -> { + Consumer streamCallback = + chunk -> { + try { + Map frame = new HashMap<>(); + frame.put("message", chunk); + sendSseFrame(emitter, frame); + } catch (Exception e) { + throw new RuntimeException("Failed to write SSE frame", e); + } + }; + + try { + BufferedInputSource inputs = new BufferedInputSource<>(); + if (data != null && !data.isNull()) { + inputs.offer(data); + } + inputs.end(); + + ActionContext context = + ActionContext.builder().registry(registry).context(userContext).build(); + JsonNode result = action.runBidiJson(context, init, inputs, streamCallback); + + Map resultFrame = new HashMap<>(); + resultFrame.put("result", result); + sendSseFrame(emitter, resultFrame); + emitter.complete(); + } catch (Exception e) { + logger.error("Error streaming agent request: {}", agentName, e); + try { + Map errorFrame = new HashMap<>(); + errorFrame.put("error", errorBody(errorCodeFromError(e), e)); + sendSseFrame(emitter, errorFrame); + emitter.complete(); + } catch (Exception writeError) { + emitter.completeWithError(writeError); + } + } + }, + "genkit-agent-sse-" + agentName); + worker.setDaemon(true); + worker.start(); + + return sseResponse(emitter); + } + + /** + * Wraps a {@link ResponseBodyEmitter} in a {@link ResponseEntity} carrying the {@code + * X-Genkit-Stream-Id} header and {@code text/event-stream} content type, mirroring Jetty's {@code + * AgentHandler#handleStreaming}. Headers must be set before the emitter is returned from the + * controller method, since headers cannot be added once the async response body starts streaming. + */ + private ResponseEntity sseResponse(ResponseBodyEmitter emitter) { + return ResponseEntity.ok() + .header("X-Genkit-Stream-Id", UUID.randomUUID().toString()) + .header("Cache-Control", "no-cache") + .contentType(MediaType.TEXT_EVENT_STREAM) + .body(emitter); + } + + /** + * Writes one SSE frame in Jetty's exact wire format: {@code "data: " + json + "\n\n"}, sent as a + * raw {@code TEXT_PLAIN} chunk so no additional SSE-specific encoding is applied. + */ + private void sendSseFrame(ResponseBodyEmitter emitter, Object payload) throws Exception { + String frame = "data: " + objectMapper.writeValueAsString(payload) + "\n\n"; + emitter.send(frame, MediaType.TEXT_PLAIN); + } + + /** Sends a single error SSE frame (used before the turn even starts) and completes. */ + private void completeWithEnvelopeError(ResponseBodyEmitter emitter, String code, String message) { + try { + Map errorFrame = new HashMap<>(); + Map error = new HashMap<>(); + error.put("status", code); + error.put("message", message); + error.put("details", Map.of()); + errorFrame.put("error", error); + sendSseFrame(emitter, errorFrame); + emitter.complete(); + } catch (Exception e) { + emitter.completeWithError(e); + } + } + + // --------------------------------------------------------------------------- + // Companion endpoints: getSnapshot / abort + // --------------------------------------------------------------------------- + + /** + * Companion endpoint for retrieving a previously-saved agent session snapshot. + * + * @param agentName the name of the agent + * @param body the request envelope: {@code {"data": , "context": }} + * @param request the servlet request, whose headers are merged into the run's user context + * @return the JSON result envelope, or an error envelope with a mapped HTTP status + */ + @PostMapping(value = "/{agentName}/getSnapshot", consumes = MediaType.APPLICATION_JSON_VALUE) + public ResponseEntity getSnapshot( + @PathVariable String agentName, + @RequestBody(required = false) JsonNode body, + HttpServletRequest request) { + return runCompanion(ActionType.AGENT_SNAPSHOT, agentName, body, request); + } + + /** + * Companion endpoint for aborting a pending agent session snapshot. + * + * @param agentName the name of the agent + * @param body the request envelope: {@code {"data": , "context": }} + * @param request the servlet request, whose headers are merged into the run's user context + * @return the JSON result envelope, or an error envelope with a mapped HTTP status + */ + @PostMapping(value = "/{agentName}/abort", consumes = MediaType.APPLICATION_JSON_VALUE) + public ResponseEntity abort( + @PathVariable String agentName, + @RequestBody(required = false) JsonNode body, + HttpServletRequest request) { + return runCompanion(ActionType.AGENT_ABORT, agentName, body, request); + } + + private ResponseEntity runCompanion( + ActionType type, String agentName, JsonNode body, HttpServletRequest request) { + Registry registry = getRegistry(); + if (registry == null) { + return ResponseEntity.status(HttpStatus.SERVICE_UNAVAILABLE) + .body(Map.of("error", "Registry not initialized")); + } + + Action action = findCompanion(type, agentName); + if (action == null) { + return ResponseEntity.status(HttpStatus.NOT_FOUND) + .body(Map.of("error", "Agent companion not found: " + agentName)); + } + + JsonNode data = body != null ? body.get("data") : null; + Map userContext = mergeHeadersIntoContext(parseContext(body), request); + + try { + ActionContext context = + ActionContext.builder().registry(registry).context(userContext).build(); + JsonNode result = action.runJson(context, data, null); + + Map envelope = new HashMap<>(); + envelope.put("result", result); + return ResponseEntity.ok(envelope); + } catch (Exception e) { + logger.error("Error handling agent companion request: {}", agentName, e); + return errorResponse(e); + } + } + + // --------------------------------------------------------------------------- + // Shared helpers + // --------------------------------------------------------------------------- + + /** + * Parses the optional {@code context} object from a request body envelope into a {@code + * Map}. Threaded into the run's ActionContext so tools/flows can read it. + * + * @param body the parsed request body (may be null) + * @return the parsed context map, or null if absent/blank + */ + private Map parseContext(JsonNode body) { + if (body == null || !body.has("context") || body.get("context").isNull()) { + return null; + } + JsonNode contextNode = body.get("context"); + if (!contextNode.isObject()) { + return null; + } + return objectMapper.convertValue(contextNode, new TypeReference>() {}); + } + + /** + * HTTP framing/transport headers that are excluded from the {@code "headers"} sub-map threaded + * into the run's {@link ActionContext}. These describe the HTTP message itself (body encoding, + * connection lifecycle, routing) rather than application-level data a tool/flow would care about. + * Everything else — including custom headers like {@code Authorization} or {@code X-*} — is + * included; when in doubt we err on the side of including a header rather than filtering it. + * + *

Kept identical to {@code plugins/jetty}'s {@code JettyPlugin.EXCLUDED_HEADERS} so a tool + * reading {@code ctx.getContext().get("headers")} behaves the same regardless of which server + * plugin served the request. + */ + private static final Set EXCLUDED_HEADERS = + Set.of("content-type", "content-length", "accept", "accept-encoding", "connection", "host"); + + /** + * Merges incoming HTTP request headers into the request-scoped user context returned by {@link + * #parseContext(JsonNode)}. + * + *

Design: headers are exposed to tools/flows the same way the JSON-body {@code context} + * field already is — via {@link ActionContext#getContext()} — by nesting them under a reserved + * {@code "headers"} key as a {@code Map} (single-valued, to match the shape of + * {@code RemoteAgentOptions.headers()} on the client). This is additive and cannot silently + * collide with arbitrary body-context keys other than a literal top-level {@code "headers"} key. + * Precedence: if the body-supplied {@code context} already defines a {@code "headers"} + * entry, that body value wins and incoming HTTP headers are dropped for that key (body {@code + * context} always takes precedence over transport-level data); otherwise the HTTP headers + * populate {@code context.get("headers")}. + * + * @param bodyContext the context map parsed from the request body (may be null) + * @param request the incoming servlet request, whose headers are folded in + * @return a merged context map (never null if headers are present; may be null if both the body + * context and the header set are empty) + */ + private Map mergeHeadersIntoContext( + Map bodyContext, HttpServletRequest request) { + Map headers = parseHeaders(request); + if (headers.isEmpty()) { + return bodyContext; + } + Map merged = bodyContext != null ? new HashMap<>(bodyContext) : new HashMap<>(); + merged.putIfAbsent("headers", headers); + return merged; + } + + /** + * Extracts incoming HTTP request headers (excluding standard framing headers, see {@link + * #EXCLUDED_HEADERS}) into a {@code Map}. If a header name repeats, the last value + * wins (mirroring {@code HttpServletRequest#getHeader}). + * + * @param request the incoming servlet request + * @return a mutable map of header name to value (never null; may be empty) + */ + private static Map parseHeaders(HttpServletRequest request) { + Map result = new HashMap<>(); + Enumeration names = request.getHeaderNames(); + if (names == null) { + return result; + } + while (names.hasMoreElements()) { + String name = names.nextElement(); + if (name == null || EXCLUDED_HEADERS.contains(name.toLowerCase(Locale.ROOT))) { + continue; + } + result.put(name, request.getHeader(name)); + } + return result; + } + + /** Returns true if the request asks for SSE streaming. */ + private static boolean isStreamingRequested(HttpServletRequest request) { + String accept = request.getHeader("Accept"); + if (accept != null && accept.contains("text/event-stream")) { + return true; + } + String query = request.getQueryString(); + return query != null && query.contains("stream=true"); + } + + /** Builds a {@code {"error": {...}}} envelope response with a status-mapped HTTP code. */ + private ResponseEntity errorResponse(Exception e) { + int httpStatus = statusFromError(e); + String code = errorCodeFromError(e); + Map envelope = new HashMap<>(); + envelope.put("error", errorBody(code, e)); + return ResponseEntity.status(httpStatus).body(envelope); + } + + /** Builds a structured error body: {@code {status, message, details: {stack}}}. */ + private static Map errorBody(String code, Throwable e) { + String message = e.getMessage() != null ? e.getMessage() : "Unknown error"; + StringWriter sw = new StringWriter(); + e.printStackTrace(new PrintWriter(sw)); + Map body = new HashMap<>(); + body.put("status", code); + body.put("message", message); + body.put("details", Map.of("stack", sw.toString())); + return body; + } + + /** Derives a status string (mirroring the error's status field) from a thrown error. */ + private static String errorCodeFromError(Throwable e) { + if (e instanceof GenkitException) { + String c = ((GenkitException) e).getErrorCode(); + if (c != null && !c.isEmpty()) { + return c; + } + } + return "INTERNAL"; + } + + /** Maps a thrown error to an HTTP status code derived from its status field. */ + private static int statusFromError(Throwable e) { + String code = errorCodeFromError(e); + switch (code) { + case "NOT_FOUND": + return 404; + case "INVALID_ARGUMENT": + case "FAILED_PRECONDITION": + case "OUT_OF_RANGE": + return 400; + case "UNAUTHENTICATED": + return 401; + case "PERMISSION_DENIED": + return 403; + case "ALREADY_EXISTS": + case "ABORTED": + return 409; + case "RESOURCE_EXHAUSTED": + return 429; + case "UNIMPLEMENTED": + return 501; + case "UNAVAILABLE": + return 503; + case "DEADLINE_EXCEEDED": + return 504; + default: + return 500; + } + } +} diff --git a/plugins/spring/src/main/java/com/google/genkit/plugins/spring/GenkitSpringApplication.java b/plugins/spring/src/main/java/com/google/genkit/plugins/spring/GenkitSpringApplication.java index d0ec5c6bc..04130560f 100644 --- a/plugins/spring/src/main/java/com/google/genkit/plugins/spring/GenkitSpringApplication.java +++ b/plugins/spring/src/main/java/com/google/genkit/plugins/spring/GenkitSpringApplication.java @@ -19,8 +19,12 @@ package com.google.genkit.plugins.spring; import com.fasterxml.jackson.databind.ObjectMapper; +import java.util.List; import org.springframework.boot.autoconfigure.SpringBootApplication; import org.springframework.context.annotation.Bean; +import org.springframework.http.converter.HttpMessageConverter; +import org.springframework.http.converter.json.MappingJackson2HttpMessageConverter; +import org.springframework.web.servlet.config.annotation.WebMvcConfigurer; /** * Spring Boot application class for Genkit. @@ -51,4 +55,41 @@ public ObjectMapper objectMapper() { public GenkitFlowController genkitFlowController(ObjectMapper objectMapper) { return new GenkitFlowController(objectMapper); } + + /** + * Creates the Genkit agent controller bean. + * + * @param objectMapper the ObjectMapper for JSON serialization + * @return the agent controller + */ + @Bean + public GenkitAgentController genkitAgentController(ObjectMapper objectMapper) { + return new GenkitAgentController(objectMapper); + } + + /** + * Configures Spring MVC to deserialize {@code @RequestBody} JSON using the Jackson 2 ({@code + * com.fasterxml.jackson.databind}) {@link ObjectMapper} that {@link GenkitFlowController} and + * {@link GenkitAgentController} use for {@code JsonNode} handling. + * + *

{@code spring-boot-starter-web} on Spring Boot 4 auto-configures a Jackson 3 ({@code + * tools.jackson.databind}) message converter as the default JSON converter. Jackson 3's binder + * cannot construct instances of the Jackson 2 {@code com.fasterxml.jackson.databind.JsonNode} + * type used throughout Genkit's core action APIs, so {@code @RequestBody JsonNode} parameters + * would otherwise fail to bind. Prepending a {@link MappingJackson2HttpMessageConverter} backed + * by the Jackson 2 {@code ObjectMapper} makes it the preferred converter for {@code + * application/json} bodies. + * + * @param objectMapper the Jackson 2 ObjectMapper for JSON serialization + * @return the MVC configurer + */ + @Bean + public WebMvcConfigurer genkitJackson2MessageConverterConfigurer(ObjectMapper objectMapper) { + return new WebMvcConfigurer() { + @Override + public void extendMessageConverters(List> converters) { + converters.add(0, new MappingJackson2HttpMessageConverter(objectMapper)); + } + }; + } } diff --git a/plugins/spring/src/test/java/com/google/genkit/plugins/spring/AgentAbortHttpTest.java b/plugins/spring/src/test/java/com/google/genkit/plugins/spring/AgentAbortHttpTest.java new file mode 100644 index 000000000..a60423c55 --- /dev/null +++ b/plugins/spring/src/test/java/com/google/genkit/plugins/spring/AgentAbortHttpTest.java @@ -0,0 +1,245 @@ +/* + * 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. + * + * SPDX-License-Identifier: Apache-2.0 + */ + +package com.google.genkit.plugins.spring; + +import static org.junit.jupiter.api.Assertions.*; + +import com.fasterxml.jackson.databind.JsonNode; +import com.fasterxml.jackson.databind.ObjectMapper; +import com.google.genkit.ai.Message; +import com.google.genkit.ai.agent.AgentFinishReason; +import com.google.genkit.ai.agent.AgentResult; +import com.google.genkit.ai.agent.CustomAgentConfig; +import com.google.genkit.ai.agent.InMemorySessionStore; +import com.google.genkit.ai.agent.internal.AgentActions; +import com.google.genkit.core.DefaultRegistry; +import java.io.IOException; +import java.net.HttpURLConnection; +import java.net.URI; +import java.net.URL; +import java.net.http.HttpClient; +import java.net.http.HttpRequest; +import java.net.http.HttpResponse; +import java.util.Map; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.TimeUnit; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.Test; + +/** + * Integration test verifying the abort companion endpoint over the HTTP wire using SpringPlugin. + * + *

Mirrors {@code plugins/jetty}'s {@code AgentAbortHttpTest}. Cross-referenced with {@code + * Agent.abort(String)} / {@code AgentActions.buildAbortAction}: the abort companion only flips a + * stored snapshot's status from {@code PENDING} to {@code ABORTED}; it never interrupts a running + * foreground {@code AgentFn} call (there is no cancellation hook threaded into a synchronous {@code + * AgentFn.run}). This test therefore targets what is actually implemented: aborting a snapshot that + * is genuinely in the {@code PENDING} window opened by a detached turn, which synchronously + * persists the {@code PENDING} row before the HTTP response returns and whose finalizer never + * overwrites an {@code ABORTED} status (race-safe by design) — then confirms via {@code + * getSnapshot} that the status is {@code ABORTED} and stays that way even after the background turn + * body completes. + */ +class AgentAbortHttpTest { + + private static final ObjectMapper MAPPER = new ObjectMapper(); + + private SpringPlugin plugin; + + @AfterEach + void tearDown() throws Exception { + if (plugin != null) { + plugin.stop(); + } + } + + /** + * Starts a SpringPlugin on the given port with a server-managed custom agent named {@code + * abortAgent} whose AgentFn blocks on {@code latch} until released, giving the test a + * deterministic window in which the snapshot is guaranteed to still be {@code PENDING}. + */ + private void startWithBlockingAgent(int port, CountDownLatch releaseLatch) throws Exception { + DefaultRegistry registry = new DefaultRegistry(); + + AgentActions.defineCustomAgent( + registry, + CustomAgentConfig.>builder() + .name("abortAgent") + .store(new InMemorySessionStore<>()) + .build(), + (sess, fnCtx) -> { + // Block the detached background turn body until the test releases it, so the + // PENDING window is deterministic and long enough to reliably call abort(). + releaseLatch.await(10, TimeUnit.SECONDS); + return AgentResult.builder() + .message(Message.model("done")) + .finishReason(AgentFinishReason.STOP) + .build(); + }); + + SpringPluginOptions options = SpringPluginOptions.builder().port(port).build(); + plugin = new SpringPlugin(options); + plugin.init(registry); + + Thread serverThread = + new Thread( + () -> { + try { + plugin.start(); + } catch (Exception e) { + // Server stopped. + } + }); + serverThread.setDaemon(true); + serverThread.start(); + + waitForServer(port); + } + + /** + * Integration test: detach a turn (leaving its snapshot PENDING while the AgentFn blocks) -> POST + * /abort for that snapshot -> assert the abort response reports status ABORTED -> release the + * blocked turn so it finalizes in the background -> poll getSnapshot and assert the status + * remains ABORTED (proving the finalizer's abort-aware guard truly won the race, not just that + * the abort call itself returned ABORTED). + */ + @Test + void testAbortOverHttp() throws Exception { + int port = findAvailablePort(); + CountDownLatch releaseLatch = new CountDownLatch(1); + startWithBlockingAgent(port, releaseLatch); + + HttpClient client = HttpClient.newHttpClient(); + + // Step 1: POST a detach turn. The AgentFn blocks on releaseLatch, so the snapshot remains + // PENDING until we count down the latch below. + String detachBody = + "{\"data\":{\"detach\":true,\"message\":{\"role\":\"user\"," + + "\"content\":[{\"text\":\"go\"}]}},\"init\":{}}"; + + HttpRequest detachRequest = + HttpRequest.newBuilder() + .uri(URI.create("http://localhost:" + port + "/abortAgent")) + .header("Content-Type", "application/json") + .POST(HttpRequest.BodyPublishers.ofString(detachBody)) + .build(); + + HttpResponse detachResponse = + client.send(detachRequest, HttpResponse.BodyHandlers.ofString()); + assertEquals(200, detachResponse.statusCode(), "detach POST body: " + detachResponse.body()); + + JsonNode detachRoot = MAPPER.readTree(detachResponse.body()); + JsonNode result = detachRoot.path("result"); + assertEquals( + "detached", + result.path("finishReason").asText(""), + "expected finishReason=detached: " + detachResponse.body()); + String snapshotId = result.path("snapshotId").asText(""); + assertFalse(snapshotId.isEmpty(), "expected non-empty snapshotId"); + + // Step 2: POST /abort for the still-PENDING snapshot. + String abortBody = "{\"data\":{\"snapshotId\":\"" + snapshotId + "\"}}"; + HttpRequest abortRequest = + HttpRequest.newBuilder() + .uri(URI.create("http://localhost:" + port + "/abortAgent/abort")) + .header("Content-Type", "application/json") + .POST(HttpRequest.BodyPublishers.ofString(abortBody)) + .build(); + + HttpResponse abortResponse = + client.send(abortRequest, HttpResponse.BodyHandlers.ofString()); + assertEquals(200, abortResponse.statusCode(), "abort POST body: " + abortResponse.body()); + + JsonNode abortRoot = MAPPER.readTree(abortResponse.body()); + assertTrue(abortRoot.has("result"), "expected result envelope: " + abortResponse.body()); + JsonNode abortResult = abortRoot.path("result"); + assertEquals(snapshotId, abortResult.path("snapshotId").asText("")); + assertEquals( + "aborted", + abortResult.path("status").asText(""), + "expected abort response status=aborted: " + abortResponse.body()); + + // Step 3: getSnapshot immediately (still PENDING-turned-ABORTED, before the background turn + // body has been released) should already reflect ABORTED. + String getSnapshotBody = "{\"data\":{\"snapshotId\":\"" + snapshotId + "\"}}"; + HttpRequest snapshotRequest = + HttpRequest.newBuilder() + .uri(URI.create("http://localhost:" + port + "/abortAgent/getSnapshot")) + .header("Content-Type", "application/json") + .POST(HttpRequest.BodyPublishers.ofString(getSnapshotBody)) + .build(); + + HttpResponse snapResponse1 = + client.send(snapshotRequest, HttpResponse.BodyHandlers.ofString()); + assertEquals(200, snapResponse1.statusCode(), "getSnapshot body: " + snapResponse1.body()); + JsonNode snapResult1 = MAPPER.readTree(snapResponse1.body()).path("result"); + assertEquals( + "aborted", + snapResult1.path("status").asText(""), + "expected snapshot status=aborted before finalize: " + snapResult1); + + // Step 4: release the blocked background turn so it runs to completion and attempts to + // finalize the (now ABORTED) snapshot to COMPLETED. + releaseLatch.countDown(); + + // Step 5: poll getSnapshot for up to 3s and assert the status NEVER reverts from "aborted" + // (proving DetachController's finalizer really never overwrites an ABORTED row). + long deadline = System.currentTimeMillis() + 3000; + String lastStatus = "aborted"; + while (System.currentTimeMillis() < deadline) { + HttpResponse snapResponse = + client.send(snapshotRequest, HttpResponse.BodyHandlers.ofString()); + assertEquals(200, snapResponse.statusCode(), "getSnapshot body: " + snapResponse.body()); + JsonNode snapResult = MAPPER.readTree(snapResponse.body()).path("result"); + lastStatus = snapResult.path("status").asText(""); + assertEquals( + "aborted", + lastStatus, + "snapshot status must remain aborted (finalizer must not overwrite it): " + snapResult); + Thread.sleep(50); + } + assertEquals("aborted", lastStatus); + } + + private static int findAvailablePort() throws IOException { + try (java.net.ServerSocket socket = new java.net.ServerSocket(0)) { + return socket.getLocalPort(); + } + } + + private static void waitForServer(int port) throws Exception { + int maxRetries = 100; + for (int i = 0; i < maxRetries; i++) { + try { + HttpURLConnection conn = + (HttpURLConnection) new URL("http://localhost:" + port + "/health").openConnection(); + conn.setRequestMethod("GET"); + conn.setConnectTimeout(200); + conn.setReadTimeout(200); + if (conn.getResponseCode() == 200) { + return; + } + } catch (IOException e) { + // Server not ready yet. + } + Thread.sleep(100); + } + fail("Server did not start within timeout"); + } +} diff --git a/plugins/spring/src/test/java/com/google/genkit/plugins/spring/AgentDetachHttpTest.java b/plugins/spring/src/test/java/com/google/genkit/plugins/spring/AgentDetachHttpTest.java new file mode 100644 index 000000000..3a7dc8410 --- /dev/null +++ b/plugins/spring/src/test/java/com/google/genkit/plugins/spring/AgentDetachHttpTest.java @@ -0,0 +1,222 @@ +/* + * 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. + * + * SPDX-License-Identifier: Apache-2.0 + */ + +package com.google.genkit.plugins.spring; + +import static org.junit.jupiter.api.Assertions.*; + +import com.fasterxml.jackson.databind.JsonNode; +import com.fasterxml.jackson.databind.ObjectMapper; +import com.google.genkit.ai.Message; +import com.google.genkit.ai.agent.AgentFinishReason; +import com.google.genkit.ai.agent.AgentResult; +import com.google.genkit.ai.agent.CustomAgentConfig; +import com.google.genkit.ai.agent.InMemorySessionStore; +import com.google.genkit.ai.agent.internal.AgentActions; +import com.google.genkit.core.DefaultRegistry; +import java.io.IOException; +import java.net.HttpURLConnection; +import java.net.URI; +import java.net.URL; +import java.net.http.HttpClient; +import java.net.http.HttpRequest; +import java.net.http.HttpResponse; +import java.util.Map; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.Test; + +/** + * Integration test verifying the detach feature works over the HTTP wire using SpringPlugin. + * + *

Mirrors {@code plugins/jetty}'s {@code AgentDetachHttpTest}. Proves that a DETACH turn works + * through the Spring agent HTTP endpoint: POST a turn with {@code detach:true} → server returns + * {@code finishReason: "detached"} + a pending {@code snapshotId} immediately → the background work + * finalizes → polling the {@code getSnapshot} companion shows {@code status: "completed"} with the + * accumulated state. + */ +class AgentDetachHttpTest { + + private static final ObjectMapper MAPPER = new ObjectMapper(); + + private SpringPlugin plugin; + + @AfterEach + void tearDown() throws Exception { + if (plugin != null) { + plugin.stop(); + } + } + + /** + * Starts a SpringPlugin on the given port with a server-managed custom agent named {@code + * detachAgent} that returns immediately (so background work finalizes quickly). + */ + private void startWithDetachAgent(int port) throws Exception { + DefaultRegistry registry = new DefaultRegistry(); + + // Register a server-managed custom agent. The AgentFn returns immediately with a STOP + // finish reason — detach background work will finalize quickly. + AgentActions.defineCustomAgent( + registry, + CustomAgentConfig.>builder() + .name("detachAgent") + .store(new InMemorySessionStore<>()) + .build(), + (sess, fnCtx) -> + AgentResult.builder() + .message(Message.model("done")) + .finishReason(AgentFinishReason.STOP) + .build()); + + SpringPluginOptions options = SpringPluginOptions.builder().port(port).build(); + plugin = new SpringPlugin(options); + plugin.init(registry); + + Thread serverThread = + new Thread( + () -> { + try { + plugin.start(); + } catch (Exception e) { + // Server stopped. + } + }); + serverThread.setDaemon(true); + serverThread.start(); + + waitForServer(port); + } + + /** + * Integration test: POST with detach:true → assert finishReason=="detached" + snapshotId → poll + * getSnapshot until status=="completed" → assert messages non-empty. + */ + @Test + void testDetachOverHttp() throws Exception { + int port = findAvailablePort(); + startWithDetachAgent(port); + + HttpClient client = HttpClient.newHttpClient(); + + // Step 1: POST a detach turn. + String detachBody = + "{\"data\":{\"detach\":true,\"message\":{\"role\":\"user\"," + + "\"content\":[{\"text\":\"go\"}]}},\"init\":{}}"; + + HttpRequest detachRequest = + HttpRequest.newBuilder() + .uri(URI.create("http://localhost:" + port + "/detachAgent")) + .header("Content-Type", "application/json") + .POST(HttpRequest.BodyPublishers.ofString(detachBody)) + .build(); + + HttpResponse detachResponse = + client.send(detachRequest, HttpResponse.BodyHandlers.ofString()); + assertEquals(200, detachResponse.statusCode(), "detach POST body: " + detachResponse.body()); + + JsonNode detachRoot = MAPPER.readTree(detachResponse.body()); + assertTrue(detachRoot.has("result"), "expected result envelope, got: " + detachResponse.body()); + + JsonNode result = detachRoot.path("result"); + + // Assert: finishReason is "detached". + String finishReason = result.path("finishReason").asText(""); + assertEquals( + "detached", + finishReason, + "expected finishReason=detached, got: " + finishReason + " body: " + detachResponse.body()); + + // Assert: snapshotId is present and non-empty. + String snapshotId = result.path("snapshotId").asText(""); + assertFalse( + snapshotId.isEmpty(), + "expected non-empty snapshotId, got: " + snapshotId + " body: " + detachResponse.body()); + + // Step 2: Poll getSnapshot until status becomes "completed" (or timeout at 5s). + String getSnapshotBody = "{\"data\":{\"snapshotId\":\"" + snapshotId + "\"}}"; + HttpRequest snapshotRequest = + HttpRequest.newBuilder() + .uri(URI.create("http://localhost:" + port + "/detachAgent/getSnapshot")) + .header("Content-Type", "application/json") + .POST(HttpRequest.BodyPublishers.ofString(getSnapshotBody)) + .build(); + + JsonNode snapResult = null; + String snapStatus = ""; + long deadline = System.currentTimeMillis() + 5000; + while (System.currentTimeMillis() < deadline) { + HttpResponse snapResponse = + client.send(snapshotRequest, HttpResponse.BodyHandlers.ofString()); + assertEquals(200, snapResponse.statusCode(), "getSnapshot body: " + snapResponse.body()); + + JsonNode snapRoot = MAPPER.readTree(snapResponse.body()); + assertTrue( + snapRoot.has("result"), + "expected result envelope from getSnapshot, got: " + snapResponse.body()); + + snapResult = snapRoot.path("result"); + snapStatus = snapResult.path("status").asText(""); + + if ("completed".equals(snapStatus) + || "failed".equals(snapStatus) + || "aborted".equals(snapStatus)) { + break; + } + Thread.sleep(100); + } + + // Assert: snapshot finalized to "completed". + assertNotNull(snapResult, "snapResult should not be null"); + assertEquals( + "completed", + snapStatus, + "expected snapshot status=completed, got: " + snapStatus + " snap: " + snapResult); + + // Assert: state.messages is non-empty (background turn ran and accumulated messages). + JsonNode messages = snapResult.path("state").path("messages"); + assertFalse( + messages.isMissingNode() || messages.isEmpty(), + "expected non-empty state.messages in completed snapshot, got: " + snapResult); + } + + private static int findAvailablePort() throws IOException { + try (java.net.ServerSocket socket = new java.net.ServerSocket(0)) { + return socket.getLocalPort(); + } + } + + private static void waitForServer(int port) throws Exception { + int maxRetries = 100; + for (int i = 0; i < maxRetries; i++) { + try { + HttpURLConnection conn = + (HttpURLConnection) new URL("http://localhost:" + port + "/health").openConnection(); + conn.setRequestMethod("GET"); + conn.setConnectTimeout(200); + conn.setReadTimeout(200); + if (conn.getResponseCode() == 200) { + return; + } + } catch (IOException e) { + // Server not ready yet. + } + Thread.sleep(100); + } + fail("Server did not start within timeout"); + } +} diff --git a/plugins/spring/src/test/java/com/google/genkit/plugins/spring/AgentHeaderPropagationTest.java b/plugins/spring/src/test/java/com/google/genkit/plugins/spring/AgentHeaderPropagationTest.java new file mode 100644 index 000000000..8cef05163 --- /dev/null +++ b/plugins/spring/src/test/java/com/google/genkit/plugins/spring/AgentHeaderPropagationTest.java @@ -0,0 +1,209 @@ +/* + * 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. + * + * SPDX-License-Identifier: Apache-2.0 + */ + +package com.google.genkit.plugins.spring; + +import static org.junit.jupiter.api.Assertions.*; + +import com.google.genkit.ai.Message; +import com.google.genkit.ai.agent.AgentChat; +import com.google.genkit.ai.agent.AgentFinishReason; +import com.google.genkit.ai.agent.AgentResponse; +import com.google.genkit.ai.agent.AgentResult; +import com.google.genkit.ai.agent.CustomAgentConfig; +import com.google.genkit.ai.agent.InMemorySessionStore; +import com.google.genkit.ai.agent.internal.AgentActions; +import com.google.genkit.client.RemoteAgent; +import com.google.genkit.client.RemoteAgentOptions; +import com.google.genkit.core.ActionContext; +import com.google.genkit.core.DefaultRegistry; +import java.io.IOException; +import java.net.HttpURLConnection; +import java.net.URI; +import java.net.URL; +import java.util.HashMap; +import java.util.Map; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.Test; + +/** + * End-to-end proof that custom HTTP request headers set via {@link RemoteAgentOptions#headers()} on + * the client actually reach a served {@code AgentFn} through {@code GenkitAgentController}. + * + *

Mirrors {@code plugins/jetty}'s {@code AgentHeaderPropagationTest}, proving that the same + * reserved {@code "headers"} context key populated by {@code JettyPlugin} is populated identically + * by {@code GenkitAgentController} — a tool/flow reading {@code ctx.getContext().get("headers")} + * behaves the same whether served by Jetty or Spring. + * + *

This is the full client-to-server round trip that was previously broken: {@code + * HttpAgentTransport} already sent {@code RemoteAgentOptions.headers()} as real HTTP headers, but + * {@code GenkitAgentController} never read incoming HTTP headers on the server side, so a tool/flow + * had no way to observe them. These tests start a real Spring server, send a real HTTP request (via + * {@link RemoteAgent}/{@link RemoteAgentOptions#headers()}) carrying a custom header, and assert a + * no-model custom {@code AgentFn} can read that header's value out of {@code + * AgentFnContext#context()} and reflect it back in its response. + */ +class AgentHeaderPropagationTest { + + private SpringPlugin plugin; + + @AfterEach + void tearDown() throws Exception { + if (plugin != null) { + plugin.stop(); + } + } + + /** + * Starts a Spring server on {@code port} hosting a server-managed, no-model custom agent named + * {@code headerEchoAgent}. The agent's {@code AgentFn} reads {@code + * fnCtx.context().getContext().get("headers")} (the reserved key that {@code + * GenkitAgentController} merges incoming HTTP headers into) and echoes the requested header's + * value back in its reply text, so the test can assert on it without any model call. + */ + private void startHeaderEchoAgent(int port) throws Exception { + DefaultRegistry registry = new DefaultRegistry(); + + AgentActions.defineCustomAgent( + registry, + CustomAgentConfig.>builder() + .name("headerEchoAgent") + .store(new InMemorySessionStore<>()) + .build(), + (sess, fnCtx) -> { + String headerValue = "(missing)"; + ActionContext actionContext = fnCtx.context(); + if (actionContext != null && actionContext.getContext() != null) { + Object headersObj = actionContext.getContext().get("headers"); + if (headersObj instanceof Map headers) { + Object v = headers.get("X-Custom-Auth"); + if (v != null) { + headerValue = String.valueOf(v); + } + } + } + return AgentResult.builder() + .message(Message.model("header=" + headerValue)) + .finishReason(AgentFinishReason.STOP) + .build(); + }); + + SpringPluginOptions options = SpringPluginOptions.builder().port(port).build(); + plugin = new SpringPlugin(options); + plugin.init(registry); + + Thread serverThread = + new Thread( + () -> { + try { + plugin.start(); + } catch (Exception e) { + // Server stopped. + } + }); + serverThread.setDaemon(true); + serverThread.start(); + + waitForServer(port); + } + + /** + * Full round trip: a {@link RemoteAgent} configured with {@link + * RemoteAgentOptions.Builder#headers(Map)} sends a real HTTP POST to a live Spring server; the + * server-side {@code AgentFn} reads the header back out of {@code ActionContext.getContext()} and + * reflects it in its reply. Proves the client-to-server header pipe end-to-end, not just that a + * map got populated somewhere. + */ + @Test + void testCustomHeaderReachesAgentFnOverHttp() throws Exception { + int port = findAvailablePort(); + startHeaderEchoAgent(port); + + Map headers = new HashMap<>(); + headers.put("X-Custom-Auth", "secret-token-123"); + + AgentChat> chat = + RemoteAgent.chat( + RemoteAgentOptions.builder() + .url("http://localhost:" + port + "/headerEchoAgent") + .headers(headers) + .build()); + + AgentResponse> resp = chat.send("hello"); + + assertNotNull(resp.text(), "expected non-null text"); + assertEquals( + "header=secret-token-123", + resp.text(), + "expected the AgentFn to read the custom HTTP header via ActionContext.getContext()"); + } + + /** + * Without any custom header configured, the AgentFn should observe the reserved {@code "headers"} + * key as either absent or not containing {@code X-Custom-Auth} — i.e. the header plumbing does + * not fabricate values, and the (missing) fallback path is exercised. + */ + @Test + void testNoCustomHeaderMeansMissingInAgentFn() throws Exception { + int port = findAvailablePort(); + startHeaderEchoAgent(port); + + AgentChat> chat = + RemoteAgent.chat( + RemoteAgentOptions.builder() + .url("http://localhost:" + port + "/headerEchoAgent") + .build()); + + AgentResponse> resp = chat.send("hello"); + + assertEquals( + "header=(missing)", + resp.text(), + "expected no X-Custom-Auth header to be observed when the client sets none"); + } + + // --------------------------------------------------------------------------- + // Helpers + // --------------------------------------------------------------------------- + + private static int findAvailablePort() throws IOException { + try (java.net.ServerSocket socket = new java.net.ServerSocket(0)) { + return socket.getLocalPort(); + } + } + + private static void waitForServer(int port) throws Exception { + int maxRetries = 100; + for (int i = 0; i < maxRetries; i++) { + try { + URL healthUrl = new URI("http://localhost:" + port + "/health").toURL(); + HttpURLConnection conn = (HttpURLConnection) healthUrl.openConnection(); + conn.setRequestMethod("GET"); + conn.setConnectTimeout(200); + conn.setReadTimeout(200); + if (conn.getResponseCode() == 200) { + return; + } + } catch (IOException e) { + // Server not ready yet. + } + Thread.sleep(100); + } + fail("Server did not start within timeout"); + } +} diff --git a/plugins/spring/src/test/java/com/google/genkit/plugins/spring/AgentHttpServingTest.java b/plugins/spring/src/test/java/com/google/genkit/plugins/spring/AgentHttpServingTest.java new file mode 100644 index 000000000..f78d0eef8 --- /dev/null +++ b/plugins/spring/src/test/java/com/google/genkit/plugins/spring/AgentHttpServingTest.java @@ -0,0 +1,248 @@ +/* + * 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. + * + * SPDX-License-Identifier: Apache-2.0 + */ + +package com.google.genkit.plugins.spring; + +import static org.junit.jupiter.api.Assertions.*; + +import com.fasterxml.jackson.databind.JsonNode; +import com.fasterxml.jackson.databind.ObjectMapper; +import com.fasterxml.jackson.databind.node.ObjectNode; +import com.google.genkit.core.ActionDef; +import com.google.genkit.core.ActionType; +import com.google.genkit.core.BidiActionImpl; +import com.google.genkit.core.DefaultRegistry; +import java.io.IOException; +import java.net.HttpURLConnection; +import java.net.URI; +import java.net.URL; +import java.net.http.HttpClient; +import java.net.http.HttpRequest; +import java.net.http.HttpResponse; +import java.util.Optional; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.Test; + +/** + * Integration tests for agent HTTP serving in the Spring plugin. + * + *

Mirrors {@code plugins/jetty}'s {@code AgentHttpServingTest}: verifies the + * one-turn-per-request transport for {@code ActionType.AGENT} bidi actions served by {@link + * GenkitAgentController} — a non-streaming {@code {data, init}} request, an SSE streaming request, + * and a companion {@code getSnapshot} endpoint — proving wire-format parity between the Spring and + * Jetty server plugins. + */ +class AgentHttpServingTest { + + private static final ObjectMapper MAPPER = new ObjectMapper(); + + private SpringPlugin plugin; + + @AfterEach + void tearDown() throws Exception { + if (plugin != null) { + plugin.stop(); + } + } + + /** + * Builds a registry containing a server-managed bidi agent named {@code chatAgent} plus a + * companion {@code agent-snapshot} action, then starts a SpringPlugin on the given port. + */ + private void startWithChatAgent(int port) throws Exception { + DefaultRegistry registry = new DefaultRegistry(); + + // A minimal server-managed bidi agent: reads one input, emits one stream chunk, returns a + // final output containing snapshotId + an echoed message. + BidiActionImpl agent = + BidiActionImpl.builder() + .name("chatAgent") + .inputClass(JsonNode.class) + .outputClass(JsonNode.class) + .streamClass(JsonNode.class) + .initClass(JsonNode.class) + .handler( + (ctx, init, inputs, cb) -> { + Optional first = inputs.next(); + JsonNode data = first.orElse(MAPPER.nullNode()); + // The AgentInput envelope carries the turn message under "message". + JsonNode message = data.path("message"); + // Emit one streamed chunk. + if (cb != null) { + ObjectNode chunk = MAPPER.createObjectNode(); + chunk.put("text", "thinking..."); + cb.accept(chunk); + } + // Final output echoes the message back. + ObjectNode result = MAPPER.createObjectNode(); + result.put("snapshotId", "s1"); + result.set("message", message); + return result; + }) + .build(); + agent.register(registry); + + // Companion agent-snapshot action looked up at key "/agent-snapshot/chatAgent". + ActionDef snapshot = + new ActionDef<>( + "chatAgent", + ActionType.AGENT_SNAPSHOT, + null, + null, + JsonNode.class, + JsonNode.class, + (ctx, input, cb) -> { + ObjectNode snap = MAPPER.createObjectNode(); + snap.put("snapshotId", input != null ? input.path("snapshotId").asText("?") : "?"); + snap.put("status", "captured"); + return snap; + }); + snapshot.register(registry); + + SpringPluginOptions options = SpringPluginOptions.builder().port(port).build(); + plugin = new SpringPlugin(options); + plugin.init(registry); + + Thread serverThread = + new Thread( + () -> { + try { + plugin.start(); + } catch (Exception e) { + // Server stopped. + } + }); + serverThread.setDaemon(true); + serverThread.start(); + + waitForServer(port); + } + + @Test + void testAgentNonStreaming() throws Exception { + int port = findAvailablePort(); + startWithChatAgent(port); + + String body = + "{\"data\":{\"message\":{\"role\":\"user\",\"content\":[{\"text\":\"hi\"}]}},\"init\":{}}"; + + HttpClient client = HttpClient.newHttpClient(); + HttpRequest request = + HttpRequest.newBuilder() + .uri(URI.create("http://localhost:" + port + "/chatAgent")) + .header("Content-Type", "application/json") + .POST(HttpRequest.BodyPublishers.ofString(body)) + .build(); + + HttpResponse response = client.send(request, HttpResponse.BodyHandlers.ofString()); + assertEquals(200, response.statusCode(), "body: " + response.body()); + + JsonNode root = MAPPER.readTree(response.body()); + assertTrue(root.has("result"), "expected result envelope: " + response.body()); + assertEquals("s1", root.path("result").path("snapshotId").asText()); + // Echoed message should be present. + assertEquals( + "hi", root.path("result").path("message").path("content").path(0).path("text").asText()); + } + + @Test + void testAgentStreamingSse() throws Exception { + int port = findAvailablePort(); + startWithChatAgent(port); + + String body = + "{\"data\":{\"message\":{\"role\":\"user\",\"content\":[{\"text\":\"hi\"}]}},\"init\":{}}"; + + HttpClient client = HttpClient.newHttpClient(); + HttpRequest request = + HttpRequest.newBuilder() + .uri(URI.create("http://localhost:" + port + "/chatAgent")) + .header("Content-Type", "application/json") + .header("Accept", "text/event-stream") + .POST(HttpRequest.BodyPublishers.ofString(body)) + .build(); + + HttpResponse response = client.send(request, HttpResponse.BodyHandlers.ofString()); + assertEquals(200, response.statusCode(), "body: " + response.body()); + + String contentType = response.headers().firstValue("Content-Type").orElse(""); + assertTrue(contentType.contains("text/event-stream"), "content-type: " + contentType); + assertTrue( + response.headers().firstValue("X-Genkit-Stream-Id").isPresent(), + "expected X-Genkit-Stream-Id header"); + + String text = response.body(); + // A message frame then a result frame. + assertTrue(text.contains("\"message\""), "expected a message frame: " + text); + assertTrue(text.contains("thinking..."), "expected streamed chunk text: " + text); + assertTrue(text.contains("\"result\""), "expected a result frame: " + text); + assertTrue(text.contains("\"snapshotId\":\"s1\""), "expected snapshotId in result: " + text); + // Frames are SSE-formatted. + assertTrue(text.contains("data:"), "expected SSE data prefix: " + text); + } + + @Test + void testAgentSnapshotCompanion() throws Exception { + int port = findAvailablePort(); + startWithChatAgent(port); + + String body = "{\"data\":{\"snapshotId\":\"s1\"}}"; + + HttpClient client = HttpClient.newHttpClient(); + HttpRequest request = + HttpRequest.newBuilder() + .uri(URI.create("http://localhost:" + port + "/chatAgent/getSnapshot")) + .header("Content-Type", "application/json") + .POST(HttpRequest.BodyPublishers.ofString(body)) + .build(); + + HttpResponse response = client.send(request, HttpResponse.BodyHandlers.ofString()); + assertEquals(200, response.statusCode(), "body: " + response.body()); + + JsonNode root = MAPPER.readTree(response.body()); + assertTrue(root.has("result"), "expected result envelope: " + response.body()); + assertEquals("s1", root.path("result").path("snapshotId").asText()); + assertEquals("captured", root.path("result").path("status").asText()); + } + + private static int findAvailablePort() throws IOException { + try (java.net.ServerSocket socket = new java.net.ServerSocket(0)) { + return socket.getLocalPort(); + } + } + + private static void waitForServer(int port) throws Exception { + int maxRetries = 100; + for (int i = 0; i < maxRetries; i++) { + try { + HttpURLConnection conn = + (HttpURLConnection) new URL("http://localhost:" + port + "/health").openConnection(); + conn.setRequestMethod("GET"); + conn.setConnectTimeout(200); + conn.setReadTimeout(200); + if (conn.getResponseCode() == 200) { + return; + } + } catch (IOException e) { + // Server not ready yet. + } + Thread.sleep(100); + } + fail("Server did not start within timeout"); + } +} diff --git a/plugins/spring/src/test/java/com/google/genkit/plugins/spring/RemoteAgentTest.java b/plugins/spring/src/test/java/com/google/genkit/plugins/spring/RemoteAgentTest.java new file mode 100644 index 000000000..50331b3fd --- /dev/null +++ b/plugins/spring/src/test/java/com/google/genkit/plugins/spring/RemoteAgentTest.java @@ -0,0 +1,492 @@ +/* + * 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. + * + * SPDX-License-Identifier: Apache-2.0 + */ + +package com.google.genkit.plugins.spring; + +import static org.junit.jupiter.api.Assertions.*; + +import com.fasterxml.jackson.databind.JsonNode; +import com.fasterxml.jackson.databind.ObjectMapper; +import com.fasterxml.jackson.databind.node.ObjectNode; +import com.google.genkit.ai.Message; +import com.google.genkit.ai.agent.AgentChat; +import com.google.genkit.ai.agent.AgentFinishReason; +import com.google.genkit.ai.agent.AgentResponse; +import com.google.genkit.ai.agent.AgentResult; +import com.google.genkit.ai.agent.CustomAgentConfig; +import com.google.genkit.ai.agent.GetSnapshotRequest; +import com.google.genkit.ai.agent.InMemorySessionStore; +import com.google.genkit.ai.agent.SessionSnapshot; +import com.google.genkit.ai.agent.internal.AgentActions; +import com.google.genkit.client.HttpAgentTransport; +import com.google.genkit.client.RemoteAgent; +import com.google.genkit.client.RemoteAgentOptions; +import com.google.genkit.core.ActionDef; +import com.google.genkit.core.ActionType; +import com.google.genkit.core.BidiActionImpl; +import com.google.genkit.core.DefaultRegistry; +import java.io.IOException; +import java.net.HttpURLConnection; +import java.net.URI; +import java.net.URL; +import java.util.ArrayList; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.Optional; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicInteger; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.Test; + +/** + * Integration tests for {@link RemoteAgent} / {@link HttpAgentTransport} against a live Spring + * server running a server-managed echo agent. + * + *

Mirrors {@code plugins/jetty}'s {@code RemoteAgentTest}, proving that the same {@code + * com.google.genkit.client.RemoteAgent} / {@code AgentChat} client used against a Jetty server also + * works unmodified against a {@link SpringPlugin}-backed server — i.e. true cross-plugin wire + * compatibility for the agent HTTP contract. + * + *

Placed in the {@code plugins/spring} module because that module depends on {@code genkit} + * (which contains {@code RemoteAgent}) and can also start a {@link SpringPlugin} — giving us both + * halves of the test in one module without a circular dependency. + */ +class RemoteAgentTest { + + private static final ObjectMapper MAPPER = new ObjectMapper(); + + private SpringPlugin plugin; + + @AfterEach + void tearDown() throws Exception { + if (plugin != null) { + plugin.stop(); + } + } + + // --------------------------------------------------------------------------- + // Server setup + // --------------------------------------------------------------------------- + + /** + * Starts a Spring server on {@code port} hosting a server-managed echo agent named {@code + * echoAgent}. + * + *

The agent: + * + *

    + *
  • Emits one SSE chunk with {@code {"text":"streaming..."}}. + *
  • Returns a final output whose {@code snapshotId} is {@code "snap-"} (turn counter + * increments on each invocation), {@code sessionId} is {@code "session-1"}, and {@code + * message} echoes the user message content back. + *
+ * + *

A companion {@code agent-snapshot} action is registered so {@code getSnapshot} works. + */ + private void startEchoAgent(int port) throws Exception { + DefaultRegistry registry = new DefaultRegistry(); + + AtomicInteger turnCounter = new AtomicInteger(0); + + BidiActionImpl agent = + BidiActionImpl.builder() + .name("echoAgent") + .inputClass(JsonNode.class) + .outputClass(JsonNode.class) + .streamClass(JsonNode.class) + .initClass(JsonNode.class) + .handler( + (ctx, init, inputs, cb) -> { + Optional first = inputs.next(); + JsonNode data = first.orElse(MAPPER.nullNode()); + JsonNode messageNode = data.path("message"); + + int turn = turnCounter.incrementAndGet(); + + // Emit one streamed chunk. + if (cb != null) { + ObjectNode chunk = MAPPER.createObjectNode(); + chunk.put("text", "streaming..."); + cb.accept(chunk); + } + + // Build final output. + ObjectNode result = MAPPER.createObjectNode(); + result.put("snapshotId", "snap-" + turn); + result.put("sessionId", "session-1"); + result.put("finishReason", AgentFinishReason.STOP.getValue()); + // Echo message back. + ObjectNode message = MAPPER.createObjectNode(); + message.put("role", "model"); + ObjectNode part = MAPPER.createObjectNode(); + // Echo the first text part of the user message. + String userText = + messageNode.path("content").path(0).path("text").asText("(empty)"); + part.put("text", "echo: " + userText); + message.putArray("content").add(part); + result.set("message", message); + return result; + }) + .build(); + agent.register(registry); + + // Companion snapshot action. + ActionDef snapshot = + new ActionDef<>( + "echoAgent", + ActionType.AGENT_SNAPSHOT, + null, + null, + JsonNode.class, + JsonNode.class, + (ctx, input, cb) -> { + ObjectNode snap = MAPPER.createObjectNode(); + String sid = + (input != null && input.has("snapshotId")) + ? input.get("snapshotId").asText("?") + : "?"; + snap.put("snapshotId", sid); + snap.put("sessionId", "session-1"); + snap.put("status", "completed"); + return snap; + }); + snapshot.register(registry); + + SpringPluginOptions options = SpringPluginOptions.builder().port(port).build(); + plugin = new SpringPlugin(options); + plugin.init(registry); + + Thread serverThread = + new Thread( + () -> { + try { + plugin.start(); + } catch (Exception e) { + // Server stopped. + } + }); + serverThread.setDaemon(true); + serverThread.start(); + + waitForServer(port); + } + + // --------------------------------------------------------------------------- + // Tests + // --------------------------------------------------------------------------- + + @Test + void testSendFirstTurn() throws Exception { + int port = findAvailablePort(); + startEchoAgent(port); + + AgentChat> chat = + RemoteAgent.chat( + RemoteAgentOptions.builder().url("http://localhost:" + port + "/echoAgent").build()); + + AgentResponse> resp = chat.send("hello"); + + // Response text must be non-empty (echoed user message). + assertNotNull(resp.text(), "expected non-null text"); + assertFalse(resp.text().isEmpty(), "expected non-empty text, got: " + resp.text()); + assertTrue(resp.text().contains("echo: hello"), "expected echo in text, got: " + resp.text()); + + // snapshotId must be set after the first turn. + assertNotNull(chat.snapshotId(), "expected snapshotId to be set after first turn"); + assertEquals("snap-1", chat.snapshotId()); + } + + @Test + void testSendSecondTurnResumes() throws Exception { + int port = findAvailablePort(); + startEchoAgent(port); + + AgentChat> chat = + RemoteAgent.chat( + RemoteAgentOptions.builder().url("http://localhost:" + port + "/echoAgent").build()); + + chat.send("first"); + AgentResponse> resp2 = chat.send("second"); + + // Second turn should have incremented the counter. + assertEquals("snap-2", chat.snapshotId()); + assertTrue( + resp2.text().contains("echo: second"), "expected echo of second turn: " + resp2.text()); + } + + @Test + void testStreamingChunksDelivered() throws Exception { + int port = findAvailablePort(); + startEchoAgent(port); + + AgentChat> chat = + RemoteAgent.chat( + RemoteAgentOptions.builder().url("http://localhost:" + port + "/echoAgent").build()); + + List chunks = new ArrayList<>(); + chat.sendStream( + "streaming test", + chunk -> { + // Chunks may contain a modelChunk; collect whatever the server sends. + chunks.add(chunk.toString()); + }); + + // The server emits one streaming chunk; the onChunk callback must have been called. + // (AgentStreamChunk is not a plain text — it wraps modelChunk; we just verify it fired.) + assertFalse(chunks.isEmpty(), "expected at least one SSE chunk"); + } + + @Test + void testGetSnapshot() throws Exception { + int port = findAvailablePort(); + startEchoAgent(port); + + AgentChat> chat = + RemoteAgent.chat( + RemoteAgentOptions.builder().url("http://localhost:" + port + "/echoAgent").build()); + + chat.send("hi"); + String snapId = chat.snapshotId(); + assertNotNull(snapId, "expected snapshotId after send"); + + // Retrieve snapshot via the transport directly. + HttpAgentTransport> transport = + new HttpAgentTransport<>( + RemoteAgentOptions.builder().url("http://localhost:" + port + "/echoAgent").build()); + + SessionSnapshot> snap = + transport.getSnapshot(GetSnapshotRequest.builder().snapshotId(snapId).build()); + + assertNotNull(snap, "expected non-null snapshot"); + assertEquals(snapId, snap.getSnapshotId(), "snapshot ID should match requested ID"); + } + + /** + * Starts a Spring server on {@code port} hosting a server-managed custom agent named {@code + * blockingAgent} whose {@code AgentFn} blocks on {@code releaseLatch} until released. Used to + * open a deterministic {@code PENDING} window for {@link #testAbortOverHttp()}. + */ + private void startBlockingAgent(int port, CountDownLatch releaseLatch) throws Exception { + DefaultRegistry registry = new DefaultRegistry(); + + AgentActions.defineCustomAgent( + registry, + CustomAgentConfig.>builder() + .name("blockingAgent") + .store(new InMemorySessionStore<>()) + .build(), + (sess, fnCtx) -> { + releaseLatch.await(10, TimeUnit.SECONDS); + return AgentResult.builder() + .message(Message.model("done")) + .finishReason(AgentFinishReason.STOP) + .build(); + }); + + SpringPluginOptions options = SpringPluginOptions.builder().port(port).build(); + plugin = new SpringPlugin(options); + plugin.init(registry); + + Thread serverThread = + new Thread( + () -> { + try { + plugin.start(); + } catch (Exception e) { + // Server stopped. + } + }); + serverThread.setDaemon(true); + serverThread.start(); + + waitForServer(port); + } + + /** + * Starts a Spring server on {@code port} hosting a client-managed (no store) custom agent named + * {@code counterAgent} that increments an integer counter in its custom state each turn and + * echoes back the running total in its reply text. + */ + private void startClientManagedCounterAgent(int port) throws Exception { + DefaultRegistry registry = new DefaultRegistry(); + + AgentActions.defineCustomAgent( + registry, + CustomAgentConfig.>builder() + .name("counterAgent") + // No store() call => client-managed: state round-trips via AgentInit/AgentOutput. + .build(), + (sess, fnCtx) -> { + Map custom = sess.getCustom(); + int count = 0; + if (custom != null && custom.get("count") instanceof Number) { + count = ((Number) custom.get("count")).intValue(); + } + count++; + Map updated = new HashMap<>(); + updated.put("count", count); + sess.updateCustom(c -> updated); + return AgentResult.builder() + .message(Message.model("count=" + count)) + .finishReason(AgentFinishReason.STOP) + .build(); + }); + + SpringPluginOptions options = SpringPluginOptions.builder().port(port).build(); + plugin = new SpringPlugin(options); + plugin.init(registry); + + Thread serverThread = + new Thread( + () -> { + try { + plugin.start(); + } catch (Exception e) { + // Server stopped. + } + }); + serverThread.setDaemon(true); + serverThread.start(); + + waitForServer(port); + } + + /** + * Verifies {@link AgentChat#abort()} (backed by {@link HttpAgentTransport#abort}) reaches the + * server's {@code /abort} companion endpoint and that a subsequent {@code getSnapshot} (via the + * transport directly, mirroring how {@code AgentChat.loadChat} would read it back) reflects the + * flipped status. + * + *

Scoped to what is actually implemented: aborting a snapshot that is genuinely {@code + * PENDING} (opened by a detached turn whose {@code AgentFn} we block deterministically) — not + * interrupting a live foreground call, which {@code Agent.abort()} / the abort companion action + * do not support (no cancellation hook is threaded into a running {@code AgentFn}). + */ + @Test + void testAbortOverHttp() throws Exception { + int port = findAvailablePort(); + CountDownLatch releaseLatch = new CountDownLatch(1); + startBlockingAgent(port, releaseLatch); + + RemoteAgentOptions opts = + RemoteAgentOptions.builder().url("http://localhost:" + port + "/blockingAgent").build(); + AgentChat> chat = RemoteAgent.chat(opts); + + // Detach so the snapshot is written PENDING and the background AgentFn blocks on the latch. + chat.sendStream( + com.google.genkit.ai.agent.AgentInput.builder() + .message(Message.user("go")) + .detach(true) + .build(), + c -> {}); + + String snapshotId = chat.snapshotId(); + assertNotNull(snapshotId, "expected snapshotId after detach"); + + // Abort via the client — reaches the server's /abort companion endpoint. + com.google.genkit.ai.agent.SnapshotStatus status = chat.abort(); + assertEquals( + com.google.genkit.ai.agent.SnapshotStatus.ABORTED, + status, + "expected client-side abort() to report ABORTED"); + + // Read the snapshot back directly via the transport (as AgentChat.loadChat would internally). + HttpAgentTransport> transport = new HttpAgentTransport<>(opts); + SessionSnapshot> snap = + transport.getSnapshot(GetSnapshotRequest.builder().snapshotId(snapshotId).build()); + assertNotNull(snap, "expected non-null snapshot"); + assertEquals( + com.google.genkit.ai.agent.SnapshotStatus.ABORTED, + snap.getStatus(), + "expected stored snapshot status to be ABORTED"); + + // Release the blocked background turn; its finalizer must not revert the ABORTED status. + releaseLatch.countDown(); + Thread.sleep(300); + SessionSnapshot> snapAfter = + transport.getSnapshot(GetSnapshotRequest.builder().snapshotId(snapshotId).build()); + assertEquals( + com.google.genkit.ai.agent.SnapshotStatus.ABORTED, + snapAfter.getStatus(), + "expected status to remain ABORTED after the background turn finalizes"); + } + + /** + * Verifies a client-managed agent (no {@code SessionStore}) genuinely round-trips its session + * state over real HTTP: {@link AgentChat} carries the full {@link + * com.google.genkit.ai.agent.SessionState} (including custom state) in {@code AgentInit} on each + * turn, and the served agent's reply on turn 2 reflects state seeded by turn 1 — proving + * serialization/deserialization across the wire, not just in-process object sharing. + */ + @Test + void testClientManagedRemoteAgent() throws Exception { + int port = findAvailablePort(); + startClientManagedCounterAgent(port); + + AgentChat> chat = + RemoteAgent.chat( + RemoteAgentOptions.builder() + .url("http://localhost:" + port + "/counterAgent") + .serverManaged(false) + .build()); + + AgentResponse> resp1 = chat.send("first"); + assertEquals("count=1", resp1.text(), "expected first turn to start the counter at 1"); + + AgentResponse> resp2 = chat.send("second"); + assertEquals( + "count=2", + resp2.text(), + "expected second turn's reply to reflect state seeded from turn 1 over the wire"); + + AgentResponse> resp3 = chat.send("third"); + assertEquals("count=3", resp3.text(), "expected third turn to continue incrementing"); + } + + // --------------------------------------------------------------------------- + // Helpers + // --------------------------------------------------------------------------- + + private static int findAvailablePort() throws IOException { + try (java.net.ServerSocket socket = new java.net.ServerSocket(0)) { + return socket.getLocalPort(); + } + } + + private static void waitForServer(int port) throws Exception { + int maxRetries = 100; + for (int i = 0; i < maxRetries; i++) { + try { + URL healthUrl = new URI("http://localhost:" + port + "/health").toURL(); + HttpURLConnection conn = (HttpURLConnection) healthUrl.openConnection(); + conn.setRequestMethod("GET"); + conn.setConnectTimeout(200); + conn.setReadTimeout(200); + if (conn.getResponseCode() == 200) { + return; + } + } catch (IOException e) { + // Server not ready yet. + } + Thread.sleep(100); + } + fail("Server did not start within timeout"); + } +} diff --git a/pom.xml b/pom.xml index ede78d90b..ad5798a91 100644 --- a/pom.xml +++ b/pom.xml @@ -84,6 +84,7 @@ plugins/localvec plugins/mcp plugins/firebase + plugins/middleware plugins/weaviate plugins/postgresql plugins/pinecone @@ -106,9 +107,15 @@ samples/middleware samples/middleware-v2 samples/mcp - samples/chat-session - samples/multi-agent samples/interrupts + samples/agents-weather + samples/agents-orchestrator + samples/agents-stateless + samples/agents-remote + samples/agents-human-in-the-loop + samples/agents-firestore-session + samples/agents-dynamodb-session + samples/agents-cosmos-session samples/spring samples/firebase samples/weaviate @@ -135,6 +142,10 @@ 1.5.37 12.1.6 1.63.0 + + 4.33.2 + 1.48.0 6.1.1 5.23.0 4.38.0 @@ -185,6 +196,11 @@ jackson-datatype-jsr310 ${jackson.version} + + com.fasterxml.jackson.dataformat + jackson-dataformat-yaml + ${jackson.version} + @@ -309,6 +325,31 @@ ${okhttp.version} test + + + + com.google.protobuf + protobuf-java + ${protobuf.version} + + + com.google.auth + google-auth-library-oauth2-http + ${google-auth.version} + + + com.google.auth + google-auth-library-credentials + ${google-auth.version} + diff --git a/samples/README.md b/samples/README.md index e61701209..f258a0715 100644 --- a/samples/README.md +++ b/samples/README.md @@ -77,6 +77,10 @@ The Dev UI will be available at `http://localhost:4000` and allows you to: | [weaviate](./weaviate) | Weaviate vector database RAG sample | `OPENAI_API_KEY` | | [postgresql](./postgresql) | PostgreSQL pgvector RAG sample | `OPENAI_API_KEY` | | [pinecone](./pinecone) | Pinecone vector database RAG sample | `OPENAI_API_KEY` + `PINECONE_API_KEY` | +| [agents-human-in-the-loop](./agents-human-in-the-loop) | Agent interrupts + human-in-the-loop resume | `GEMINI_API_KEY` | +| [agents-firestore-session](./agents-firestore-session) | Agent session persistence backed by Firestore | `GEMINI_API_KEY` + Firestore | +| [agents-dynamodb-session](./agents-dynamodb-session) | Agent session persistence backed by DynamoDB | AWS credentials + DynamoDB | +| [agents-cosmos-session](./agents-cosmos-session) | Agent session persistence backed by Azure Cosmos DB | Azure credentials + Cosmos DB | ## Sample Details @@ -413,6 +417,59 @@ export PINECONE_API_KEY=your-key ./run.sh ``` +### Agents Human-in-the-Loop Sample + +A banking assistant agent that pauses on a money transfer and waits for approval: +- Interrupt tool via `genkit.defineInterrupt(...)` +- Detecting `INTERRUPTED` turns and resolving them with `AgentChat.resume(...)` +- Server-managed state via `FileSessionStore` + +```bash +cd java/samples/agents-human-in-the-loop +export GEMINI_API_KEY=your-key +mvn -q exec:java # CLI demo: prompts you to approve/reject the transfer +``` + +### Agents Firestore Session Sample + +A server-managed agent whose sessions are persisted in Firestore: +- `FirestoreSessionStore` wired via `AgentConfig.store(...)` +- Two-turn conversation, then reads the snapshot back from Firestore +- Runs against the Firestore emulator or a real project + +```bash +cd java/samples/agents-firestore-session +export GEMINI_API_KEY=your-key +export FIRESTORE_EMULATOR_HOST=localhost:8080 # or configure GCLOUD_PROJECT +mvn -q exec:java -Dexec.args=demo # or `./run.sh` to serve over HTTP +``` + +### Agents DynamoDB Session Sample + +A server-managed agent whose sessions are persisted in DynamoDB, using an AWS Bedrock model: +- `DynamoDbSessionStore` wired via `AgentConfig.store(...)` +- Runs against DynamoDB Local or real AWS + +```bash +cd java/samples/agents-dynamodb-session +export DYNAMODB_LOCAL_ENDPOINT=http://localhost:8000 # docker run -p 8000:8000 amazon/dynamodb-local +export AWS_REGION=us-east-1 +mvn -q exec:java -Dexec.args=demo # or `./run.sh` to serve over HTTP +``` + +### Agents Cosmos DB Session Sample + +A server-managed agent whose sessions are persisted in Azure Cosmos DB, using an Azure AI Foundry model: +- `CosmosSessionStore` wired via `AgentConfig.store(...)` +- Runs against the Cosmos DB emulator or a real account + +```bash +cd java/samples/agents-cosmos-session +export COSMOS_ENDPOINT=https://localhost:8081 +export COSMOS_KEY=your-key +mvn -q exec:java -Dexec.args=demo # or `./run.sh` to serve over HTTP +``` + ## Building All Samples From the Java root directory: diff --git a/samples/agents-cosmos-session/README.md b/samples/agents-cosmos-session/README.md new file mode 100644 index 000000000..1294954db --- /dev/null +++ b/samples/agents-cosmos-session/README.md @@ -0,0 +1,92 @@ +# Agents: Cosmos DB Session Sample + +A server-managed assistant agent whose conversation state is persisted in **Azure Cosmos DB** via `CosmosSessionStore`, using an **Azure AI Foundry** model for generation — a self-contained, single-cloud sample. + +## Prerequisites + +- Java 21+ and Maven 3.6+ +- A Cosmos DB backend — the emulator for dev, or an Azure account +- Azure AI Foundry endpoint + credentials (for live model calls) + +### Run the Cosmos DB emulator (dev) + +Use the **next-gen Linux emulator** (`vnext-latest`) — it's ARM-native and runs on Apple Silicon. (The older `azure-cosmos-emulator:latest` image is x86-only with no `arm64` build, so it fails to pull on M-series Macs.) + +**1. Start the emulator with HTTPS enabled.** It defaults to HTTP, but the Java Cosmos SDK requires HTTPS. We keep the Data Explorer on HTTP so it opens in the browser without a cert warning: + +```bash +docker run --detach \ + --publish 8081:8081 --publish 8080:8080 --publish 1234:1234 \ + --name cosmos-emulator \ + mcr.microsoft.com/cosmosdb/linux/azure-cosmos-emulator:vnext-latest \ + --protocol https --explorer-protocol http + +# wait until ready (8080 is the health probe) +until curl -sf http://localhost:8080/ready >/dev/null; do sleep 1; done; echo "ready" +``` + +Ports: `8081` = Cosmos endpoint (HTTPS), `1234` = Data Explorer, `8080` = health probe. The emulator supports the NoSQL API in **gateway mode**, which is exactly what `CosmosSessionStore` uses. + +**2. Trust the emulator's TLS certificate in your JVM** (required — the Java SDK rejects the self-signed cert otherwise): + +```bash +openssl s_client -connect localhost:8081 /tmp/cosmos_emulator.cert +keytool -cacerts -importcert -alias cosmos_emulator -file /tmp/cosmos_emulator.cert \ + -storepass changeit -noprompt +# remove later with: keytool -cacerts -delete -alias cosmos_emulator -storepass changeit +``` + +`keytool -cacerts` targets the truststore of the `keytool` on your `PATH` — make sure it's the same JDK you run the sample with (`which java`). + +**3. Point the sample at it** (the key below is the emulator's fixed well-known key — not a secret): + +```bash +export COSMOS_ENDPOINT=https://localhost:8081 +export COSMOS_KEY='' + +# For live model calls (Azure OpenAI). Use the *.openai.azure.com endpoint (NOT a +# services.ai.azure.com /openai/v1/... URL — the plugin speaks legacy chat/completions): +export AZURE_AI_FOUNDRY_ENDPOINT=https://.openai.azure.com +export AZURE_API_KEY= +# Your Azure OpenAI *deployment* name (Azure routes by deployment, not model name): +export AZURE_MODEL= +# Optional: override the api-version if your resource rejects the plugin default +# (2024-10-01-preview) with 400 "API version not supported" — e.g. a GA resource: +export AZURE_API_VERSION=2024-10-21 +``` + +The agent's model is `azure-foundry/${AZURE_MODEL}` (default `gpt-4o-mini`). `AZURE_MODEL` must be an existing **deployment** name in your Azure OpenAI resource — the sample registers it automatically. + +With `COSMOS_ENDPOINT`/`COSMOS_KEY` set, the database (`genkit`) and container (`genkit-sessions`, partition key `/pk`) are created automatically. + +### Inspecting the stored data + +Open the **Data Explorer** at `http://localhost:1234` → database `genkit` → container `genkit-sessions`. Documents are keyed by `id`: `SNAP_` (snapshot metadata), `SHARD__` (the state JSON), and `PTR_` (the current-leaf pointer). Real Azure accounts show the same under the portal's Data Explorer. + +## Run + +**Serve over HTTP (default):** + +```bash +mvn -q exec:java +# assistant -> POST http://localhost:8082/assistant +``` + +**Genkit Dev UI:** + +```bash +genkit start -- mvn -q exec:java +``` + +**Persistence demo** (requires Cosmos DB + Azure AI Foundry credentials): + +```bash +mvn -q exec:java -Dexec.args=demo +``` + +The demo runs a two-turn conversation, then reads the latest snapshot **back from Cosmos DB** by session id and prints how many messages were persisted. + +## Configuration + +`CosmosSessionStore` uses a sharded checkpoint + diff + pointer layout in a single container. Tune it with `CosmosSessionStoreOptions` (database/container names, checkpoint interval, shard size — default 1 MiB, under the 2 MB document cap — per-tenant prefix, `createIfNotExists`). diff --git a/samples/agents-cosmos-session/pom.xml b/samples/agents-cosmos-session/pom.xml new file mode 100644 index 000000000..453869e98 --- /dev/null +++ b/samples/agents-cosmos-session/pom.xml @@ -0,0 +1,82 @@ + + + + 4.0.0 + + + com.google.genkit + genkit-parent + 1.0.0-SNAPSHOT + ../../pom.xml + + + com.google.genkit.samples + genkit-sample-agents-cosmos-session + jar + Genkit Agents Cosmos DB Session Sample + Sample demonstrating agent session persistence backed by Azure Cosmos DB + + + UTF-8 + 21 + 21 + 1.0.0-SNAPSHOT + true + true + + + + + com.google.genkit + genkit + ${genkit.version} + + + com.google.genkit + genkit-plugin-azure-foundry + ${genkit.version} + + + com.google.genkit + genkit-plugin-jetty + ${genkit.version} + + + ch.qos.logback + logback-classic + 1.5.37 + + + + + + + org.codehaus.mojo + exec-maven-plugin + 3.6.3 + + com.google.genkit.samples.CosmosSessionAgentApp + + + + + diff --git a/samples/agents-cosmos-session/run.sh b/samples/agents-cosmos-session/run.sh new file mode 100755 index 000000000..d10fb3e2a --- /dev/null +++ b/samples/agents-cosmos-session/run.sh @@ -0,0 +1,7 @@ +#!/bin/bash + +# Genkit Sample Runner +# This script runs the sample application + +echo "Building and running the sample..." +mvn clean compile exec:java diff --git a/samples/agents-cosmos-session/src/main/java/com/google/genkit/samples/CosmosSessionAgentApp.java b/samples/agents-cosmos-session/src/main/java/com/google/genkit/samples/CosmosSessionAgentApp.java new file mode 100644 index 000000000..033d62b50 --- /dev/null +++ b/samples/agents-cosmos-session/src/main/java/com/google/genkit/samples/CosmosSessionAgentApp.java @@ -0,0 +1,223 @@ +/* + * 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. + * + * SPDX-License-Identifier: Apache-2.0 + */ + +package com.google.genkit.samples; + +import com.azure.cosmos.CosmosClient; +import com.azure.cosmos.CosmosClientBuilder; +import com.google.genkit.Genkit; +import com.google.genkit.GenkitOptions; +import com.google.genkit.agent.AgentConfig; +import com.google.genkit.ai.agent.Agent; +import com.google.genkit.ai.agent.AgentChat; +import com.google.genkit.ai.agent.AgentResponse; +import com.google.genkit.ai.agent.GetSnapshotOptions; +import com.google.genkit.ai.agent.SessionSnapshot; +import com.google.genkit.ai.agent.SessionStore; +import com.google.genkit.plugins.azurefoundry.AzureFoundryPlugin; +import com.google.genkit.plugins.azurefoundry.AzureFoundryPluginOptions; +import com.google.genkit.plugins.azurefoundry.session.CosmosSessionStore; +import com.google.genkit.plugins.azurefoundry.session.CosmosSessionStoreOptions; +import com.google.genkit.plugins.jetty.JettyPlugin; +import java.util.Map; + +/** + * Agent session persistence backed by Azure Cosmos DB. + * + *

Defines a server-managed assistant whose conversation snapshots are stored in Cosmos DB via + * {@link CosmosSessionStore}, and uses an Azure AI Foundry model for generation — a self-contained, + * single-cloud sample. + * + *

Point it at the Cosmos DB emulator + * or a real account: + * + * 

+ *   export COSMOS_ENDPOINT=https://localhost:8081
+ *   export COSMOS_KEY=<emulator-or-account-key>
+ *
+ * + *

Live model calls additionally require {@code AZURE_AI_FOUNDRY_ENDPOINT} (and credentials). + * + *

Two run modes: + * + *

    + *
  • Serve (default) — starts Jetty to expose the agent over HTTP and keep the process + * alive; also discoverable in the Genkit Dev UI. {@code POST + * http://localhost:8082/assistant}. + *
  • Demo — runs a two-turn conversation then reads the persisted snapshot back from + * Cosmos DB to prove server-side persistence: {@code mvn -q exec:java -Dexec.args=demo}. + *
+ */ +public class CosmosSessionAgentApp { + + public static void main(String[] args) throws Exception { + // The model reference for Azure OpenAI is the *deployment* name. Set it to your deployment via + // AZURE_MODEL. Defaults to "gpt-4o-mini". + String model = azureModel(); + + // ── 1. Build Genkit with the beta agents API + Azure AI Foundry model ──── + Genkit genkit = + Genkit.builder() + .options(GenkitOptions.builder().experimental(true).devMode(true).build()) + .plugin(buildAzurePlugin(model)) + .build(); + + // ── 2. Build the Cosmos-backed session store (null if not configured) ──── + SessionStore> store = buildStore(); + + // ── 3. Define a server-managed agent using the store ──────────────────── + AgentConfig.Builder> cfg = + AgentConfig.>builder() + .name("assistant") + .description("A helpful assistant with Cosmos DB-backed memory") + .system( + "You are a helpful assistant. Keep answers concise and remember what the user tells you.") + .model("azure-foundry/" + model); + if (store != null) { + cfg.store(store); + } + Agent> assistant = genkit.beta().defineAgent(cfg.build()); + + // ── 4. Default mode: serve over HTTP + keep the process alive ─────────── + boolean runDemo = args.length > 0 && "demo".equalsIgnoreCase(args[0]); + if (!runDemo) { + serve(genkit); + return; + } + + // ── 5. Demo mode: two turns + read-back proof ──────────────────────────── + if (store == null) { + System.out.println( + "Cosmos DB is not configured (set COSMOS_ENDPOINT and COSMOS_KEY) — " + + "cannot demonstrate server-side persistence."); + return; + } + try { + System.out.println("=== Cosmos DB-backed session agent ==="); + AgentChat> chat = assistant.chat(); + + AgentResponse> r1 = + chat.send("My name is Ada Lovelace. Please remember it."); + System.out.println("Turn 1: " + r1.text()); + AgentResponse> r2 = chat.send("What is my name?"); + System.out.println("Turn 2: " + r2.text()); + + String sessionId = chat.sessionId(); + System.out.println("Session id: " + sessionId); + + SessionSnapshot> persisted = + store.getSnapshot(GetSnapshotOptions.builder().sessionId(sessionId).build()); + int messageCount = + (persisted != null + && persisted.getState() != null + && persisted.getState().getMessages() != null) + ? persisted.getState().getMessages().size() + : 0; + System.out.println("Messages persisted in Cosmos DB for this session: " + messageCount); + } catch (Exception e) { + System.err.println( + "Demo failed. It needs a reachable Cosmos DB and Azure AI Foundry credentials: " + + e.getMessage()); + } + } + + /** + * Returns the Azure model/deployment name from {@code AZURE_MODEL}, defaulting to gpt-4o-mini. + */ + private static String azureModel() { + String m = System.getenv("AZURE_MODEL"); + return (m != null && !m.isBlank()) ? m : "gpt-4o-mini"; + } + + /** + * Builds the Azure AI Foundry plugin. The endpoint comes from {@code AZURE_AI_FOUNDRY_ENDPOINT}; + * a placeholder is used when it is unset so the sample still boots in serve mode (live calls then + * require a real endpoint + key). The {@code model} (deployment) is registered as a custom model + * when it isn't one of the plugin's built-in names. + */ + private static AzureFoundryPlugin buildAzurePlugin(String model) { + String endpoint = System.getenv("AZURE_AI_FOUNDRY_ENDPOINT"); + String apiKey = System.getenv("AZURE_API_KEY"); + String apiVersion = System.getenv("AZURE_API_VERSION"); // optional; must match your resource + if (endpoint == null || endpoint.isBlank()) { + endpoint = "https://placeholder.openai.azure.com"; + } + AzureFoundryPluginOptions.Builder options = + AzureFoundryPluginOptions.builder() + .endpoint(endpoint) + .apiKey(apiKey != null ? apiKey : "placeholder-key"); + // Azure returns 400 "API version not supported" when the api-version doesn't match the + // resource. The plugin default is 2024-10-01-preview; override it via AZURE_API_VERSION + // (e.g. 2024-10-21 for a GA Azure OpenAI resource) if your endpoint rejects the default. + if (apiVersion != null && !apiVersion.isBlank()) { + options.apiVersion(apiVersion); + } + AzureFoundryPlugin plugin = new AzureFoundryPlugin(options.build()); + // Azure OpenAI routes by deployment name; register non-built-in deployment names as custom + // models (guarding against double-registration of a built-in name). + if (!AzureFoundryPlugin.SUPPORTED_MODELS.contains(model)) { + plugin.customModel(model); + } + return plugin; + } + + /** + * Builds a Cosmos DB-backed session store, or returns {@code null} (running client-managed) when + * Cosmos DB is not configured / cannot be reached. + */ + private static SessionStore> buildStore() { + String endpoint = System.getenv("COSMOS_ENDPOINT"); + String key = System.getenv("COSMOS_KEY"); + if (endpoint == null || endpoint.isBlank() || key == null || key.isBlank()) { + System.out.println( + "Cosmos DB not configured (set COSMOS_ENDPOINT and COSMOS_KEY); running client-managed."); + return null; + } + try { + CosmosClient client = + new CosmosClientBuilder().endpoint(endpoint).key(key).gatewayMode().buildClient(); + return new CosmosSessionStore<>( + client, CosmosSessionStoreOptions.builder().createIfNotExists(true).build()); + } catch (Exception e) { + System.out.println( + "Cosmos DB not reachable (" + e.getMessage() + "); running client-managed."); + return null; + } + } + + /** Starts Jetty to serve the agent over HTTP and blocks until the process is stopped. */ + private static void serve(Genkit genkit) throws Exception { + int port = System.getenv("PORT") != null ? Integer.parseInt(System.getenv("PORT")) : 8082; + JettyPlugin jetty = JettyPlugin.create(port); + jetty.init(genkit.getRegistry()); + System.out.println("Serving the assistant on http://localhost:" + port); + System.out.println(" assistant -> POST http://localhost:" + port + "/assistant"); + System.out.println( + "Tip: run under `genkit start -- mvn -q exec:java`" + + " to open the Dev UI, or pass `demo` for the persistence demo."); + try { + jetty.start(); // blocks until the process is stopped (Ctrl-C) + } catch (Exception e) { + System.err.println( + "ERROR: could not start the HTTP server on port " + port + ": " + e.getMessage()); + System.err.println("Port " + port + " is likely already in use. Free it or set PORT."); + throw e; + } + } +} diff --git a/samples/agents-cosmos-session/src/main/resources/logback.xml b/samples/agents-cosmos-session/src/main/resources/logback.xml new file mode 100644 index 000000000..073738862 --- /dev/null +++ b/samples/agents-cosmos-session/src/main/resources/logback.xml @@ -0,0 +1,24 @@ + + + + + + %d{HH:mm:ss.SSS} %-5level %logger{24} - %msg%n + + + + + + + + + + + + + + + + + + diff --git a/samples/agents-dynamodb-session/README.md b/samples/agents-dynamodb-session/README.md new file mode 100644 index 000000000..d4d3a3b36 --- /dev/null +++ b/samples/agents-dynamodb-session/README.md @@ -0,0 +1,47 @@ +# Agents: DynamoDB Session Sample + +A server-managed assistant agent whose conversation state is persisted in **Amazon DynamoDB** via `DynamoDbSessionStore`, using an **AWS Bedrock** model for generation — a self-contained, single-cloud sample. + +## Prerequisites + +- Java 21+ and Maven 3.6+ +- AWS credentials (for the Bedrock model, and for DynamoDB unless using DynamoDB Local) +- A DynamoDB backend — DynamoDB Local for dev, or real AWS + +### Point at DynamoDB Local (dev) + +```bash +docker run -p 8000:8000 amazon/dynamodb-local +export DYNAMODB_LOCAL_ENDPOINT=http://localhost:8000 +export AWS_REGION=us-east-1 +``` + +With `DYNAMODB_LOCAL_ENDPOINT` set, the table (`genkit-sessions`) is created automatically. Against real AWS, create the table beforehand (partition key `pk` string, sort key `sk` string) or run once with credentials that allow `CreateTable`. + +## Run + +**Serve over HTTP (default)** — starts without touching AWS: + +```bash +mvn -q exec:java +# assistant -> POST http://localhost:8080/assistant +``` + +**Genkit Dev UI:** + +```bash +genkit start -- mvn -q exec:java +``` + +**Persistence demo** (requires AWS credentials for Bedrock + a reachable DynamoDB): + +```bash +export DYNAMODB_LOCAL_ENDPOINT=http://localhost:8000 # or use real AWS +mvn -q exec:java -Dexec.args=demo +``` + +The demo runs a two-turn conversation, then reads the latest snapshot **back from DynamoDB** by session id and prints how many messages were persisted. + +## Configuration + +`DynamoDbSessionStore` uses a sharded checkpoint + diff + pointer layout in a single table. Tune it with `DynamoDbSessionStoreOptions` (table name, checkpoint interval, shard size — default 350 KiB, under the 400 KB item cap — per-tenant prefix, `createTableIfNotExists`). diff --git a/samples/chat-session/pom.xml b/samples/agents-dynamodb-session/pom.xml similarity index 87% rename from samples/chat-session/pom.xml rename to samples/agents-dynamodb-session/pom.xml index 2169404b3..8e5bcaeb3 100644 --- a/samples/chat-session/pom.xml +++ b/samples/agents-dynamodb-session/pom.xml @@ -30,10 +30,10 @@ com.google.genkit.samples - genkit-sample-chat-session + genkit-sample-agents-dynamodb-session jar - Genkit Chat Session Sample - Sample application demonstrating session-based multi-turn chat with persistence + Genkit Agents DynamoDB Session Sample + Sample demonstrating agent session persistence backed by DynamoDB UTF-8 @@ -52,7 +52,7 @@ com.google.genkit - genkit-plugin-openai + genkit-plugin-aws-bedrock ${genkit.version} @@ -74,7 +74,7 @@ exec-maven-plugin 3.6.3 - com.google.genkit.samples.ChatSessionApp + com.google.genkit.samples.DynamoDbSessionAgentApp diff --git a/samples/agents-dynamodb-session/run.sh b/samples/agents-dynamodb-session/run.sh new file mode 100755 index 000000000..d10fb3e2a --- /dev/null +++ b/samples/agents-dynamodb-session/run.sh @@ -0,0 +1,7 @@ +#!/bin/bash + +# Genkit Sample Runner +# This script runs the sample application + +echo "Building and running the sample..." +mvn clean compile exec:java diff --git a/samples/agents-dynamodb-session/src/main/java/com/google/genkit/samples/DynamoDbSessionAgentApp.java b/samples/agents-dynamodb-session/src/main/java/com/google/genkit/samples/DynamoDbSessionAgentApp.java new file mode 100644 index 000000000..bee942b34 --- /dev/null +++ b/samples/agents-dynamodb-session/src/main/java/com/google/genkit/samples/DynamoDbSessionAgentApp.java @@ -0,0 +1,200 @@ +/* + * 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. + * + * SPDX-License-Identifier: Apache-2.0 + */ + +package com.google.genkit.samples; + +import com.google.genkit.Genkit; +import com.google.genkit.GenkitOptions; +import com.google.genkit.agent.AgentConfig; +import com.google.genkit.ai.agent.Agent; +import com.google.genkit.ai.agent.AgentChat; +import com.google.genkit.ai.agent.AgentResponse; +import com.google.genkit.ai.agent.GetSnapshotOptions; +import com.google.genkit.ai.agent.SessionSnapshot; +import com.google.genkit.ai.agent.SessionStore; +import com.google.genkit.plugins.awsbedrock.AwsBedrockPlugin; +import com.google.genkit.plugins.awsbedrock.session.DynamoDbSessionStore; +import com.google.genkit.plugins.awsbedrock.session.DynamoDbSessionStoreOptions; +import com.google.genkit.plugins.jetty.JettyPlugin; +import java.net.URI; +import java.util.Map; +import software.amazon.awssdk.auth.credentials.AwsBasicCredentials; +import software.amazon.awssdk.auth.credentials.StaticCredentialsProvider; +import software.amazon.awssdk.regions.Region; +import software.amazon.awssdk.services.dynamodb.DynamoDbClient; + +/** + * Agent session persistence backed by Amazon DynamoDB. + * + *

Defines a server-managed assistant whose conversation snapshots are stored in DynamoDB via + * {@link DynamoDbSessionStore}, and uses an AWS Bedrock model for generation — a self-contained, + * single-cloud sample. + * + *

Point it at DynamoDB Local for + * development: + * + * 

+ *   docker run -p 8000:8000 amazon/dynamodb-local
+ *   export DYNAMODB_LOCAL_ENDPOINT=http://localhost:8000
+ *
+ * + * or at real AWS via the default credential chain (the {@code genkit-sessions} table must exist). + * Live model calls always require AWS credentials for Bedrock. + * + *

Two run modes: + * + *

    + *
  • Serve (default) — starts Jetty to expose the agent over HTTP and keep the process + * alive; also discoverable in the Genkit Dev UI. {@code POST + * http://localhost:8080/assistant}. + *
  • Demo — runs a two-turn conversation then reads the persisted snapshot back from + * DynamoDB to prove server-side persistence: {@code mvn -q -pl + * samples/agents-dynamodb-session exec:java -Dexec.args=demo}. + *
+ */ +public class DynamoDbSessionAgentApp { + + public static void main(String[] args) throws Exception { + // ── 1. Build Genkit with the beta agents API + AWS Bedrock model plugin ── + // + // Newer Claude models on Bedrock (Sonnet 4.6, Opus 4.x, Sonnet 5, ...) cannot be invoked with + // on-demand throughput — they require a cross-region *inference profile* ID (a geo prefix like + // "us.", "eu.", "jp.", "au."). We register the US profile ID as a custom model so it can be + // referenced below. Change the prefix to match your region, or use a base model that supports + // on-demand invocation in your region. + Genkit genkit = + Genkit.builder() + .options(GenkitOptions.builder().experimental(true).devMode(true).build()) + .plugin(AwsBedrockPlugin.create().customModel("us.anthropic.claude-sonnet-4-6")) + .build(); + + // ── 2. Build the DynamoDB-backed session store (null if not configured) ── + SessionStore> store = buildStore(); + + // ── 3. Define a server-managed agent using the store ──────────────────── + AgentConfig.Builder> cfg = + AgentConfig.>builder() + .name("assistant") + .description("A helpful assistant with DynamoDB-backed memory") + .system( + "You are a helpful assistant. Keep answers concise and remember what the user tells you.") + .model("aws-bedrock/us.anthropic.claude-sonnet-4-6"); + if (store != null) { + cfg.store(store); + } + Agent> assistant = genkit.beta().defineAgent(cfg.build()); + + // ── 4. Default mode: serve over HTTP + keep the process alive ─────────── + boolean runDemo = args.length > 0 && "demo".equalsIgnoreCase(args[0]); + if (!runDemo) { + serve(genkit); + return; + } + + // ── 5. Demo mode: two turns + read-back proof ──────────────────────────── + if (store == null) { + System.out.println( + "DynamoDB is not configured (set DYNAMODB_LOCAL_ENDPOINT for local, or AWS credentials" + + " with an existing table) — cannot demonstrate server-side persistence."); + return; + } + try { + System.out.println("=== DynamoDB-backed session agent ==="); + AgentChat> chat = assistant.chat(); + + AgentResponse> r1 = + chat.send("My name is Ada Lovelace. Please remember it."); + System.out.println("Turn 1: " + r1.text()); + AgentResponse> r2 = chat.send("What is my name?"); + System.out.println("Turn 2: " + r2.text()); + + String sessionId = chat.sessionId(); + System.out.println("Session id: " + sessionId); + + SessionSnapshot> persisted = + store.getSnapshot(GetSnapshotOptions.builder().sessionId(sessionId).build()); + int messageCount = + (persisted != null + && persisted.getState() != null + && persisted.getState().getMessages() != null) + ? persisted.getState().getMessages().size() + : 0; + System.out.println("Messages persisted in DynamoDB for this session: " + messageCount); + } catch (Exception e) { + System.err.println( + "Demo failed. It needs AWS credentials for Bedrock and a reachable DynamoDB " + + "(set DYNAMODB_LOCAL_ENDPOINT for local): " + + e.getMessage()); + } + } + + /** + * Builds a DynamoDB-backed session store. Building the client is offline; the table is + * auto-created only when running against DynamoDB Local ({@code DYNAMODB_LOCAL_ENDPOINT} set), so + * serve mode starts without touching AWS. + */ + private static SessionStore> buildStore() { + try { + String localEndpoint = System.getenv("DYNAMODB_LOCAL_ENDPOINT"); + String region = System.getenv("AWS_REGION"); + if (region == null || region.isBlank()) { + region = "us-east-1"; + } + DynamoDbClient client; + boolean autoCreate = false; + if (localEndpoint != null && !localEndpoint.isBlank()) { + client = + DynamoDbClient.builder() + .region(Region.of(region)) + .endpointOverride(URI.create(localEndpoint)) + .credentialsProvider( + StaticCredentialsProvider.create(AwsBasicCredentials.create("local", "local"))) + .build(); + autoCreate = true; // safe to auto-create against DynamoDB Local + } else { + client = DynamoDbClient.builder().region(Region.of(region)).build(); + } + return new DynamoDbSessionStore<>( + client, DynamoDbSessionStoreOptions.builder().createTableIfNotExists(autoCreate).build()); + } catch (Exception e) { + System.out.println( + "DynamoDB not configured (" + e.getMessage() + "); running client-managed."); + return null; + } + } + + /** Starts Jetty to serve the agent over HTTP and blocks until the process is stopped. */ + private static void serve(Genkit genkit) throws Exception { + int port = System.getenv("PORT") != null ? Integer.parseInt(System.getenv("PORT")) : 8080; + JettyPlugin jetty = JettyPlugin.create(port); + jetty.init(genkit.getRegistry()); + System.out.println("Serving the assistant on http://localhost:" + port); + System.out.println(" assistant -> POST http://localhost:" + port + "/assistant"); + System.out.println( + "Tip: run under `genkit start -- mvn -q exec:java`" + + " to open the Dev UI, or pass `demo` for the persistence demo."); + try { + jetty.start(); // blocks until the process is stopped (Ctrl-C) + } catch (Exception e) { + System.err.println( + "ERROR: could not start the HTTP server on port " + port + ": " + e.getMessage()); + System.err.println("Port " + port + " is likely already in use. Free it or set PORT."); + throw e; + } + } +} diff --git a/samples/agents-dynamodb-session/src/main/resources/logback.xml b/samples/agents-dynamodb-session/src/main/resources/logback.xml new file mode 100644 index 000000000..ce538eacb --- /dev/null +++ b/samples/agents-dynamodb-session/src/main/resources/logback.xml @@ -0,0 +1,24 @@ + + + + + + %d{HH:mm:ss.SSS} %-5level %logger{24} - %msg%n + + + + + + + + + + + + + + + + + + diff --git a/samples/agents-firestore-session/README.md b/samples/agents-firestore-session/README.md new file mode 100644 index 000000000..cbbb49312 --- /dev/null +++ b/samples/agents-firestore-session/README.md @@ -0,0 +1,47 @@ +# Agents: Firestore Session Sample + +A server-managed assistant agent whose conversation state is persisted in **Firestore** via `FirestoreSessionStore`. Because the state lives in Firestore rather than in the process, the conversation survives restarts and can be resumed from any instance. + +## Prerequisites + +- Java 21+ and Maven 3.6+ +- `GEMINI_API_KEY` (for live model calls in demo mode) +- A Firestore backend — the emulator for local dev, or a real project + +### Point at the Firestore emulator (local dev) + +```bash +gcloud emulators firestore start --host-port=localhost:8080 +export FIRESTORE_EMULATOR_HOST=localhost:8080 +export GCLOUD_PROJECT=demo-genkit +``` + +Or use a real project by setting `GCLOUD_PROJECT` (with application default credentials configured). + +## Run + +**Serve over HTTP (default):** + +```bash +mvn -q exec:java +# assistant -> POST http://localhost:8080/assistant +``` + +**Genkit Dev UI:** + +```bash +genkit start -- mvn -q exec:java +``` + +**Persistence demo** (requires `GEMINI_API_KEY` + a configured Firestore): + +```bash +export GEMINI_API_KEY=your-key +mvn -q exec:java -Dexec.args=demo +``` + +The demo runs a two-turn conversation, then reads the latest snapshot **back from Firestore** by session id and prints how many messages were persisted — proving the state is stored server-side. + +## Configuration + +`FirestoreSessionStore` uses a sharded checkpoint + diff + pointer layout. Tune it with `FirestoreSessionStoreOptions` (collection name, checkpoint interval, shard size, per-tenant prefix). diff --git a/samples/agents-firestore-session/pom.xml b/samples/agents-firestore-session/pom.xml new file mode 100644 index 000000000..44fafe284 --- /dev/null +++ b/samples/agents-firestore-session/pom.xml @@ -0,0 +1,87 @@ + + + + 4.0.0 + + + com.google.genkit + genkit-parent + 1.0.0-SNAPSHOT + ../../pom.xml + + + com.google.genkit.samples + genkit-sample-agents-firestore-session + jar + Genkit Agents Firestore Session Sample + Sample demonstrating agent session persistence backed by Firestore + + + UTF-8 + 21 + 21 + 1.0.0-SNAPSHOT + true + true + + + + + com.google.genkit + genkit + ${genkit.version} + + + com.google.genkit + genkit-plugin-google-genai + ${genkit.version} + + + com.google.genkit + genkit-plugin-firebase + ${genkit.version} + + + com.google.genkit + genkit-plugin-jetty + ${genkit.version} + + + ch.qos.logback + logback-classic + 1.5.37 + + + + + + + org.codehaus.mojo + exec-maven-plugin + 3.6.3 + + com.google.genkit.samples.FirestoreSessionAgentApp + + + + + diff --git a/samples/agents-firestore-session/run.sh b/samples/agents-firestore-session/run.sh new file mode 100755 index 000000000..d10fb3e2a --- /dev/null +++ b/samples/agents-firestore-session/run.sh @@ -0,0 +1,7 @@ +#!/bin/bash + +# Genkit Sample Runner +# This script runs the sample application + +echo "Building and running the sample..." +mvn clean compile exec:java diff --git a/samples/agents-firestore-session/src/main/java/com/google/genkit/samples/FirestoreSessionAgentApp.java b/samples/agents-firestore-session/src/main/java/com/google/genkit/samples/FirestoreSessionAgentApp.java new file mode 100644 index 000000000..4f751133a --- /dev/null +++ b/samples/agents-firestore-session/src/main/java/com/google/genkit/samples/FirestoreSessionAgentApp.java @@ -0,0 +1,196 @@ +/* + * 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. + * + * SPDX-License-Identifier: Apache-2.0 + */ + +package com.google.genkit.samples; + +import com.google.cloud.firestore.Firestore; +import com.google.cloud.firestore.FirestoreOptions; +import com.google.genkit.Genkit; +import com.google.genkit.GenkitOptions; +import com.google.genkit.agent.AgentConfig; +import com.google.genkit.ai.agent.Agent; +import com.google.genkit.ai.agent.AgentChat; +import com.google.genkit.ai.agent.AgentResponse; +import com.google.genkit.ai.agent.GetSnapshotOptions; +import com.google.genkit.ai.agent.SessionSnapshot; +import com.google.genkit.ai.agent.SessionStore; +import com.google.genkit.plugins.firebase.session.FirestoreSessionStore; +import com.google.genkit.plugins.googlegenai.GoogleGenAIPlugin; +import com.google.genkit.plugins.jetty.JettyPlugin; +import java.util.Map; + +/** + * Agent session persistence backed by Firestore. + * + *

Defines a server-managed assistant whose conversation snapshots are stored in Firestore via + * {@link FirestoreSessionStore}. Because the state lives in Firestore (not in the process), the + * conversation survives restarts and can be resumed from any instance. + * + *

Point it at the Firestore emulator for local development: + * + * 

+ *   gcloud emulators firestore start --host-port=localhost:8080
+ *   export FIRESTORE_EMULATOR_HOST=localhost:8080
+ *
+ * + * or at a real project by setting {@code GCLOUD_PROJECT} (and application default credentials). + * + *

Two run modes: + * + *

    + *
  • Serve (default) — starts Jetty to expose the agent over HTTP and keep the process + * alive; also discoverable in the Genkit Dev UI. {@code POST + * http://localhost:8080/assistant}. + *
  • Demo — runs a two-turn conversation then reads the persisted snapshot back from + * Firestore to prove server-side persistence. Requires {@code GEMINI_API_KEY} and a + * configured Firestore: {@code mvn -q exec:java -Dexec.args=demo}. + *
+ */ +public class FirestoreSessionAgentApp { + + public static void main(String[] args) throws Exception { + // ── 1. Build Genkit with the beta agents API enabled ──────────────────── + Genkit genkit = + Genkit.builder() + .options(GenkitOptions.builder().experimental(true).devMode(true).build()) + .plugin(GoogleGenAIPlugin.create()) + .build(); + + // ── 2. Build the Firestore-backed session store (null if not configured) ─ + SessionStore> store = buildStore(); + + // ── 3. Define a server-managed agent using the store ──────────────────── + AgentConfig.Builder> cfg = + AgentConfig.>builder() + .name("assistant") + .description("A helpful assistant with Firestore-backed memory") + .system( + "You are a helpful assistant. Keep answers concise and remember what the user tells you.") + .model("googleai/gemini-2.5-flash"); + if (store != null) { + cfg.store(store); + } + Agent> assistant = genkit.beta().defineAgent(cfg.build()); + + // ── 4. Default mode: serve over HTTP + keep the process alive ─────────── + boolean runDemo = args.length > 0 && "demo".equalsIgnoreCase(args[0]); + if (!runDemo) { + serve(genkit); + return; + } + + // ── 5. Demo mode: two turns + read-back proof (needs GEMINI_API_KEY) ───── + String apiKey = System.getenv("GEMINI_API_KEY"); + if (apiKey == null || apiKey.isBlank()) { + System.out.println( + "GEMINI_API_KEY is not set — agent defined successfully but skipping live calls."); + return; + } + if (store == null) { + System.out.println( + "Firestore is not configured (set FIRESTORE_EMULATOR_HOST or GCLOUD_PROJECT) — " + + "cannot demonstrate server-side persistence."); + return; + } + + System.out.println("=== Firestore-backed session agent ==="); + AgentChat> chat = assistant.chat(); + + AgentResponse> r1 = + chat.send("My name is Ada Lovelace. Please remember it."); + System.out.println("Turn 1: " + r1.text()); + AgentResponse> r2 = chat.send("What is my name?"); + System.out.println("Turn 2: " + r2.text()); + + String sessionId = chat.sessionId(); + System.out.println("Session id: " + sessionId); + + // Read the latest snapshot straight from Firestore to prove it was persisted. + SessionSnapshot> persisted = + store.getSnapshot(GetSnapshotOptions.builder().sessionId(sessionId).build()); + int messageCount = + (persisted != null + && persisted.getState() != null + && persisted.getState().getMessages() != null) + ? persisted.getState().getMessages().size() + : 0; + System.out.println("Messages persisted in Firestore for this session: " + messageCount); + } + + /** + * Builds a Firestore-backed session store, or returns {@code null} (running client-managed) when + * Firestore is not configured. Construction is gated on a config signal because building a + * Firestore client eagerly resolves application-default credentials (which can block probing the + * metadata server when nothing is configured), so serve mode stays fast without configuration. + */ + private static SessionStore> buildStore() { + String emulatorHost = System.getenv("FIRESTORE_EMULATOR_HOST"); + String projectId = System.getenv("GCLOUD_PROJECT"); + if (projectId == null || projectId.isBlank()) { + projectId = System.getenv("GOOGLE_CLOUD_PROJECT"); + } + boolean hasEmulator = emulatorHost != null && !emulatorHost.isBlank(); + boolean hasProject = projectId != null && !projectId.isBlank(); + if (!hasEmulator && !hasProject) { + System.out.println( + "Firestore not configured (set FIRESTORE_EMULATOR_HOST or GCLOUD_PROJECT); running" + + " client-managed."); + return null; + } + try { + String effectiveProject = hasProject ? projectId : "demo-genkit"; + Firestore firestore; + if (hasEmulator) { + firestore = + FirestoreOptions.getDefaultInstance().toBuilder() + .setProjectId(effectiveProject) + .setEmulatorHost(emulatorHost) + .build() + .getService(); + } else { + firestore = + FirestoreOptions.newBuilder().setProjectId(effectiveProject).build().getService(); + } + return new FirestoreSessionStore<>(firestore); + } catch (Exception e) { + System.out.println( + "Firestore not reachable (" + e.getMessage() + "); running client-managed."); + return null; + } + } + + /** Starts Jetty to serve the agent over HTTP and blocks until the process is stopped. */ + private static void serve(Genkit genkit) throws Exception { + int port = System.getenv("PORT") != null ? Integer.parseInt(System.getenv("PORT")) : 8080; + JettyPlugin jetty = JettyPlugin.create(port); + jetty.init(genkit.getRegistry()); + System.out.println("Serving the assistant on http://localhost:" + port); + System.out.println(" assistant -> POST http://localhost:" + port + "/assistant"); + System.out.println( + "Tip: run under `genkit start -- mvn -q exec:java`" + + " to open the Dev UI, or pass `demo` for the persistence demo."); + try { + jetty.start(); // blocks until the process is stopped (Ctrl-C) + } catch (Exception e) { + System.err.println( + "ERROR: could not start the HTTP server on port " + port + ": " + e.getMessage()); + System.err.println("Port " + port + " is likely already in use. Free it or set PORT."); + throw e; + } + } +} diff --git a/samples/agents-firestore-session/src/main/resources/logback.xml b/samples/agents-firestore-session/src/main/resources/logback.xml new file mode 100644 index 000000000..ea1ff8c1b --- /dev/null +++ b/samples/agents-firestore-session/src/main/resources/logback.xml @@ -0,0 +1,23 @@ + + + + + + %d{HH:mm:ss.SSS} %-5level %logger{24} - %msg%n + + + + + + + + + + + + + + + + + diff --git a/samples/agents-human-in-the-loop/README.md b/samples/agents-human-in-the-loop/README.md new file mode 100644 index 000000000..d847ce860 --- /dev/null +++ b/samples/agents-human-in-the-loop/README.md @@ -0,0 +1,32 @@ +# Agents: Human-in-the-Loop (Interrupts) Sample + +A command-line banking assistant agent that **pauses on a sensitive action** — a money transfer — and waits for you to approve or reject it before continuing. + +It demonstrates the agent-level interrupt/resume flow: + +- A tool created with `genkit.defineInterrupt(...)` pauses the turn instead of executing. +- The turn finishes with `AgentFinishReason.INTERRUPTED`, and `AgentResponse.interrupts()` surfaces the pending tool request. +- The caller resolves it with `tool.respond(interrupt.part(), output)` and resumes the turn with `AgentChat.resume(...)`. + +Conversation state is persisted server-side with `FileSessionStore` (under `./.snapshots`). + +## Prerequisites + +- Java 21+ and Maven 3.6+ +- `GEMINI_API_KEY` (for live model calls) + +## Run + +```bash +export GEMINI_API_KEY=your-key +mvn -q exec:java +``` + +The demo asks the agent to transfer money, prints the pending confirmation, prompts you for `yes`/`no` on the command line, then resumes the turn with your decision. Without `GEMINI_API_KEY` set, the agent is defined but live calls are skipped. + +## How it works + +1. `send("Transfer $150 to Alice for dinner")` → the model calls `confirmTransfer`, which interrupts the turn. +2. `response.finishReason()` is `INTERRUPTED`; `response.interrupts().get(0)` is the pending tool request. +3. `confirmTransfer.respond(interrupt.part(), new ConfirmationOutput(approved, ...))` builds the response part. +4. `chat.resume(List.of(respond))` resumes the turn; the agent finishes with a normal (`STOP`) response. diff --git a/samples/agents-human-in-the-loop/pom.xml b/samples/agents-human-in-the-loop/pom.xml new file mode 100644 index 000000000..12a368de5 --- /dev/null +++ b/samples/agents-human-in-the-loop/pom.xml @@ -0,0 +1,77 @@ + + + + 4.0.0 + + + com.google.genkit + genkit-parent + 1.0.0-SNAPSHOT + ../../pom.xml + + + com.google.genkit.samples + genkit-sample-agents-human-in-the-loop + jar + Genkit Agents Human-in-the-Loop Sample + Sample demonstrating agent interrupts and human-in-the-loop resume + + + UTF-8 + 21 + 21 + 1.0.0-SNAPSHOT + true + true + + + + + com.google.genkit + genkit + ${genkit.version} + + + com.google.genkit + genkit-plugin-google-genai + ${genkit.version} + + + ch.qos.logback + logback-classic + 1.5.37 + + + + + + + org.codehaus.mojo + exec-maven-plugin + 3.6.3 + + com.google.genkit.samples.HumanInTheLoopAgentApp + + + + + diff --git a/samples/agents-human-in-the-loop/run.sh b/samples/agents-human-in-the-loop/run.sh new file mode 100755 index 000000000..d10fb3e2a --- /dev/null +++ b/samples/agents-human-in-the-loop/run.sh @@ -0,0 +1,7 @@ +#!/bin/bash + +# Genkit Sample Runner +# This script runs the sample application + +echo "Building and running the sample..." +mvn clean compile exec:java diff --git a/samples/agents-human-in-the-loop/src/main/java/com/google/genkit/samples/HumanInTheLoopAgentApp.java b/samples/agents-human-in-the-loop/src/main/java/com/google/genkit/samples/HumanInTheLoopAgentApp.java new file mode 100644 index 000000000..e56a0ef4e --- /dev/null +++ b/samples/agents-human-in-the-loop/src/main/java/com/google/genkit/samples/HumanInTheLoopAgentApp.java @@ -0,0 +1,233 @@ +/* + * 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. + * + * SPDX-License-Identifier: Apache-2.0 + */ + +package com.google.genkit.samples; + +import com.google.genkit.Genkit; +import com.google.genkit.GenkitOptions; +import com.google.genkit.agent.AgentConfig; +import com.google.genkit.ai.InterruptConfig; +import com.google.genkit.ai.Part; +import com.google.genkit.ai.Tool; +import com.google.genkit.ai.agent.Agent; +import com.google.genkit.ai.agent.AgentChat; +import com.google.genkit.ai.agent.AgentFinishReason; +import com.google.genkit.ai.agent.AgentInterrupt; +import com.google.genkit.ai.agent.AgentResponse; +import com.google.genkit.ai.agent.FileSessionStore; +import com.google.genkit.plugins.googlegenai.GoogleGenAIPlugin; +import java.util.List; +import java.util.Map; +import java.util.Scanner; + +/** + * Human-in-the-loop agent CLI demo — a banking assistant that pauses on a sensitive action (a money + * transfer) and waits for the user to approve or reject it on the command line before continuing. + * + *

This demonstrates the agent-level interrupt/resume flow: + * + *

    + *
  • A tool created with {@code genkit.defineInterrupt(...)} pauses the turn instead of running. + *
  • The turn finishes with {@link AgentFinishReason#INTERRUPTED} and {@link + * AgentResponse#interrupts()} surfaces the pending tool request. + *
  • The caller resolves it with {@code tool.respond(interrupt.part(), output)} and resumes via + * {@link AgentChat#resume(java.util.List)}; the agent then completes the turn. + *
+ * + *

State is persisted server-side via {@link FileSessionStore} (under {@code ./.snapshots}), so + * the paused turn survives across the resume. + * + *

Run it (requires {@code GEMINI_API_KEY}): + * + * 

+ *   export GEMINI_API_KEY=your-key
+ *   mvn -q exec:java
+ *
+ */ +public class HumanInTheLoopAgentApp { + + /** Input for the money-transfer confirmation interrupt. */ + public static class TransferRequest { + private String recipient; + private double amount; + private String reason; + + public TransferRequest() {} + + public String getRecipient() { + return recipient; + } + + public void setRecipient(String recipient) { + this.recipient = recipient; + } + + public double getAmount() { + return amount; + } + + public void setAmount(double amount) { + this.amount = amount; + } + + public String getReason() { + return reason; + } + + public void setReason(String reason) { + this.reason = reason; + } + } + + /** The caller's decision returned to the interrupted tool. */ + public static class ConfirmationOutput { + private boolean confirmed; + private String reason; + + public ConfirmationOutput() {} + + public ConfirmationOutput(boolean confirmed, String reason) { + this.confirmed = confirmed; + this.reason = reason; + } + + public boolean isConfirmed() { + return confirmed; + } + + public void setConfirmed(boolean confirmed) { + this.confirmed = confirmed; + } + + public String getReason() { + return reason; + } + + public void setReason(String reason) { + this.reason = reason; + } + } + + // Scanner wraps System.in for the CLI prompt; we intentionally do not close it (closing + // System.in would break stdin for the rest of the JVM) — the process exits right after. + @SuppressWarnings("resource") + public static void main(String[] args) throws Exception { + // ── 1. Build Genkit with the beta agents API enabled ──────────────────── + Genkit genkit = + Genkit.builder() + .options(GenkitOptions.builder().experimental(true).build()) + .plugin(GoogleGenAIPlugin.create()) + .build(); + + // ── 2. Define an interrupt tool ───────────────────────────────────────── + // + // defineInterrupt creates a tool that PAUSES the turn (throwing the proper + // interrupt internally) instead of executing. The model calls it like any + // tool; the turn then finishes INTERRUPTED so a human can decide. + Tool confirmTransfer = + genkit.defineInterrupt( + InterruptConfig.builder() + .name("confirmTransfer") + .description( + "Request user confirmation before executing a money transfer. " + + "ALWAYS use this tool before transferring money.") + .inputType(TransferRequest.class) + .outputType(ConfirmationOutput.class) + .inputSchema( + Map.of( + "type", + "object", + "properties", + Map.of( + "recipient", + Map.of("type", "string", "description", "Who to transfer to"), + "amount", + Map.of("type", "number", "description", "Amount to transfer"), + "reason", + Map.of("type", "string", "description", "Reason for transfer")), + "required", + List.of("recipient", "amount"))) + .requestMetadata( + input -> + Map.of( + "type", + "transfer_confirmation", + "recipient", + input.getRecipient() != null ? input.getRecipient() : "", + "amount", + input.getAmount(), + "reason", + input.getReason() != null ? input.getReason() : "")) + .build()); + + // ── 3. Define a server-managed agent that uses the interrupt tool ─────── + Agent> bankingAgent = + genkit + .beta() + .defineAgent( + AgentConfig.>builder() + .name("bankingAgent") + .description("A banking assistant that confirms transfers with the user") + .system( + "You are a helpful banking assistant. Whenever the user asks to transfer" + + " or send money, you MUST call the confirmTransfer tool first and" + + " only proceed once it returns a confirmation.") + .tools(confirmTransfer) + .model("googleai/gemini-2.5-flash") + .store(new FileSessionStore<>("./.snapshots")) + .build()); + + // ── 4. Run the interrupt → human decision → resume flow (needs GEMINI_API_KEY) ── + String apiKey = System.getenv("GEMINI_API_KEY"); + if (apiKey == null || apiKey.isBlank()) { + System.out.println( + "GEMINI_API_KEY is not set — agent defined successfully but skipping live calls."); + System.out.println("Set GEMINI_API_KEY and re-run to see the full human-in-the-loop flow."); + return; + } + + System.out.println("=== Human-in-the-loop banking agent ==="); + AgentChat> chat = bankingAgent.chat(); + + AgentResponse> turn1 = chat.send("Transfer $150 to Alice for dinner"); + System.out.println("Finish reason: " + turn1.finishReason()); + + if (turn1.finishReason() == AgentFinishReason.INTERRUPTED && !turn1.interrupts().isEmpty()) { + AgentInterrupt interrupt = turn1.interrupts().get(0); + System.out.println("Agent paused awaiting approval for tool: " + interrupt.name()); + System.out.println(" Requested transfer: " + interrupt.input()); + + System.out.print("Approve this transfer? (yes/no): "); + Scanner scanner = new Scanner(System.in); + String answer = scanner.hasNextLine() ? scanner.nextLine().trim().toLowerCase() : "no"; + boolean approved = answer.equals("yes") || answer.equals("y"); + + // Build the response part for the interrupted tool and resume the turn. + Part respond = + confirmTransfer.respond( + interrupt.part(), + new ConfirmationOutput(approved, approved ? "User approved" : "User declined")); + AgentResponse> resumed = chat.resume(List.of(respond)); + + System.out.println("Resumed finish reason: " + resumed.finishReason()); + System.out.println("Final response: " + resumed.text()); + } else { + System.out.println("Agent did not interrupt. Response: " + turn1.text()); + } + } +} diff --git a/samples/agents-human-in-the-loop/src/main/resources/logback.xml b/samples/agents-human-in-the-loop/src/main/resources/logback.xml new file mode 100644 index 000000000..ec9b691df --- /dev/null +++ b/samples/agents-human-in-the-loop/src/main/resources/logback.xml @@ -0,0 +1,26 @@ + + + + + + %d{HH:mm:ss.SSS} %-5level %logger{24} - %msg%n + + + + + + + + + + + + + + + + + + + diff --git a/samples/agents-orchestrator/pom.xml b/samples/agents-orchestrator/pom.xml new file mode 100644 index 000000000..b03ef9d6b --- /dev/null +++ b/samples/agents-orchestrator/pom.xml @@ -0,0 +1,87 @@ + + + + 4.0.0 + + + com.google.genkit + genkit-parent + 1.0.0-SNAPSHOT + ../../pom.xml + + + com.google.genkit.samples + genkit-sample-agents-orchestrator + jar + Genkit Agents Orchestrator Sample + Sample demonstrating the Genkit Agents orchestrator pattern with sub-agent delegation + + + UTF-8 + 21 + 21 + 1.0.0-SNAPSHOT + true + true + + + + + com.google.genkit + genkit + ${genkit.version} + + + com.google.genkit + genkit-plugin-middleware + ${genkit.version} + + + com.google.genkit + genkit-plugin-openai + ${genkit.version} + + + com.google.genkit + genkit-plugin-jetty + ${genkit.version} + + + ch.qos.logback + logback-classic + 1.5.37 + + + + + + + org.codehaus.mojo + exec-maven-plugin + 3.6.3 + + com.google.genkit.samples.OrchestratorApp + + + + + diff --git a/samples/agents-orchestrator/run.sh b/samples/agents-orchestrator/run.sh new file mode 100755 index 000000000..d10fb3e2a --- /dev/null +++ b/samples/agents-orchestrator/run.sh @@ -0,0 +1,7 @@ +#!/bin/bash + +# Genkit Sample Runner +# This script runs the sample application + +echo "Building and running the sample..." +mvn clean compile exec:java diff --git a/samples/agents-orchestrator/src/main/java/com/google/genkit/samples/OrchestratorApp.java b/samples/agents-orchestrator/src/main/java/com/google/genkit/samples/OrchestratorApp.java new file mode 100644 index 000000000..48474f936 --- /dev/null +++ b/samples/agents-orchestrator/src/main/java/com/google/genkit/samples/OrchestratorApp.java @@ -0,0 +1,193 @@ +/* + * 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. + * + * SPDX-License-Identifier: Apache-2.0 + */ + +package com.google.genkit.samples; + +import com.google.genkit.Genkit; +import com.google.genkit.GenkitOptions; +import com.google.genkit.agent.AgentConfig; +import com.google.genkit.ai.Tool; +import com.google.genkit.ai.agent.Agent; +import com.google.genkit.ai.agent.AgentChat; +import com.google.genkit.ai.agent.AgentResponse; +import com.google.genkit.plugins.jetty.JettyPlugin; +import com.google.genkit.plugins.middleware.Agents; +import com.google.genkit.plugins.middleware.AgentsOptions; +import com.google.genkit.plugins.openai.OpenAIPlugin; +import java.util.ArrayList; +import java.util.List; +import java.util.Map; + +/** + * Orchestrator agent sample demonstrating sub-agent delegation via the middleware. + * + *

This mirrors the JavaScript {@code orchestrator-agent.ts} testapp. Two specialised sub-agents + * ({@code researcher} and {@code coder}) are defined first; the orchestrator is given delegation + * tools produced by {@link Agents#delegationTools(AgentsOptions)} and a system-prompt fragment + * produced by {@link Agents#systemPromptFragment(AgentsOptions)} so it can dispatch tasks to them. + * + *

Two run modes: + * + *

    + *
  • Serve (default) — starts the Jetty plugin to expose the agents over HTTP and keep + * the process alive. Run under the Genkit CLI to test the agents in the Dev UI: {@code + * genkit start -- mvn -q -pl samples/agents-orchestrator exec:java}. The agents are also + * reachable at {@code POST http://localhost:8080/} (override the port with the {@code + * PORT} env var). + *
  • Demo — runs the in-process orchestrator delegation demo once and exits. Requires + * {@code OPENAI_API_KEY}: {@code mvn -q -pl samples/agents-orchestrator exec:java + * -Dexec.args=demo}. + *
+ */ +public class OrchestratorApp { + + public static void main(String[] args) throws Exception { + // ── 1. Build Genkit with experimental (agents) enabled ────────────────── + Genkit genkit = + Genkit.builder() + .options( + GenkitOptions.builder() + .experimental(true) // required for the beta agents API + .devMode(true) + .build()) + .plugin(OpenAIPlugin.create()) + .build(); + + // ── 2. Define the researcher sub-agent ────────────────────────────────── + // + // A specialised agent that researches topics. + @SuppressWarnings("unused") + Agent> researcher = + genkit + .beta() + .defineAgent( + AgentConfig.>builder() + .name("researcher") + .description("Researches topics and returns a summary") + .system( + "You are a knowledgeable researcher. " + + "When given a topic, provide a concise, factual summary. " + + "Focus on key facts, recent developments, and practical implications.") + .model("openai/gpt-4o-mini") + .build()); + + // ── 3. Define the coder sub-agent ──────────────────────────────────────── + // + // A specialised agent that writes code based on a description. + @SuppressWarnings("unused") + Agent> coder = + genkit + .beta() + .defineAgent( + AgentConfig.>builder() + .name("coder") + .description("Writes code based on a description or research summary") + .system( + "You are an expert software engineer. " + + "When given a task description or research summary, write clean, " + + "well-commented code. Prefer Java when no language is specified. " + + "Include usage examples in comments.") + .model("openai/gpt-4o-mini") + .build()); + + // ── 4. Build delegation tools via the middleware Agents factory ───────── + // + // AgentsOptions lists the registered agent names the orchestrator can + // delegate to. Agents.delegationTools() produces one Tool per sub-agent + // (named "delegate_to_" by default). + // Agents.systemPromptFragment() builds a XML block that + // explains the tools to the model. + AgentsOptions delegationOpts = + AgentsOptions.builder().agents("researcher", "coder").maxDelegations(5).build(); + + List> delegationTools = Agents.delegationTools(delegationOpts); + String subAgentFragment = Agents.systemPromptFragment(delegationOpts); + + String baseSystem = + "You are an orchestrator agent that breaks complex tasks into sub-tasks " + + "and delegates them to specialised agents. " + + "First research the topic, then produce code based on the findings."; + + // ── 5. Define the orchestrator agent ──────────────────────────────────── + Agent> orchestrator = + genkit + .beta() + .defineAgent( + AgentConfig.>builder() + .name("orchestrator") + .description( + "Orchestrator that delegates tasks to researcher and coder sub-agents") + .system(baseSystem + "\n" + subAgentFragment) + .tools(new ArrayList<>(delegationTools)) + .model("openai/gpt-4o-mini") + .build()); + + int port = System.getenv("PORT") != null ? Integer.parseInt(System.getenv("PORT")) : 8080; + JettyPlugin jetty = JettyPlugin.create(port); + jetty.init(genkit.getRegistry()); + System.out.println("Serving agents on http://localhost:" + port); + System.out.println(" orchestrator -> POST http://localhost:" + port + "/orchestrator"); + System.out.println(" researcher -> POST http://localhost:" + port + "/researcher"); + System.out.println(" coder -> POST http://localhost:" + port + "/coder"); + System.out.println( + "Tip: run under `genkit start -- mvn -q -pl samples/agents-orchestrator exec:java` to" + + " open the Dev UI, or pass `demo` to run the one-shot delegation demo."); + try { + // + // The module must compile even without an API key. Real calls only execute + // when OPENAI_API_KEY is set so CI / offline builds succeed. + String apiKey = System.getenv("OPENAI_API_KEY"); + if (apiKey == null || apiKey.isBlank()) { + System.out.println( + "OPENAI_API_KEY is not set — agents defined successfully but skipping live calls."); + System.out.println("Set OPENAI_API_KEY and re-run with `demo` to see the orchestrator."); + return; + } + + System.out.println("=== Orchestrator Agent ==="); + AgentChat> chat = orchestrator.chat(); + + AgentResponse> response = + chat.send("Research the bubble sort algorithm and then write Java code to implement it."); + System.out.println("Orchestrator response:\n" + response.text()); + + jetty.start(); // blocks until the process is stopped (Ctrl-C) + } catch (Exception e) { + // The most common failure here is a port collision (something else — often + // another + // dev server — is already listening on the port). Fail loudly with an + // actionable hint + // instead of leaving the misleading "Serving agents on ..." message above + // standing while + // the agents are actually unreachable. + System.err.println(); + System.err.println( + "ERROR: could not start the agent HTTP server on port " + port + ": " + e.getMessage()); + System.err.println( + "Port " + + port + + " is likely already in use. Free it (e.g. `lsof -nP -iTCP:" + + port + + " -sTCP:LISTEN` then kill the process), or run on a different port, e.g.: PORT=" + + (port + 1) + + " mvn -q -pl samples/agents-orchestrator exec:java"); + throw e; + } + return; + } +} diff --git a/samples/agents-orchestrator/src/main/resources/logback.xml b/samples/agents-orchestrator/src/main/resources/logback.xml new file mode 100644 index 000000000..fe6b73651 --- /dev/null +++ b/samples/agents-orchestrator/src/main/resources/logback.xml @@ -0,0 +1,24 @@ + + + + + + %d{HH:mm:ss.SSS} %-5level %logger{24} - %msg%n + + + + + + + + + + + + + + + + + diff --git a/samples/agents-remote/pom.xml b/samples/agents-remote/pom.xml new file mode 100644 index 000000000..14401bfa0 --- /dev/null +++ b/samples/agents-remote/pom.xml @@ -0,0 +1,72 @@ + + + + 4.0.0 + + + com.google.genkit + genkit-parent + 1.0.0-SNAPSHOT + ../../pom.xml + + + com.google.genkit.samples + genkit-sample-agents-remote + jar + Genkit Agents Remote Client Sample + Sample demonstrating the Genkit RemoteAgent client for talking to a served agent over HTTP + + + UTF-8 + 21 + 21 + 1.0.0-SNAPSHOT + true + true + + + + + com.google.genkit + genkit + ${genkit.version} + + + ch.qos.logback + logback-classic + 1.5.37 + + + + + + + org.codehaus.mojo + exec-maven-plugin + 3.6.3 + + com.google.genkit.samples.RemoteAgentClientApp + + + + + diff --git a/samples/agents-remote/run.sh b/samples/agents-remote/run.sh new file mode 100755 index 000000000..d10fb3e2a --- /dev/null +++ b/samples/agents-remote/run.sh @@ -0,0 +1,7 @@ +#!/bin/bash + +# Genkit Sample Runner +# This script runs the sample application + +echo "Building and running the sample..." +mvn clean compile exec:java diff --git a/samples/agents-remote/src/main/java/com/google/genkit/samples/RemoteAgentClientApp.java b/samples/agents-remote/src/main/java/com/google/genkit/samples/RemoteAgentClientApp.java new file mode 100644 index 000000000..bac7ede19 --- /dev/null +++ b/samples/agents-remote/src/main/java/com/google/genkit/samples/RemoteAgentClientApp.java @@ -0,0 +1,128 @@ +/* + * 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. + * + * SPDX-License-Identifier: Apache-2.0 + */ + +package com.google.genkit.samples; + +import com.google.genkit.ai.agent.AgentChat; +import com.google.genkit.ai.agent.AgentResponse; +import com.google.genkit.client.RemoteAgent; +import com.google.genkit.client.RemoteAgentOptions; +import java.util.Map; + +/** + * Remote agent client sample. + * + *

This mirrors the JavaScript {@code remote-client.ts} testapp. It uses {@link RemoteAgent} to + * connect to an agent served over HTTP (e.g. by the Jetty plugin). The client holds the + * conversation history locally and sends it with every turn — the server itself can be stateless or + * server-managed depending on how the served agent was configured. + * + *

Running this sample

+ * + *
    + *
  1. Start the server: run {@code samples/agents-weather} with the Jetty plugin serving the + * {@code weatherAgent} at {@code http://localhost:8080/weatherAgent}. + *
  2. Set {@code AGENT_URL} (optional, defaults to {@code http://localhost:8080/weatherAgent}). + *
  3. Set {@code REMOTE_AGENT_RUN=true} to enable the actual network call (omit for offline / CI + * builds). + *
  4. {@code mvn exec:java -pl samples/agents-remote} + *
+ * + *

Architecture note

+ * + *

{@link RemoteAgent#chat(RemoteAgentOptions)} returns an {@link AgentChat} backed by an {@code + * HttpAgentTransport}. Each {@code send()} posts a turn to {@code url} and (in client-managed mode) + * reads back the updated session state so that subsequent turns include the full conversation + * history. The companion endpoints {@code url/getSnapshot} and {@code url/abort} are derived + * automatically from the base URL. + */ +public class RemoteAgentClientApp { + + /** Default agent URL — override via the {@code AGENT_URL} environment variable. */ + private static final String DEFAULT_AGENT_URL = "http://localhost:8080/weatherAgent"; + + public static void main(String[] args) { + // ── 1. Resolve the agent URL ───────────────────────────────────────────── + // + // Override with AGENT_URL env var to point at a different agent endpoint. + String agentUrl = System.getenv("AGENT_URL"); + if (agentUrl == null || agentUrl.isBlank()) { + agentUrl = DEFAULT_AGENT_URL; + } + System.out.println("Remote agent URL: " + agentUrl); + + // ── 2. Build RemoteAgentOptions ────────────────────────────────────────── + // + // RemoteAgent.chat() wraps an HttpAgentTransport configured with the options. + // getSnapshotUrl and abortUrl default to url+"/getSnapshot" and url+"/abort". + // serverManaged defaults to true — set false for client-managed mode. + RemoteAgentOptions opts = + RemoteAgentOptions.builder() + .url(agentUrl) + // serverManaged(false) // uncomment for client-managed remote agent + .build(); + + // ── 3. Create the remote chat client ───────────────────────────────────── + // + // No Genkit instance is needed on the client side — RemoteAgent is a + // pure HTTP client that speaks the Genkit agent wire format. + AgentChat> chat = RemoteAgent.chat(opts); + + // ── 4. Guard live network calls behind an env flag ─────────────────────── + // + // The module must compile and start even without a running server. + // Set REMOTE_AGENT_RUN=true to actually send requests. + String runFlag = System.getenv("REMOTE_AGENT_RUN"); + if (!"true".equalsIgnoreCase(runFlag)) { + System.out.println( + "REMOTE_AGENT_RUN is not set to 'true' — RemoteAgent client created successfully " + + "but skipping live network calls."); + System.out.println( + "Start the weather agent server and set REMOTE_AGENT_RUN=true to run the demo."); + return; + } + + // ── 5. Multi-turn remote chat ──────────────────────────────────────────── + System.out.println("=== Remote agent client ==="); + + try { + AgentResponse> turn1 = chat.send("What is the weather in Tokyo?"); + System.out.println("Turn 1: " + turn1.text()); + + AgentResponse> turn2 = chat.send("And in London?"); + System.out.println("Turn 2: " + turn2.text()); + } catch (RuntimeException e) { + String msg = e.getMessage() == null ? "" : e.getMessage(); + System.err.println("Remote call to " + agentUrl + " failed: " + msg); + if (msg.contains("404") || msg.contains("Cannot POST") || msg.contains("Express")) { + // A 404 / "Cannot POST" (Express) reply means the URL is NOT the Java agent server — + // some other process answered. Almost always a port mismatch / the Java server didn't bind. + System.err.println( + "That looks like a different server answered (e.g. an Express dev server), not the Java" + + " agent endpoint. Make sure the weather sample is running AND that its port" + + " matches this URL."); + System.err.println( + " - Server: look for `Jetty server started on 0.0.0.0:` in its logs (if the port" + + " was busy it fails to bind)."); + System.err.println( + " - Then point this client at the same port, e.g. AGENT_URL=http://localhost:/weatherAgent"); + } + throw e; + } + } +} diff --git a/samples/agents-remote/src/main/resources/logback.xml b/samples/agents-remote/src/main/resources/logback.xml new file mode 100644 index 000000000..fe6b73651 --- /dev/null +++ b/samples/agents-remote/src/main/resources/logback.xml @@ -0,0 +1,24 @@ + + + + + + %d{HH:mm:ss.SSS} %-5level %logger{24} - %msg%n + + + + + + + + + + + + + + + + + diff --git a/samples/agents-stateless/pom.xml b/samples/agents-stateless/pom.xml new file mode 100644 index 000000000..22231e56a --- /dev/null +++ b/samples/agents-stateless/pom.xml @@ -0,0 +1,77 @@ + + + + 4.0.0 + + + com.google.genkit + genkit-parent + 1.0.0-SNAPSHOT + ../../pom.xml + + + com.google.genkit.samples + genkit-sample-agents-stateless + jar + Genkit Agents Stateless Sample + Sample demonstrating the Genkit Agents API with client-managed (stateless) session state + + + UTF-8 + 21 + 21 + 1.0.0-SNAPSHOT + true + true + + + + + com.google.genkit + genkit + ${genkit.version} + + + com.google.genkit + genkit-plugin-openai + ${genkit.version} + + + ch.qos.logback + logback-classic + 1.5.37 + + + + + + + org.codehaus.mojo + exec-maven-plugin + 3.6.3 + + com.google.genkit.samples.StatelessAgentApp + + + + + diff --git a/samples/agents-stateless/run.sh b/samples/agents-stateless/run.sh new file mode 100755 index 000000000..d10fb3e2a --- /dev/null +++ b/samples/agents-stateless/run.sh @@ -0,0 +1,7 @@ +#!/bin/bash + +# Genkit Sample Runner +# This script runs the sample application + +echo "Building and running the sample..." +mvn clean compile exec:java diff --git a/samples/agents-stateless/src/main/java/com/google/genkit/samples/StatelessAgentApp.java b/samples/agents-stateless/src/main/java/com/google/genkit/samples/StatelessAgentApp.java new file mode 100644 index 000000000..dd30365b4 --- /dev/null +++ b/samples/agents-stateless/src/main/java/com/google/genkit/samples/StatelessAgentApp.java @@ -0,0 +1,165 @@ +/* + * 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. + * + * SPDX-License-Identifier: Apache-2.0 + */ + +package com.google.genkit.samples; + +import com.google.genkit.Genkit; +import com.google.genkit.GenkitOptions; +import com.google.genkit.agent.AgentConfig; +import com.google.genkit.ai.Tool; +import com.google.genkit.ai.agent.Agent; +import com.google.genkit.ai.agent.AgentChat; +import com.google.genkit.ai.agent.AgentResponse; +import com.google.genkit.plugins.openai.OpenAIPlugin; +import java.util.Map; + +/** + * Client-managed (stateless) agent sample. + * + *

This mirrors the JavaScript {@code weather-agent-stateless.ts} testapp. The agent is defined + * WITHOUT a {@code .store(...)} call, so no session state is persisted server-side. Instead, {@link + * AgentChat} automatically round-trips the full {@code SessionState} (messages + custom state) on + * every {@code send()} call. The server processes each turn statelessly while the client holds the + * conversation history. + * + *

To run: + * + *

    + *
  1. Set {@code OPENAI_API_KEY} in the environment. + *
  2. {@code mvn exec:java -pl samples/agents-stateless} + *
+ */ +public class StatelessAgentApp { + + /** Simple input type for the weather tool. */ + public static class WeatherInput { + private String location; + + public WeatherInput() {} + + public WeatherInput(String location) { + this.location = location; + } + + public String getLocation() { + return location; + } + + public void setLocation(String location) { + this.location = location; + } + } + + /** Simple output type for the weather tool. */ + public static class WeatherOutput { + private String report; + + public WeatherOutput() {} + + public WeatherOutput(String report) { + this.report = report; + } + + public String getReport() { + return report; + } + + public void setReport(String report) { + this.report = report; + } + } + + public static void main(String[] args) { + // ── 1. Build Genkit with experimental (agents) enabled ────────────────── + Genkit genkit = + Genkit.builder() + .options( + GenkitOptions.builder() + .experimental(true) // required for the beta agents API + .devMode(true) + .build()) + .plugin(OpenAIPlugin.create()) + .build(); + + // ── 2. Define a mock weather tool ─────────────────────────────────────── + Tool getWeather = + genkit.defineTool( + "getWeather", + "Returns current weather conditions for a given location", + (ctx, input) -> { + String location = input != null ? input.getLocation() : "unknown"; + return new WeatherOutput("Sunny and 22°C in " + location); + }, + WeatherInput.class, + WeatherOutput.class); + + // ── 3. Define a client-managed (stateless) agent ───────────────────────── + // + // Key: no .store(...) call. Without a SessionStore, the agent operates in + // client-managed mode. AgentChat holds the full SessionState locally and + // sends it with every turn so the model always sees the full conversation + // history — without any server-side persistence. This is useful for + // serverless environments or when you want the client to own session state. + Agent> statelessWeather = + genkit + .beta() + .defineAgent( + AgentConfig.>builder() + .name("statelessWeather") + .description("A helpful weather assistant (client-managed state)") + .system( + "You are a helpful weather assistant. " + + "Use the getWeather tool to answer questions about the weather. " + + "Remember previous questions in the conversation.") + .tools(getWeather) + .model("openai/gpt-4o-mini") + // No .store() → client-managed: AgentChat round-trips full state each turn. + .build()); + + // ── 4. Guard live model calls behind an API-key check ─────────────────── + // + // The module must compile even without an API key. Real calls only execute + // when OPENAI_API_KEY is set so CI / offline builds succeed. + String apiKey = System.getenv("OPENAI_API_KEY"); + if (apiKey == null || apiKey.isBlank()) { + System.out.println( + "OPENAI_API_KEY is not set — agent defined successfully but skipping live calls."); + System.out.println("Set OPENAI_API_KEY and re-run to see the stateless agent demo."); + return; + } + + // ── 5. Multi-turn chat with client-managed state ───────────────────────── + // + // The AgentChat object holds the full conversation history locally. + // Each send() serialises the current state into the AgentInit and + // the server processes the turn without reading or writing any store. + System.out.println("=== Client-managed (stateless) weather agent ==="); + AgentChat> chat = statelessWeather.chat(); + + AgentResponse> turn1 = chat.send("What is the weather in Tokyo?"); + System.out.println("Turn 1: " + turn1.text()); + + // The full conversation history (including turn 1 and the tool call) is + // automatically carried in the request — no server-side store is accessed. + AgentResponse> turn2 = chat.send("How does that compare to Paris?"); + System.out.println("Turn 2: " + turn2.text()); + + AgentResponse> turn3 = chat.send("Which city had warmer weather?"); + System.out.println("Turn 3: " + turn3.text()); + } +} diff --git a/samples/agents-stateless/src/main/java/resources/logback.xml b/samples/agents-stateless/src/main/java/resources/logback.xml new file mode 100644 index 000000000..fe6b73651 --- /dev/null +++ b/samples/agents-stateless/src/main/java/resources/logback.xml @@ -0,0 +1,24 @@ + + + + + + %d{HH:mm:ss.SSS} %-5level %logger{24} - %msg%n + + + + + + + + + + + + + + + + + diff --git a/samples/multi-agent/pom.xml b/samples/agents-weather/pom.xml similarity index 90% rename from samples/multi-agent/pom.xml rename to samples/agents-weather/pom.xml index 1ff7bc8ed..aa4d3e710 100644 --- a/samples/multi-agent/pom.xml +++ b/samples/agents-weather/pom.xml @@ -30,10 +30,10 @@ com.google.genkit.samples - genkit-sample-multi-agent + genkit-sample-agents-weather jar - Genkit Multi-Agent Sample - Sample application demonstrating multi-agent patterns with agent delegation + Genkit Agents Weather Sample + Sample application demonstrating the Genkit Agents API with a weather assistant UTF-8 @@ -74,7 +74,7 @@ exec-maven-plugin 3.6.3 - com.google.genkit.samples.MultiAgentApp + com.google.genkit.samples.WeatherAgentApp diff --git a/samples/agents-weather/run.sh b/samples/agents-weather/run.sh new file mode 100755 index 000000000..d10fb3e2a --- /dev/null +++ b/samples/agents-weather/run.sh @@ -0,0 +1,7 @@ +#!/bin/bash + +# Genkit Sample Runner +# This script runs the sample application + +echo "Building and running the sample..." +mvn clean compile exec:java diff --git a/samples/agents-weather/src/main/java/com/google/genkit/samples/WeatherAgentApp.java b/samples/agents-weather/src/main/java/com/google/genkit/samples/WeatherAgentApp.java new file mode 100644 index 000000000..4debca658 --- /dev/null +++ b/samples/agents-weather/src/main/java/com/google/genkit/samples/WeatherAgentApp.java @@ -0,0 +1,274 @@ +/* + * 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. + * + * SPDX-License-Identifier: Apache-2.0 + */ + +package com.google.genkit.samples; + +import com.google.genkit.Genkit; +import com.google.genkit.GenkitOptions; +import com.google.genkit.agent.AgentConfig; +import com.google.genkit.ai.Tool; +import com.google.genkit.ai.agent.Agent; +import com.google.genkit.ai.agent.AgentChat; +import com.google.genkit.ai.agent.AgentResponse; +import com.google.genkit.ai.agent.FileSessionStore; +import com.google.genkit.plugins.jetty.JettyPlugin; +import com.google.genkit.plugins.openai.OpenAIPlugin; +import java.util.Map; + +/** + * Canonical Genkit Agents sample — a weather assistant demonstrating: + * + *
    + *
  • Server-managed session state via {@link FileSessionStore} ({@code weatherAgent}) + *
  • Client-managed (stateless) session state ({@code weatherAgentStateless}) + *
  • Multi-turn chat with automatic history carry-forward + *
  • Streaming with {@code sendStream} + *
  • Tool definition via {@code genkit.defineTool} + *
+ * + *

This mirrors the JavaScript {@code weather-agent.ts} / {@code weather-agent-stateless.ts} + * testapps. The agents API is in beta — enable it with {@code + * GenkitOptions.builder().experimental(true)}. + * + *

Two run modes: + * + *

    + *
  • Serve (default) — starts the Jetty plugin to expose the agents over HTTP and keep + * the process alive. Run under the Genkit CLI to test the agents in the Dev UI: {@code + * genkit start -- mvn -q -pl samples/agents-weather exec:java}. The agents are also reachable + * at {@code POST http://localhost:8080/weatherAgent} (override the port with the {@code PORT} + * env var). No API key is needed to start the server, but live model calls still require + * {@code OPENAI_API_KEY}. + *
  • Demo — runs the in-process multi-turn/streaming chat demo and exits. Requires {@code + * OPENAI_API_KEY}: {@code mvn -q -pl samples/agents-weather exec:java -Dexec.args=demo}. + *
+ */ +public class WeatherAgentApp { + + /** Simple input type for the weather tool. */ + public static class WeatherInput { + private String location; + + public WeatherInput() {} + + public WeatherInput(String location) { + this.location = location; + } + + public String getLocation() { + return location; + } + + public void setLocation(String location) { + this.location = location; + } + } + + /** Simple output type for the weather tool. */ + public static class WeatherOutput { + private String report; + + public WeatherOutput() {} + + public WeatherOutput(String report) { + this.report = report; + } + + public String getReport() { + return report; + } + + public void setReport(String report) { + this.report = report; + } + } + + public static void main(String[] args) throws Exception { + // ── 1. Build Genkit with experimental (agents) enabled ────────────────── + Genkit genkit = + Genkit.builder() + .options( + GenkitOptions.builder() + .experimental(true) // required for the beta agents API + .devMode(true) + .build()) + .plugin(OpenAIPlugin.create()) + .build(); + + // ── 2. Define a mock weather tool ─────────────────────────────────────── + // + // In production you would call a real weather service here. + // The tool signature uses the typed defineTool variant so JSON schema is + // auto-generated from WeatherInput / WeatherOutput. + Tool getWeather = + genkit.defineTool( + "getWeather", + "Returns current weather conditions for a given location", + (ctx, input) -> { + // Simulated weather — swap for a real API call as needed. + String location = input != null ? input.getLocation() : "unknown"; + return new WeatherOutput("Sunny and 22°C in " + location); + }, + WeatherInput.class, + WeatherOutput.class); + + // ── 3. Server-managed agent (state stored in .snapshots/) ─────────────── + // + // The FileSessionStore persists conversation snapshots to disk so the agent + // can resume across process restarts. Pass a store to .store(...) to enable + // server-managed mode; omit it for client-managed (stateless) mode. + Agent> weatherAgent = + genkit + .beta() + .defineAgent( + AgentConfig.>builder() + .name("weatherAgent") + .description("A helpful weather assistant") + .system( + "You are a helpful weather assistant. " + + "Use the getWeather tool to answer questions about the weather.") + .tools(getWeather) + .model("openai/gpt-4o-mini") + .store(new FileSessionStore<>("./.snapshots")) + .build()); + + // ── 4. Client-managed (stateless) agent ───────────────────────────────── + // + // No .store(...) → the agent does NOT persist state server-side. + // Instead the full SessionState round-trips through AgentChat automatically: + // each send() carries the current messages/state in the AgentInit so the + // model always sees the full conversation history. + Agent> weatherAgentStateless = + genkit + .beta() + .defineAgent( + AgentConfig.>builder() + .name("weatherAgentStateless") + .description("A helpful weather assistant (client-managed state)") + .system( + "You are a helpful weather assistant. " + + "Use the getWeather tool to answer questions about the weather.") + .tools(getWeather) + .model("openai/gpt-4o-mini") + // No .store() → client-managed: AgentChat round-trips full state each turn. + .build()); + + // ── 5. Default mode: serve the agents over HTTP + keep the process alive ─ + // + // Run with NO args (e.g. under `genkit start`) to expose the agents: + // • The Genkit Dev UI discovers and runs them via the reflection server, + // which starts automatically because devMode(true) is set above. Under + // `genkit start` the runtime connects to the CLI and the agents show up + // in the browser's "Agents"/"Flows" surface, ready to chat. + // • Jetty also serves each agent over plain HTTP at POST / + // (plus /getSnapshot and /abort companions for server-managed agents), + // so you can curl them or drive them with the remoteAgent client. + // + // jetty.start() blocks until the process is stopped — that is what keeps the + // runtime alive for the Dev UI. Pass the "demo" argument to instead run the + // in-process chat demo below (needs OPENAI_API_KEY). + boolean runDemo = args.length > 0 && "demo".equalsIgnoreCase(args[0]); + if (!runDemo) { + int port = System.getenv("PORT") != null ? Integer.parseInt(System.getenv("PORT")) : 8080; + JettyPlugin jetty = JettyPlugin.create(port); + jetty.init(genkit.getRegistry()); + System.out.println("Serving agents on http://localhost:" + port); + System.out.println( + " weatherAgent -> POST http://localhost:" + port + "/weatherAgent"); + System.out.println( + " weatherAgentStateless -> POST http://localhost:" + port + "/weatherAgentStateless"); + System.out.println( + "Tip: run under `genkit start -- mvn -q -pl samples/agents-weather exec:java` to open" + + " the Dev UI, or pass `demo` to run the in-process chat demo."); + try { + jetty.start(); // blocks until the process is stopped (Ctrl-C) + } catch (Exception e) { + // The most common failure here is a port collision (something else — often another + // dev server — is already listening on the port). Fail loudly with an actionable hint + // instead of leaving the misleading "Serving agents on ..." message above standing while + // the agents are actually unreachable. + System.err.println(); + System.err.println( + "ERROR: could not start the agent HTTP server on port " + port + ": " + e.getMessage()); + System.err.println( + "Port " + + port + + " is likely already in use. Free it (e.g. `lsof -nP -iTCP:" + + port + + " -sTCP:LISTEN` then kill the process), or run on a different port and point the"); + System.err.println( + "remote client at the same one, e.g.: PORT=" + + (port + 1) + + " mvn -q -pl samples/agents-weather exec:java"); + System.err.println( + " then: AGENT_URL=http://localhost:" + + (port + 1) + + "/weatherAgent REMOTE_AGENT_RUN=true mvn -q -pl samples/agents-remote exec:java"); + throw e; + } + return; + } + + // ── 6. Demo mode: in-process chat (requires OPENAI_API_KEY) ────────────── + // + // The module must compile even without an API key. Real calls only execute + // when OPENAI_API_KEY is set so CI / offline builds succeed. + String apiKey = System.getenv("OPENAI_API_KEY"); + if (apiKey == null || apiKey.isBlank()) { + System.out.println( + "OPENAI_API_KEY is not set — agents defined successfully but skipping live calls."); + System.out.println("Set OPENAI_API_KEY and re-run to see the full demo."); + return; + } + + // ── 6. Multi-turn chat with server-managed state ───────────────────────── + System.out.println("=== Server-managed agent ==="); + AgentChat> chat = weatherAgent.chat(); + + AgentResponse> res1 = chat.send("What is the weather in London?"); + System.out.println("Turn 1: " + res1.text()); + + AgentResponse> res2 = chat.send("Now say that in French"); + System.out.println("Turn 2: " + res2.text()); + + // ── 7. Streaming turn ──────────────────────────────────────────────────── + System.out.println("\n=== Streaming turn ==="); + System.out.print("Streaming: "); + AgentResponse> res3 = + chat.sendStream( + "Summarise in one sentence", + chunk -> { + // chunk.text() contains the incremental model token + if (chunk.modelChunk() != null) { + System.out.print(chunk.modelChunk().getText()); + } + }); + System.out.println("\nFull response: " + res3.text()); + + // ── 8. Client-managed (stateless) demo ────────────────────────────────── + System.out.println("\n=== Client-managed (stateless) agent ==="); + AgentChat> statelessChat = weatherAgentStateless.chat(); + + AgentResponse> sl1 = statelessChat.send("What is the weather in Tokyo?"); + System.out.println("Turn 1: " + sl1.text()); + + // State is carried automatically by AgentChat — no server-side store needed. + AgentResponse> sl2 = statelessChat.send("Compare that to Paris"); + System.out.println("Turn 2: " + sl2.text()); + } +} diff --git a/samples/agents-weather/src/main/resources/logback.xml b/samples/agents-weather/src/main/resources/logback.xml new file mode 100644 index 000000000..ec9b691df --- /dev/null +++ b/samples/agents-weather/src/main/resources/logback.xml @@ -0,0 +1,26 @@ + + + + + + %d{HH:mm:ss.SSS} %-5level %logger{24} - %msg%n + + + + + + + + + + + + + + + + + + + diff --git a/samples/aws-bedrock/README.md b/samples/aws-bedrock/README.md index 0c395b04b..372eb561e 100644 --- a/samples/aws-bedrock/README.md +++ b/samples/aws-bedrock/README.md @@ -126,7 +126,7 @@ curl -X POST http://localhost:8080/streamingDemo \ This sample uses: - `amazon.nova-pro-v1:0` - Amazon Nova Pro (multimodal) -- `anthropic.claude-3-5-sonnet-20241022-v2:0` - Claude 3.5 Sonnet +- `us.anthropic.claude-sonnet-5` - Claude Sonnet 5 (US inference profile) - `meta.llama3-3-70b-instruct-v1:0` - Llama 3.3 70B See the plugin README for the full list of supported models. diff --git a/samples/aws-bedrock/src/main/java/com/google/genkit/samples/AwsBedrockSample.java b/samples/aws-bedrock/src/main/java/com/google/genkit/samples/AwsBedrockSample.java index 2741e8957..6cbafb562 100644 --- a/samples/aws-bedrock/src/main/java/com/google/genkit/samples/AwsBedrockSample.java +++ b/samples/aws-bedrock/src/main/java/com/google/genkit/samples/AwsBedrockSample.java @@ -55,8 +55,7 @@ public static void main(String[] args) throws Exception { Genkit.builder() .options(GenkitOptions.builder().devMode(true).reflectionPort(3100).build()) .plugin( - AwsBedrockPlugin.create("us-east-1") - .customModel("us.anthropic.claude-sonnet-4-20250514-v1:0")) + AwsBedrockPlugin.create("us-east-1").customModel("us.anthropic.claude-sonnet-5")) .plugin(jetty) .build(); @@ -222,14 +221,14 @@ public static void main(String[] args) throws Exception { String.class, String.class, (ctx, prompt) -> { - // Using US inference profile for Claude 4 Sonnet + // Using US inference profile for Claude Sonnet 5 // Format: {region-prefix}.{provider}.{model-name} ModelResponse response = genkit.generate( GenerateOptions.builder() - .model("aws-bedrock/us.anthropic.claude-sonnet-4-20250514-v1:0") + .model("aws-bedrock/us.anthropic.claude-sonnet-5") .system( - "You are Claude 4 Sonnet running via inference profile for cross-region routing.") + "You are Claude Sonnet 5 running via inference profile for cross-region routing.") .prompt(prompt) .config( GenerationConfig.builder() @@ -254,14 +253,14 @@ public static void main(String[] args) throws Exception { System.out.println( " - compareModels: Compare responses from Nova Pro, Nova Lite, and Llama 3.3"); System.out.println(" - streamingDemo: Demonstrate streaming responses"); - System.out.println(" - inferenceProfileDemo: Use Claude 4 Sonnet via inference profile"); + System.out.println(" - inferenceProfileDemo: Use Claude Sonnet 5 via inference profile"); System.out.println(""); System.out.println("Models used:"); System.out.println(" - amazon.nova-pro-v1:0 (ON_DEMAND)"); System.out.println(" - amazon.nova-lite-v1:0 (ON_DEMAND)"); System.out.println( " - meta.llama3-3-70b-instruct-v1:0 (requires inference profile in some regions)"); - System.out.println(" - us.anthropic.claude-sonnet-4-20250514-v1:0 (INFERENCE_PROFILE)"); + System.out.println(" - us.anthropic.claude-sonnet-5 (INFERENCE_PROFILE)"); System.out.println(""); System.out.println( "Note: Inference profiles enable cross-region routing for better availability."); diff --git a/samples/chat-session/README.md b/samples/chat-session/README.md deleted file mode 100644 index 6df24e1e5..000000000 --- a/samples/chat-session/README.md +++ /dev/null @@ -1,91 +0,0 @@ -# Chat Session Sample - -This sample demonstrates session-based multi-turn chat with persistence in Genkit Java. - -## Features - -- **Multi-turn conversations** - Automatic conversation history management -- **Session state** - Track user preferences and conversation context -- **Session persistence** - Save and load sessions across interactions -- **Tool integration** - Using tools (note-taking) within chat sessions -- **Multiple personas** - Choose between assistant, tutor, and creative modes - -## Prerequisites - -1. Java 21 or later -2. Maven -3. OpenAI API key - -## Prerequisites - -- Java 21+ -- Maven 3.6+ -- OpenAI API key - -## Running the Sample - -### Option 1: Direct Run (Interactive Mode) - -```bash -# Set your OpenAI API key -export OPENAI_API_KEY=your-api-key-here - -# Navigate to the sample directory -cd java/samples/chat-session - -# Run the sample -./run.sh -# Or: mvn compile exec:java -``` - -### Option 2: Demo Mode - -Run the automated demo to see all features: - -```bash -cd java/samples/chat-session -mvn exec:java -Dexec.args="--demo" -``` - -### Option 3: With Genkit Dev UI - -```bash -# Set your OpenAI API key -export OPENAI_API_KEY=your-api-key-here - -# Navigate to the sample directory -cd java/samples/chat-session - -# Run with Genkit CLI -genkit start -- ./run.sh -``` - -The Dev UI will be available at http://localhost:4000 - -## Commands - -During interactive chat, you can use these commands: - -| Command | Description | -|---------|-------------| -| `/history` | Show conversation history | -| `/notes` | Show saved notes | -| `/state` | Show session state | -| `/topic X` | Set conversation topic to X | -| `/quit` | Exit the chat | - -## Example Session - -``` -What's your name? Alice - -Choose a chat persona: - 1. Assistant (general help) - 2. Tutor (learning & education) - 3. Creative (storytelling & ideas) -Enter choice (1-3): 2 - -✓ Session created: a1b2c3d4-e5f6-... -✓ Persona: tutor - -You: What is photosynthesis? \ No newline at end of file diff --git a/samples/chat-session/run.sh b/samples/chat-session/run.sh deleted file mode 100755 index 7a055a49c..000000000 --- a/samples/chat-session/run.sh +++ /dev/null @@ -1,4 +0,0 @@ -#!/bin/bash -# Run script for Genkit DotPrompt Sample -cd "$(dirname "$0")" -mvn exec:java diff --git a/samples/chat-session/src/main/java/com/google/genkit/samples/ChatSessionApp.java b/samples/chat-session/src/main/java/com/google/genkit/samples/ChatSessionApp.java deleted file mode 100644 index 357d82ccc..000000000 --- a/samples/chat-session/src/main/java/com/google/genkit/samples/ChatSessionApp.java +++ /dev/null @@ -1,451 +0,0 @@ -/* - * 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. - * - * SPDX-License-Identifier: Apache-2.0 - */ - -package com.google.genkit.samples; - -import com.google.genkit.Genkit; -import com.google.genkit.GenkitOptions; -import com.google.genkit.ai.Message; -import com.google.genkit.ai.ModelResponse; -import com.google.genkit.ai.Tool; -import com.google.genkit.ai.session.Chat; -import com.google.genkit.ai.session.ChatOptions; -import com.google.genkit.ai.session.InMemorySessionStore; -import com.google.genkit.ai.session.Session; -import com.google.genkit.ai.session.SessionOptions; -import com.google.genkit.plugins.openai.OpenAIPlugin; -import java.util.HashMap; -import java.util.List; -import java.util.Map; -import java.util.Scanner; - -/** - * Interactive Chat Application with Session Persistence. - * - *

This sample demonstrates: - * - *

    - *
  • Creating and managing chat sessions - *
  • Multi-turn conversations with automatic history - *
  • Session state management - *
  • Using tools in chat sessions - *
  • Persisting and loading sessions - *
- * - *

To run: - * - *

    - *
  1. Set the OPENAI_API_KEY environment variable - *
  2. Run: mvn exec:java -pl samples/chat-session - *
- */ -public class ChatSessionApp { - - /** Session state to track conversation context and user preferences. */ - public static class ConversationState { - private String userName; - private String topic; - private int messageCount; - - public ConversationState() { - this.messageCount = 0; - } - - public ConversationState(String userName) { - this.userName = userName; - this.messageCount = 0; - } - - public String getUserName() { - return userName; - } - - public void setUserName(String userName) { - this.userName = userName; - } - - public String getTopic() { - return topic; - } - - public void setTopic(String topic) { - this.topic = topic; - } - - public int getMessageCount() { - return messageCount; - } - - public void incrementMessageCount() { - this.messageCount++; - } - - @Override - public String toString() { - return String.format( - "User: %s, Topic: %s, Messages: %d", - userName != null ? userName : "Anonymous", - topic != null ? topic : "General", - messageCount); - } - } - - private final Genkit genkit; - private final InMemorySessionStore sessionStore; - private final Tool noteTool; - private final Map notes; - - public ChatSessionApp() { - // Initialize notes storage - this.notes = new HashMap<>(); - - // Create Genkit with OpenAI plugin - this.genkit = - Genkit.builder() - .options(GenkitOptions.builder().devMode(true).reflectionPort(3100).build()) - .plugin(OpenAIPlugin.create()) - .build(); - - // Create a shared session store - this.sessionStore = new InMemorySessionStore<>(); - - // Define a note-taking tool - this.noteTool = createNoteTool(); - } - - @SuppressWarnings("unchecked") - private Tool createNoteTool() { - return genkit.defineTool( - "saveNote", - "Saves a note for the user. Use this when the user wants to remember something.", - Map.of( - "type", - "object", - "properties", - Map.of( - "title", - Map.of("type", "string", "description", "Title of the note"), - "content", - Map.of("type", "string", "description", "Content of the note")), - "required", - new String[] {"title", "content"}), - (Class>) (Class) Map.class, - (ctx, input) -> { - String title = (String) input.get("title"); - String content = (String) input.get("content"); - notes.put(title, content); - Map result = new HashMap<>(); - result.put("status", "saved"); - result.put("message", "Note '" + title + "' has been saved."); - return result; - }); - } - - /** Creates a new chat session with the given user name. */ - public Session createSession(String userName) { - return genkit.createSession( - SessionOptions.builder() - .store(sessionStore) - .initialState(new ConversationState(userName)) - .build()); - } - - /** Loads an existing session by ID. */ - public Session loadSession(String sessionId) { - try { - return genkit - .loadSession( - sessionId, SessionOptions.builder().store(sessionStore).build()) - .get(); - } catch (Exception e) { - System.err.println("Failed to load session: " + e.getMessage()); - return null; - } - } - - /** Creates a chat instance for a session. */ - @SuppressWarnings("unchecked") - public Chat createChat(Session session, String persona) { - String systemPrompt = buildSystemPrompt(session, persona); - - return session.chat( - ChatOptions.builder() - .model("openai/gpt-4o-mini") - .system(systemPrompt) - .tools(List.of((Tool) noteTool)) - .build()); - } - - private String buildSystemPrompt(Session session, String persona) { - ConversationState state = session.getState(); - StringBuilder prompt = new StringBuilder(); - - // Base persona - switch (persona.toLowerCase()) { - case "assistant": - prompt.append("You are a helpful, friendly assistant. "); - break; - case "tutor": - prompt.append( - "You are a patient and knowledgeable tutor. Explain concepts clearly and encourage" - + " learning. "); - break; - case "creative": - prompt.append( - "You are a creative writing partner. Be imaginative and help with storytelling. "); - break; - default: - prompt.append("You are a helpful assistant. "); - } - - // Add user context if available - if (state.getUserName() != null) { - prompt.append("The user's name is ").append(state.getUserName()).append(". "); - } - - // Add topic context if set - if (state.getTopic() != null) { - prompt.append("The current topic of discussion is: ").append(state.getTopic()).append(". "); - } - - prompt.append( - "You can save notes for the user using the saveNote tool when they want to remember" - + " something important."); - - return prompt.toString(); - } - - /** Sends a message and updates session state. */ - public String chat(Chat chat, String userMessage) { - try { - // Update message count in state - Session session = chat.getSession(); - ConversationState state = session.getState(); - state.incrementMessageCount(); - session.updateState(state).join(); - - // Send message and get response - ModelResponse response = chat.send(userMessage); - return response.getText(); - } catch (Exception e) { - return "Error: " + e.getMessage(); - } - } - - /** Displays conversation history. */ - public void showHistory(Chat chat) { - System.out.println("\n--- Conversation History ---"); - List history = chat.getHistory(); - for (Message msg : history) { - String role = msg.getRole().toString(); - String text = msg.getText(); - if (text.length() > 100) { - text = text.substring(0, 100) + "..."; - } - System.out.printf("[%s]: %s%n", role, text); - } - System.out.println("--- End History ---\n"); - } - - /** Displays saved notes. */ - public void showNotes() { - System.out.println("\n--- Saved Notes ---"); - if (notes.isEmpty()) { - System.out.println("No notes saved yet."); - } else { - notes.forEach((title, content) -> System.out.printf("• %s: %s%n", title, content)); - } - System.out.println("--- End Notes ---\n"); - } - - /** Interactive chat loop. */ - public void runInteractive() { - Scanner scanner = new Scanner(System.in); - - System.out.println("╔════════════════════════════════════════════════════════════╗"); - System.out.println("║ Genkit Chat Session Demo - Interactive Chat App ║"); - System.out.println("╚════════════════════════════════════════════════════════════╝"); - System.out.println(); - - // Get user name - System.out.print("What's your name? "); - String userName = scanner.nextLine().trim(); - if (userName.isEmpty()) { - userName = "User"; - } - - // Choose persona - System.out.println("\nChoose a chat persona:"); - System.out.println(" 1. Assistant (general help)"); - System.out.println(" 2. Tutor (learning & education)"); - System.out.println(" 3. Creative (storytelling & ideas)"); - System.out.print("Enter choice (1-3): "); - String choice = scanner.nextLine().trim(); - String persona = - switch (choice) { - case "2" -> "tutor"; - case "3" -> "creative"; - default -> "assistant"; - }; - - // Create session and chat - Session session = createSession(userName); - Chat chat = createChat(session, persona); - - System.out.println("\n✓ Session created: " + session.getId()); - System.out.println("✓ Persona: " + persona); - System.out.println("\nCommands:"); - System.out.println(" /history - Show conversation history"); - System.out.println(" /notes - Show saved notes"); - System.out.println(" /state - Show session state"); - System.out.println(" /topic X - Set conversation topic to X"); - System.out.println(" /quit - Exit the chat"); - System.out.println("\nStart chatting!\n"); - - // Chat loop - while (true) { - System.out.print("You: "); - String input = scanner.nextLine().trim(); - - if (input.isEmpty()) { - continue; - } - - // Handle commands - if (input.startsWith("/")) { - if (input.equals("/quit") || input.equals("/exit")) { - System.out.println("\nGoodbye, " + userName + "! Session saved."); - break; - } else if (input.equals("/history")) { - showHistory(chat); - continue; - } else if (input.equals("/notes")) { - showNotes(); - continue; - } else if (input.equals("/state")) { - System.out.println("\nSession State: " + session.getState()); - continue; - } else if (input.startsWith("/topic ")) { - String topic = input.substring(7).trim(); - ConversationState state = session.getState(); - state.setTopic(topic); - session.updateState(state).join(); - System.out.println("✓ Topic set to: " + topic); - // Recreate chat with updated system prompt - chat = createChat(session, persona); - continue; - } else { - System.out.println("Unknown command: " + input); - continue; - } - } - - // Send message - String response = chat(chat, input); - System.out.println("\nAssistant: " + response + "\n"); - } - - scanner.close(); - } - - /** Demo mode showing various session features. */ - public void runDemo() { - System.out.println("╔════════════════════════════════════════════════════════════╗"); - System.out.println("║ Genkit Chat Session Demo - Automated Demo ║"); - System.out.println("╚════════════════════════════════════════════════════════════╝"); - System.out.println(); - - // Demo 1: Basic multi-turn conversation - System.out.println("=== Demo 1: Multi-turn Conversation ===\n"); - Session session1 = createSession("Alice"); - Chat chat1 = createChat(session1, "assistant"); - - String[] questions = { - "What are the three laws of thermodynamics?", - "Can you explain the second one in simpler terms?", - "How does this relate to entropy?" - }; - - for (String question : questions) { - System.out.println("User: " + question); - String response = chat(chat1, question); - System.out.println("Assistant: " + truncate(response, 200) + "\n"); - } - - // Demo 2: Session state - System.out.println("\n=== Demo 2: Session State ===\n"); - System.out.println("Session ID: " + session1.getId()); - System.out.println("State: " + session1.getState()); - - // Demo 3: Save and load session - System.out.println("\n=== Demo 3: Session Persistence ===\n"); - String sessionId = session1.getId(); - System.out.println("Saving session: " + sessionId); - - // Load the session - Session loadedSession = loadSession(sessionId); - if (loadedSession != null) { - System.out.println("✓ Session loaded successfully!"); - System.out.println(" Messages in history: " + loadedSession.getMessages().size()); - System.out.println(" State: " + loadedSession.getState()); - - // Continue the conversation - Chat continuedChat = createChat(loadedSession, "assistant"); - System.out.println("\nContinuing conversation..."); - System.out.println("User: Can you summarize what we discussed?"); - String summary = chat(continuedChat, "Can you summarize what we discussed?"); - System.out.println("Assistant: " + truncate(summary, 300)); - } - - // Demo 4: Using tools - System.out.println("\n\n=== Demo 4: Using Tools (Note Taking) ===\n"); - Session session2 = createSession("Bob"); - Chat chat2 = createChat(session2, "assistant"); - - System.out.println("User: Please save a note titled 'Meeting' with content 'Review Q4 goals'"); - String noteResponse = - chat(chat2, "Please save a note titled 'Meeting' with content 'Review Q4 goals'"); - System.out.println("Assistant: " + noteResponse); - showNotes(); - - System.out.println("\n=== Demo Complete ==="); - } - - private String truncate(String text, int maxLength) { - if (text == null) { - return ""; - } - if (text.length() <= maxLength) { - return text; - } - return text.substring(0, maxLength) + "..."; - } - - public static void main(String[] args) { - ChatSessionApp app = new ChatSessionApp(); - - // Check for demo mode flag - boolean demoMode = args.length > 0 && args[0].equals("--demo"); - - if (demoMode) { - app.runDemo(); - } else { - app.runInteractive(); - } - } -} diff --git a/samples/interrupts/src/main/java/com/google/genkit/samples/InterruptsApp.java b/samples/interrupts/src/main/java/com/google/genkit/samples/InterruptsApp.java index 17bc217a6..0884e8c00 100644 --- a/samples/interrupts/src/main/java/com/google/genkit/samples/InterruptsApp.java +++ b/samples/interrupts/src/main/java/com/google/genkit/samples/InterruptsApp.java @@ -22,20 +22,12 @@ import com.google.genkit.GenkitOptions; import com.google.genkit.ai.GenerateOptions; import com.google.genkit.ai.InterruptConfig; -import com.google.genkit.ai.InterruptRequest; import com.google.genkit.ai.ModelResponse; import com.google.genkit.ai.Part; import com.google.genkit.ai.ResumeOptions; import com.google.genkit.ai.Tool; -import com.google.genkit.ai.ToolResponse; -import com.google.genkit.ai.session.Chat; -import com.google.genkit.ai.session.ChatOptions; -import com.google.genkit.ai.session.InMemorySessionStore; -import com.google.genkit.ai.session.Session; -import com.google.genkit.ai.session.SessionOptions; import com.google.genkit.plugins.openai.OpenAIPlugin; import com.google.genkit.prompt.ExecutablePrompt; -import java.util.ArrayList; import java.util.List; import java.util.Map; import java.util.Scanner; @@ -43,7 +35,8 @@ /** * Human-in-the-Loop Application using Interrupts. * - *

This sample demonstrates the interrupt pattern for human-in-the-loop scenarios: + *

This sample demonstrates the interrupt pattern for human-in-the-loop scenarios using the + * {@code generate()} and {@link ExecutablePrompt} APIs: * *

    *
  • Tools that pause execution to request user confirmation @@ -60,39 +53,6 @@ */ public class InterruptsApp { - /** Confirmation input structure. */ - public static class ConfirmationInput { - private String action; - private String details; - private double amount; - - public ConfirmationInput() {} - - public String getAction() { - return action; - } - - public void setAction(String action) { - this.action = action; - } - - public String getDetails() { - return details; - } - - public void setDetails(String details) { - this.details = details; - } - - public double getAmount() { - return amount; - } - - public void setAmount(double amount) { - this.amount = amount; - } - } - /** Transfer request for the interrupt tool input. */ public static class TransferRequest { private String recipient; @@ -155,41 +115,6 @@ public void setReason(String reason) { } } - /** Bank account state. */ - public static class AccountState { - private String accountId; - private double balance; - private List transactions = new ArrayList<>(); - - public AccountState() { - this.accountId = "ACC-" + System.currentTimeMillis() % 10000; - this.balance = 5000.00; // Starting balance - } - - public String getAccountId() { - return accountId; - } - - public double getBalance() { - return balance; - } - - public void addTransaction(String transaction, double amount) { - this.balance += amount; - this.transactions.add(transaction); - } - - public List getTransactions() { - return transactions; - } - - @Override - public String toString() { - return String.format( - "Account: %s, Balance: $%.2f, Transactions: %d", accountId, balance, transactions.size()); - } - } - /** Banking request input for the prompt. */ public static class BankingInput { private String request; @@ -210,12 +135,9 @@ public void setRequest(String request) { } private final Genkit genkit; - private final InMemorySessionStore sessionStore; private final Scanner scanner; // Tools - private Tool getBalanceTool; - private Tool transferMoneyTool; private Tool confirmTransferTool; public InterruptsApp() { @@ -225,30 +147,15 @@ public InterruptsApp() { .plugin(OpenAIPlugin.create()) .build(); - this.sessionStore = new InMemorySessionStore<>(); this.scanner = new Scanner(System.in); initializeTools(); } - @SuppressWarnings("unchecked") private void initializeTools() { - // Get Balance Tool - no confirmation needed - getBalanceTool = - genkit.defineTool( - "getBalance", - "Gets the current account balance", - Map.of("type", "object", "properties", Map.of()), - (Class>) (Class) Map.class, - (ctx, input) -> { - // In a real app, we'd get this from session context - return Map.of("balance", 5000.00, "currency", "USD"); - }); - // Use defineInterrupt to create an interrupt tool that pauses for confirmation. // This is the preferred way to create interrupt tools - it automatically - // handles - // throwing ToolInterruptException with the proper metadata. + // handles throwing ToolInterruptException with the proper metadata. confirmTransferTool = genkit.defineInterrupt( InterruptConfig.builder() @@ -285,207 +192,10 @@ private void initializeTools() { "reason", input.getReason() != null ? input.getReason() : "")) .build()); - - // Transfer Money Tool - executes after confirmation - transferMoneyTool = - genkit.defineTool( - "executeTransfer", - "Executes a confirmed money transfer. Only call this after confirmation.", - Map.of( - "type", - "object", - "properties", - Map.of( - "recipient", - Map.of("type", "string", "description", "Transfer recipient"), - "amount", - Map.of("type", "number", "description", "Amount to transfer"), - "confirmationCode", - Map.of("type", "string", "description", "Confirmation code from user")), - "required", - new String[] {"recipient", "amount", "confirmationCode"}), - (Class>) (Class) Map.class, - (ctx, input) -> { - String recipient = (String) input.get("recipient"); - double amount = ((Number) input.get("amount")).doubleValue(); - String transactionId = "TXN-" + System.currentTimeMillis() % 100000; - - return Map.of( - "status", - "success", - "transactionId", - transactionId, - "recipient", - recipient, - "amount", - amount, - "message", - String.format( - "Successfully transferred $%.2f to %s. Transaction ID: %s", - amount, recipient, transactionId)); - }); - } - - /** Creates a chat session. */ - @SuppressWarnings("unchecked") - public Chat createChat() { - Session session = - genkit.createSession( - SessionOptions.builder() - .store(sessionStore) - .initialState(new AccountState()) - .build()); - - String systemPrompt = - "You are a helpful banking assistant for SecureBank. " - + "You can help customers check their balance and transfer money. " - + "IMPORTANT: For any money transfer, you MUST first use the confirmTransfer tool " - + "to get user confirmation. Never execute a transfer without confirmation. " - + "After the user confirms, use the executeTransfer tool with their confirmation code."; - - return session.chat( - ChatOptions.builder() - .model("openai/gpt-4o-mini") - .system(systemPrompt) - .tools(List.of(getBalanceTool, confirmTransferTool, transferMoneyTool)) - .build()); - } - - /** Handles an interrupt by prompting the user. */ - private ConfirmationOutput handleInterrupt(InterruptRequest interrupt) { - Map metadata = interrupt.getMetadata(); - - System.out.println("\n╔═══════════════════════════════════════════════════════════╗"); - System.out.println("║ ⚠️ CONFIRMATION REQUIRED ⚠️ ║"); - System.out.println("╠═══════════════════════════════════════════════════════════╣"); - System.out.printf( - "║ Transfer: $%.2f to %s%n", metadata.get("amount"), metadata.get("recipient")); - if (metadata.get("reason") != null) { - System.out.printf("║ Reason: %s%n", metadata.get("reason")); - } - System.out.println("╠═══════════════════════════════════════════════════════════╣"); - System.out.println("║ Type 'yes' to confirm or 'no' to cancel ║"); - System.out.println("╚═══════════════════════════════════════════════════════════╝"); - System.out.print("\nYour decision: "); - - String response = scanner.nextLine().trim().toLowerCase(); - boolean confirmed = response.equals("yes") || response.equals("y"); - - if (confirmed) { - System.out.println("✓ Transfer confirmed"); - return new ConfirmationOutput( - true, "User confirmed with code: CONF-" + System.currentTimeMillis() % 10000); - } else { - System.out.println("✗ Transfer cancelled"); - return new ConfirmationOutput(false, "User declined the transfer"); - } - } - - /** Sends a message and handles any interrupts. */ - public String sendWithInterruptHandling(Chat chat, String message) { - try { - ModelResponse response = chat.send(message); - - // Check for pending interrupts - if (chat.hasPendingInterrupts()) { - List interrupts = chat.getPendingInterrupts(); - - for (InterruptRequest interrupt : interrupts) { - // Handle the interrupt (get user confirmation) - ConfirmationOutput userResponse = handleInterrupt(interrupt); - - // Create resume options with the user's response - ToolResponse toolResponse = interrupt.respond(userResponse); - ResumeOptions resume = ResumeOptions.builder().respond(List.of(toolResponse)).build(); - - // Resume the conversation with the user's response - response = - chat.send( - userResponse.isConfirmed() - ? "User confirmed. Proceed with the transfer." - : "User declined. Cancel the transfer.", - Chat.SendOptions.builder().resumeOptions(resume).build()); - } - } - - return response.getText(); - } catch (Exception e) { - return "Error: " + e.getMessage(); - } - } - - /** Interactive chat loop. */ - public void runInteractive() { - System.out.println("╔════════════════════════════════════════════════════════════════╗"); - System.out.println("║ SecureBank - Human-in-the-Loop Banking Assistant ║"); - System.out.println("╚════════════════════════════════════════════════════════════════╝"); - System.out.println(); - System.out.println("This demo shows the interrupt pattern for sensitive operations."); - System.out.println("Money transfers require explicit confirmation before execution."); - System.out.println(); - System.out.println("Try saying:"); - System.out.println(" • 'What's my balance?'"); - System.out.println(" • 'Transfer $100 to John for lunch'"); - System.out.println(" • 'Send $500 to Alice'"); - System.out.println(); - System.out.println("Commands: /status, /quit\n"); - - Chat chat = createChat(); - - while (true) { - System.out.print("You: "); - String input = scanner.nextLine().trim(); - - if (input.isEmpty()) continue; - - if (input.equals("/quit") || input.equals("/exit")) { - System.out.println("\nThank you for banking with SecureBank!"); - break; - } - - if (input.equals("/status")) { - System.out.println("\n" + chat.getSession().getState() + "\n"); - continue; - } - - String response = sendWithInterruptHandling(chat, input); - System.out.println("\nAssistant: " + response + "\n"); - } - } - - /** Demo mode. */ - public void runDemo() { - System.out.println("╔════════════════════════════════════════════════════════════════╗"); - System.out.println("║ Interrupts Demo - Human-in-the-Loop Pattern ║"); - System.out.println("╚════════════════════════════════════════════════════════════════╝"); - System.out.println(); - System.out.println("This demo shows how interrupts work for human-in-the-loop scenarios."); - System.out.println("Watch how the system pauses for confirmation on sensitive operations.\n"); - - Chat chat = createChat(); - - // Demo 1: Check balance (no interrupt) - System.out.println("=== Demo 1: Simple Query (No Interrupt) ===\n"); - System.out.println("Customer: What's my current balance?"); - String response1 = sendWithInterruptHandling(chat, "What's my current balance?"); - System.out.println("Assistant: " + response1 + "\n"); - - // Demo 2: Transfer money (triggers interrupt) - System.out.println("\n=== Demo 2: Transfer Request (Triggers Interrupt) ===\n"); - System.out.println("Customer: Transfer $250 to John Smith for the concert tickets"); - System.out.println("\n[The system will now request confirmation...]\n"); - - // For demo, we'll use a mock confirmation - String response2 = - sendWithInterruptHandling(chat, "Transfer $250 to John Smith for the concert tickets"); - System.out.println("\nAssistant: " + response2); - - System.out.println("\n=== Demo Complete ==="); - System.out.println("Final state: " + chat.getSession().getState()); } /** - * Demo using generate() directly with interrupts (without Chat). + * Demo using generate() directly with interrupts. * *

    This shows how to use interrupts at the lower level generate() API, which is useful when you * don't need session management. @@ -552,9 +262,7 @@ public void runGenerateDemo() { genkit.generate( GenerateOptions.builder() .model(model) - .messages(response.getMessages()) // Include - // previous - // context + .messages(response.getMessages()) // Include previous context .tools(List.of(confirmTransferTool)) .resume(ResumeOptions.builder().respond(responseData.getToolResponse()).build()) .build()); @@ -655,18 +363,12 @@ public void runPromptDemo() { public static void main(String[] args) { InterruptsApp app = new InterruptsApp(); - boolean demoMode = args.length > 0 && args[0].equals("--demo"); - boolean generateDemo = args.length > 0 && args[0].equals("--generate"); boolean promptDemo = args.length > 0 && args[0].equals("--prompt"); if (promptDemo) { app.runPromptDemo(); - } else if (generateDemo) { - app.runGenerateDemo(); - } else if (demoMode) { - app.runDemo(); } else { - app.runInteractive(); + app.runGenerateDemo(); } } } diff --git a/samples/multi-agent/README.md b/samples/multi-agent/README.md deleted file mode 100644 index 52fd41169..000000000 --- a/samples/multi-agent/README.md +++ /dev/null @@ -1,204 +0,0 @@ -# Genkit Multi-Agent Sample - -This sample demonstrates multi-agent orchestration patterns using Genkit Java, where specialized agents handle different domains and a triage agent routes requests. - -## Features Demonstrated - -- **Multi-Agent Architecture** - Triage agent routing to specialized agents -- **Specialized Agents** - Reservation, menu, and order agents -- **Agent-as-Tool Pattern** - Agents can be used as tools for delegation -- **Session Management** - Track customer state across interactions -- **Tool Integration** - Agents with domain-specific tools - -## Prerequisites - -- Java 21+ -- Maven 3.6+ -- OpenAI API key - -## Running the Sample - -### Option 1: Direct Run - -```bash -# Set your OpenAI API key -export OPENAI_API_KEY=your-api-key-here - -# Navigate to the sample directory -cd java/samples/multi-agent - -# Run the sample -./run.sh -# Or: mvn compile exec:java -``` - -### Option 2: With Genkit Dev UI - -```bash -# Set your OpenAI API key -export OPENAI_API_KEY=your-api-key-here - -# Navigate to the sample directory -cd java/samples/multi-agent - -# Run with Genkit CLI -genkit start -- ./run.sh -``` - -## Architecture - -``` -┌─────────────────────────────────────────────────────────────┐ -│ Customer Request │ -└─────────────────────────────────────────────────────────────┘ - │ - ▼ -┌─────────────────────────────────────────────────────────────┐ -│ Triage Agent │ -│ Routes requests to specialized agents based on intent │ -└─────────────────────────────────────────────────────────────┘ - │ │ │ - ▼ ▼ ▼ -┌─────────────────┐ ┌─────────────────┐ ┌─────────────────┐ -│ Reservation │ │ Menu │ │ Order │ -│ Agent │ │ Agent │ │ Agent │ -│ │ │ │ │ │ -│ • makeRes │ │ • getMenu │ │ • placeOrder │ -│ • cancelRes │ │ • getDietInfo │ │ • getOrderStatus│ -└─────────────────┘ └─────────────────┘ └─────────────────┘ -``` - -## Agents - -### Triage Agent -The main entry point that analyzes customer requests and routes them to the appropriate specialized agent. - -### Reservation Agent -Handles table reservations: -- Make new reservations -- Cancel existing reservations -- Check availability - -### Menu Agent -Handles menu-related queries: -- Get menu items -- Dietary information -- Recommendations - -### Order Agent -Handles food orders: -- Place orders -- Check order status -- Modify orders - -## Available Tools - -| Tool | Agent | Description | -|------|-------|-------------| -| `makeReservation` | Reservation | Makes a new reservation | -| `cancelReservation` | Reservation | Cancels an existing reservation | -| `getMenu` | Menu | Returns menu items | -| `placeOrder` | Order | Places a food order | - -## Example Interactions - -The sample runs as an interactive CLI application: - -``` -🍽️ Welcome to the Restaurant! -Type 'quit' to exit. - -You: I'd like to make a reservation for 4 people tomorrow at 7pm - -Agent: I'd be happy to help you with your reservation. Let me set that up for you. - -[Reservation Agent handles the request] - -Reservation confirmed! Your confirmation number is RES-1234. -- Date: 2024-01-16 -- Time: 19:00 -- Party size: 4 - -You: What's on the menu? - -Agent: [Menu Agent handles the request] - -Here's our current menu: -- Appetizers: ... -- Main Courses: ... -- Desserts: ... -``` - -## Session State - -The sample tracks customer state across interactions: - -```java -public class CustomerState { - private String customerId; - private String currentAgent; - private List reservations; - private List orders; -} -``` - -## Code Highlights - -### Defining an Agent - -```java -Agent reservationAgent = genkit.defineAgent( - AgentConfig.builder() - .name("reservationAgent") - .model("openai/gpt-4o") - .system("You are a helpful reservation agent for a restaurant...") - .tools(List.of(makeReservationTool, cancelReservationTool)) - .build()); -``` - -### Agent-as-Tool Pattern - -```java -// Agents can be used as tools for delegation -Tool reservationAgentTool = reservationAgent.asTool(); - -Agent triageAgent = genkit.defineAgent( - AgentConfig.builder() - .name("triageAgent") - .model("openai/gpt-4o") - .system("Route requests to the appropriate agent...") - .tools(List.of(reservationAgentTool, menuAgentTool, orderAgentTool)) - .build()); -``` - -### Session-Based Chat - -```java -Session session = genkit.createSession( - SessionOptions.builder() - .sessionStore(sessionStore) - .initialState(new CustomerState()) - .build()); - -Chat chat = session.chat(ChatOptions.builder() - .model("openai/gpt-4o") - .agent(triageAgent) - .build()); - -String response = chat.send("I'd like to make a reservation"); -``` - -## Development UI - -When running with `genkit start`, access the Dev UI at http://localhost:4000 to: - -- View registered agents and tools -- Test individual agents -- Inspect traces showing agent routing -- View tool calls and responses - -## See Also - -- [Genkit Java README](../../README.md) -- [Chat Sessions Sample](../chat-session/README.md) -- [Interrupts Sample](../interrupts/README.md) diff --git a/samples/multi-agent/run.sh b/samples/multi-agent/run.sh deleted file mode 100755 index 7a055a49c..000000000 --- a/samples/multi-agent/run.sh +++ /dev/null @@ -1,4 +0,0 @@ -#!/bin/bash -# Run script for Genkit DotPrompt Sample -cd "$(dirname "$0")" -mvn exec:java diff --git a/samples/multi-agent/src/main/java/com/google/genkit/samples/MultiAgentApp.java b/samples/multi-agent/src/main/java/com/google/genkit/samples/MultiAgentApp.java deleted file mode 100644 index a7fab34ff..000000000 --- a/samples/multi-agent/src/main/java/com/google/genkit/samples/MultiAgentApp.java +++ /dev/null @@ -1,560 +0,0 @@ -/* - * 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. - * - * SPDX-License-Identifier: Apache-2.0 - */ - -package com.google.genkit.samples; - -import com.google.genkit.Genkit; -import com.google.genkit.GenkitOptions; -import com.google.genkit.ai.Agent; -import com.google.genkit.ai.AgentConfig; -import com.google.genkit.ai.GenerationConfig; -import com.google.genkit.ai.Tool; -import com.google.genkit.ai.session.Chat; -import com.google.genkit.ai.session.ChatOptions; -import com.google.genkit.ai.session.InMemorySessionStore; -import com.google.genkit.ai.session.Session; -import com.google.genkit.ai.session.SessionOptions; -import com.google.genkit.plugins.openai.OpenAIPlugin; -import java.util.ArrayList; -import java.util.HashMap; -import java.util.List; -import java.util.Map; -import java.util.Scanner; - -/** - * Multi-Agent Customer Service Application. - * - *

    This sample demonstrates the multi-agent pattern where: - * - *

      - *
    • A triage agent routes requests to specialized agents - *
    • Specialized agents handle specific domains (reservations, menu, etc.) - *
    • Agents can be used as tools for delegation - *
    - * - *

    To run: - * - *

      - *
    1. Set the OPENAI_API_KEY environment variable - *
    2. Run: mvn exec:java -pl samples/multi-agent - *
    - */ -public class MultiAgentApp { - - /** Customer state for tracking context. */ - public static class CustomerState { - private String customerId; - private String currentAgent; - private List reservations = new ArrayList<>(); - private List orders = new ArrayList<>(); - - public CustomerState() { - this.customerId = "customer-" + System.currentTimeMillis(); - this.currentAgent = "triage"; - } - - public String getCustomerId() { - return customerId; - } - - public String getCurrentAgent() { - return currentAgent; - } - - public void setCurrentAgent(String agent) { - this.currentAgent = agent; - } - - public List getReservations() { - return reservations; - } - - public void addReservation(String reservation) { - this.reservations.add(reservation); - } - - public List getOrders() { - return orders; - } - - public void addOrder(String order) { - this.orders.add(order); - } - - @Override - public String toString() { - return String.format( - "Customer: %s, Agent: %s, Reservations: %d, Orders: %d", - customerId, currentAgent, reservations.size(), orders.size()); - } - } - - private final Genkit genkit; - private final InMemorySessionStore sessionStore; - - // Agents - private Agent triageAgent; - private Agent reservationAgent; - private Agent menuAgent; - private Agent orderAgent; - - // Tools - private Tool makeReservationTool; - private Tool cancelReservationTool; - private Tool getMenuTool; - private Tool placeOrderTool; - - public MultiAgentApp() { - // Initialize Genkit - this.genkit = - Genkit.builder() - .options(GenkitOptions.builder().devMode(true).reflectionPort(3101).build()) - .plugin(OpenAIPlugin.create()) - .build(); - - this.sessionStore = new InMemorySessionStore<>(); - - // Initialize tools and agents - initializeTools(); - initializeAgents(); - } - - @SuppressWarnings("unchecked") - private void initializeTools() { - // Reservation Tool - makeReservationTool = - genkit.defineTool( - "makeReservation", - "Makes a restaurant reservation for the customer", - Map.of( - "type", - "object", - "properties", - Map.of( - "date", - Map.of("type", "string", "description", "Date in YYYY-MM-DD format"), - "time", - Map.of("type", "string", "description", "Time in HH:MM format"), - "partySize", - Map.of("type", "integer", "description", "Number of guests")), - "required", - new String[] {"date", "time", "partySize"}), - (Class>) (Class) Map.class, - (ctx, input) -> { - String date = (String) input.get("date"); - String time = (String) input.get("time"); - Integer partySize = (Integer) input.get("partySize"); - String confirmationId = "RES-" + System.currentTimeMillis() % 10000; - - Map result = new HashMap<>(); - result.put("status", "confirmed"); - result.put("confirmationId", confirmationId); - result.put("date", date); - result.put("time", time); - result.put("partySize", partySize); - result.put( - "message", - String.format( - "Reservation confirmed for %d guests on %s at %s. Confirmation: %s", - partySize, date, time, confirmationId)); - return result; - }); - - // Cancel Reservation Tool - cancelReservationTool = - genkit.defineTool( - "cancelReservation", - "Cancels an existing reservation", - Map.of( - "type", - "object", - "properties", - Map.of( - "confirmationId", - Map.of("type", "string", "description", "The reservation confirmation ID")), - "required", - new String[] {"confirmationId"}), - (Class>) (Class) Map.class, - (ctx, input) -> { - String confirmationId = (String) input.get("confirmationId"); - Map result = new HashMap<>(); - result.put("status", "cancelled"); - result.put("confirmationId", confirmationId); - result.put("message", "Reservation " + confirmationId + " has been cancelled."); - return result; - }); - - // Menu Tool - getMenuTool = - genkit.defineTool( - "getMenu", - "Gets the current restaurant menu", - Map.of( - "type", - "object", - "properties", - Map.of( - "category", - Map.of( - "type", - "string", - "description", - "Menu category: appetizers, mains, desserts, drinks, or all", - "enum", - new String[] {"appetizers", "mains", "desserts", "drinks", "all"}))), - (Class>) (Class) Map.class, - (ctx, input) -> { - String category = - input.get("category") != null ? (String) input.get("category") : "all"; - Map menu = new HashMap<>(); - - if (category.equals("all") || category.equals("appetizers")) { - menu.put( - "appetizers", - List.of( - Map.of( - "name", - "Bruschetta", - "price", - 8.99, - "description", - "Toasted bread with tomatoes"), - Map.of( - "name", - "Calamari", - "price", - 12.99, - "description", - "Fried squid rings"))); - } - if (category.equals("all") || category.equals("mains")) { - menu.put( - "mains", - List.of( - Map.of( - "name", - "Grilled Salmon", - "price", - 24.99, - "description", - "Atlantic salmon with herbs"), - Map.of( - "name", - "Ribeye Steak", - "price", - 32.99, - "description", - "12oz prime ribeye"), - Map.of( - "name", - "Pasta Primavera", - "price", - 18.99, - "description", - "Seasonal vegetables"))); - } - if (category.equals("all") || category.equals("desserts")) { - menu.put( - "desserts", - List.of( - Map.of( - "name", - "Tiramisu", - "price", - 9.99, - "description", - "Classic Italian dessert"), - Map.of("name", "Cheesecake", "price", 8.99, "description", "NY style"))); - } - if (category.equals("all") || category.equals("drinks")) { - menu.put( - "drinks", - List.of( - Map.of("name", "House Wine", "price", 8.99, "description", "Red or white"), - Map.of( - "name", - "Craft Beer", - "price", - 6.99, - "description", - "Local selection"))); - } - - return menu; - }); - - // Order Tool - placeOrderTool = - genkit.defineTool( - "placeOrder", - "Places a food order for pickup or delivery", - Map.of( - "type", - "object", - "properties", - Map.of( - "items", - Map.of( - "type", - "array", - "items", - Map.of("type", "string"), - "description", - "List of menu item names to order"), - "orderType", - Map.of( - "type", - "string", - "description", - "pickup or delivery", - "enum", - new String[] {"pickup", "delivery"})), - "required", - new String[] {"items", "orderType"}), - (Class>) (Class) Map.class, - (ctx, input) -> { - @SuppressWarnings("unchecked") - List items = (List) input.get("items"); - String orderType = (String) input.get("orderType"); - String orderId = "ORD-" + System.currentTimeMillis() % 10000; - - Map result = new HashMap<>(); - result.put("status", "confirmed"); - result.put("orderId", orderId); - result.put("items", items); - result.put("orderType", orderType); - result.put("estimatedTime", orderType.equals("pickup") ? "20 minutes" : "45 minutes"); - result.put( - "message", - String.format( - "Order %s placed for %s. Items: %s. Ready in %s.", - orderId, - orderType, - String.join(", ", items), - orderType.equals("pickup") ? "20 minutes" : "45 minutes")); - return result; - }); - } - - @SuppressWarnings("unchecked") - private void initializeAgents() { - // Reservation Agent - handles booking and cancellation - // Note: genkit.defineAgent automatically registers the agent - reservationAgent = - genkit.defineAgent( - AgentConfig.builder() - .name("reservationAgent") - .description( - "Handles restaurant reservations. Transfer to this agent when the customer " - + "wants to make, modify, or cancel a reservation.") - .system( - "You are a reservation specialist for an upscale restaurant. " - + "Help customers make, modify, or cancel reservations. " - + "Always confirm the date, time, and party size before making a reservation. " - + "Be professional and courteous.") - .model("openai/gpt-4o-mini") - .tools(List.of(makeReservationTool, cancelReservationTool)) - .config(GenerationConfig.builder().temperature(0.3).build()) - .build()); - - // Menu Agent - provides menu information - menuAgent = - genkit.defineAgent( - AgentConfig.builder() - .name("menuAgent") - .description( - "Provides menu information. Transfer to this agent when the customer " - + "wants to know about menu items, prices, or recommendations.") - .system( - "You are a menu expert at an upscale restaurant. " - + "Help customers explore the menu, understand dishes, and get recommendations. " - + "Use the getMenu tool to retrieve current menu items. " - + "Be knowledgeable about ingredients and preparation methods.") - .model("openai/gpt-4o-mini") - .tools(List.of(getMenuTool)) - .config(GenerationConfig.builder().temperature(0.5).build()) - .build()); - - // Order Agent - handles food orders - orderAgent = - genkit.defineAgent( - AgentConfig.builder() - .name("orderAgent") - .description( - "Handles food orders for pickup or delivery. Transfer to this agent when " - + "the customer wants to place an order.") - .system( - "You are an order specialist for a restaurant. " - + "Help customers place orders for pickup or delivery. " - + "Confirm all items before placing the order. " - + "Provide accurate time estimates.") - .model("openai/gpt-4o-mini") - .tools(List.of(placeOrderTool, getMenuTool)) - .config(GenerationConfig.builder().temperature(0.3).build()) - .build()); - - // Triage Agent - routes to specialized agents - triageAgent = - genkit.defineAgent( - AgentConfig.builder() - .name("triageAgent") - .description("Main customer service agent that routes requests to specialists") - .system( - "You are the main customer service agent for The Golden Fork restaurant. " - + "Your job is to understand what the customer needs and transfer them to the right specialist.\n\n" - + "IMPORTANT: To transfer to another agent, you MUST call the appropriate agent tool. " - + "Do NOT just say you are transferring - you must actually invoke the tool:\n" - + "- reservationAgent: for reservations (booking, canceling, modifying)\n" - + "- menuAgent: for menu questions, recommendations, or dietary info\n" - + "- orderAgent: for placing orders (pickup or delivery)\n\n" - + "When a customer needs help with a specific task, call the corresponding agent tool immediately. " - + "You can handle general greetings and questions, but for specific tasks, always use the tools.") - .model("openai/gpt-4o-mini") - .agents( - List.of( - reservationAgent.getConfig(), - menuAgent.getConfig(), - orderAgent.getConfig())) - .config(GenerationConfig.builder().temperature(0.7).build()) - .build()); - } - - /** Creates a chat session with the triage agent. */ - @SuppressWarnings("unchecked") - public Chat createChat() { - Session session = - genkit.createSession( - SessionOptions.builder() - .store(sessionStore) - .initialState(new CustomerState()) - .build()); - - // Get all tools including sub-agents as tools - Genkit handles the registry - List> allTools = genkit.getAllToolsForAgent(triageAgent); - - // Agent registry is automatically available from the session - no need to pass - // explicitly - return session.chat( - ChatOptions.builder() - .model("openai/gpt-4o-mini") - .system(triageAgent.getSystem()) - .tools(allTools) - .build()); - } - - /** Interactive chat loop. */ - public void runInteractive() { - Scanner scanner = new Scanner(System.in); - - System.out.println("╔════════════════════════════════════════════════════════════════╗"); - System.out.println("║ The Golden Fork Restaurant - Multi-Agent Customer Service ║"); - System.out.println("╚════════════════════════════════════════════════════════════════╝"); - System.out.println(); - System.out.println("Available agents:"); - System.out.println(" • Triage Agent - Routes your requests"); - System.out.println(" • Reservation Agent - Handles bookings"); - System.out.println(" • Menu Agent - Menu information and recommendations"); - System.out.println(" • Order Agent - Pickup and delivery orders"); - System.out.println(); - System.out.println("Commands:"); - System.out.println(" /status - Show current state"); - System.out.println(" /quit - Exit"); - System.out.println(); - System.out.println("How can we help you today?\n"); - - Chat chat = createChat(); - - while (true) { - System.out.print("You: "); - String input = scanner.nextLine().trim(); - - if (input.isEmpty()) continue; - - if (input.equals("/quit") || input.equals("/exit")) { - System.out.println("\nThank you for visiting The Golden Fork!"); - break; - } - - if (input.equals("/status")) { - String currentAgent = chat.getCurrentAgentName(); - System.out.println("\nState: " + chat.getSession().getState()); - System.out.println( - "Current Agent: " + (currentAgent != null ? currentAgent : "triage (default)") + "\n"); - continue; - } - - try { - String response = chat.send(input).getText(); - System.out.println("\nAssistant: " + response + "\n"); - } catch (Exception e) { - System.out.println("\nError: " + e.getMessage() + "\n"); - } - } - - scanner.close(); - } - - /** Demo mode. */ - public void runDemo() { - System.out.println("╔════════════════════════════════════════════════════════════════╗"); - System.out.println("║ Multi-Agent Demo - Restaurant Customer Service ║"); - System.out.println("╚════════════════════════════════════════════════════════════════╝"); - System.out.println(); - - Chat chat = createChat(); - - // Demo conversation - String[] messages = { - "Hi, I'd like to make a reservation for this weekend", - "Saturday at 7pm for 4 people", - "Thanks! Also, what's on your dessert menu?", - "I'd like to place a pickup order for the Tiramisu and Cheesecake" - }; - - for (String message : messages) { - System.out.println("Customer: " + message); - try { - String response = chat.send(message).getText(); - System.out.println("\nAssistant: " + truncate(response, 300) + "\n"); - Thread.sleep(1000); // Pause for readability - } catch (Exception e) { - System.out.println("Error: " + e.getMessage()); - } - } - - System.out.println("\n=== Demo Complete ==="); - System.out.println("Final state: " + chat.getSession().getState()); - } - - private String truncate(String text, int maxLength) { - if (text == null || text.length() <= maxLength) return text; - return text.substring(0, maxLength) + "..."; - } - - public static void main(String[] args) { - MultiAgentApp app = new MultiAgentApp(); - - boolean demoMode = args.length > 0 && args[0].equals("--demo"); - - if (demoMode) { - app.runDemo(); - } else { - app.runInteractive(); - } - } -} diff --git a/samples/multi-agent/src/main/resources/logback.xml b/samples/multi-agent/src/main/resources/logback.xml deleted file mode 100644 index d63c14f8a..000000000 --- a/samples/multi-agent/src/main/resources/logback.xml +++ /dev/null @@ -1,13 +0,0 @@ - - - - %d{HH:mm:ss.SSS} [%thread] %-5level %logger{36} - %msg%n - - - - - - - - - diff --git a/samples/openai/src/main/java/com/google/genkit/samples/SessionSample.java b/samples/openai/src/main/java/com/google/genkit/samples/SessionSample.java deleted file mode 100644 index 967898467..000000000 --- a/samples/openai/src/main/java/com/google/genkit/samples/SessionSample.java +++ /dev/null @@ -1,345 +0,0 @@ -/* - * 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. - * - * SPDX-License-Identifier: Apache-2.0 - */ - -package com.google.genkit.samples; - -import com.google.genkit.Genkit; -import com.google.genkit.GenkitOptions; -import com.google.genkit.ai.*; -import com.google.genkit.ai.session.*; -import com.google.genkit.plugins.openai.OpenAIPlugin; -import java.util.HashMap; -import java.util.List; -import java.util.Map; - -/** - * Sample application demonstrating session-based multi-turn conversations. - * - *

    This example shows how to: - Create sessions with persistent state - Conduct multi-turn - * conversations with automatic history management - Use multiple conversation threads within a - * session - Implement custom session stores for persistence - Use tools within session-based chats - * - *

    To run: 1. Set the OPENAI_API_KEY environment variable 2. Run: mvn exec:java - * -Dexec.mainClass=com.google.genkit.samples.SessionSample - */ -public class SessionSample { - - /** Custom session state to track user preferences and conversation context. */ - public static class UserState { - private String userName; - private String preferredLanguage; - private int messageCount; - - public UserState() {} - - public UserState(String userName) { - this.userName = userName; - this.preferredLanguage = "English"; - this.messageCount = 0; - } - - public String getUserName() { - return userName; - } - - public void setUserName(String userName) { - this.userName = userName; - } - - public String getPreferredLanguage() { - return preferredLanguage; - } - - public void setPreferredLanguage(String preferredLanguage) { - this.preferredLanguage = preferredLanguage; - } - - public int getMessageCount() { - return messageCount; - } - - public void incrementMessageCount() { - this.messageCount++; - } - } - - public static void main(String[] args) throws Exception { - // Create Genkit with OpenAI plugin - Genkit genkit = - Genkit.builder() - .options(GenkitOptions.builder().devMode(true).reflectionPort(3100).build()) - .plugin(OpenAIPlugin.create()) - .build(); - - // Define a tool for the conversation - @SuppressWarnings("unchecked") - Tool, Map> reminderTool = - genkit.defineTool( - "setReminder", - "Sets a reminder for the user", - Map.of( - "type", - "object", - "properties", - Map.of( - "message", - Map.of("type", "string", "description", "The reminder message"), - "time", - Map.of( - "type", - "string", - "description", - "When to remind (e.g., '5 minutes', 'tomorrow')")), - "required", - new String[] {"message", "time"}), - (Class>) (Class) Map.class, - (ctx, input) -> { - Map result = new HashMap<>(); - result.put("status", "success"); - result.put( - "message", "Reminder set: " + input.get("message") + " at " + input.get("time")); - return result; - }); - - System.out.println("=== Session-Based Chat Demo ===\n"); - - // Example 1: Basic session with multi-turn conversation - basicSessionExample(genkit); - - // Example 2: Session with custom state - sessionWithStateExample(genkit); - - // Example 3: Multiple conversation threads - multiThreadExample(genkit); - - // Example 4: Session with tools - sessionWithToolsExample(genkit, reminderTool); - - // Example 5: Loading existing sessions - sessionPersistenceExample(genkit); - - System.out.println("\n=== Demo Complete ==="); - } - - /** Demonstrates basic session creation and multi-turn conversation. */ - private static void basicSessionExample(Genkit genkit) throws Exception { - System.out.println("--- Example 1: Basic Multi-Turn Conversation ---\n"); - - // Create a session - Session session = genkit.createSession(); - System.out.println("Created session: " + session.getId()); - - // Create a chat with system prompt - Chat chat = - session.chat( - ChatOptions.builder() - .model("openai/gpt-4o-mini") - .system("You are a helpful assistant. Keep your responses brief and friendly.") - .build()); - - // Multi-turn conversation - history is automatically managed - System.out.println("\nUser: What is the capital of France?"); - ModelResponse response1 = chat.send("What is the capital of France?"); - System.out.println("Assistant: " + response1.getText()); - - System.out.println("\nUser: What's the population?"); - ModelResponse response2 = chat.send("What's the population?"); - System.out.println("Assistant: " + response2.getText()); - - System.out.println("\nUser: What language do they speak there?"); - ModelResponse response3 = chat.send("What language do they speak there?"); - System.out.println("Assistant: " + response3.getText()); - - // Show conversation history - System.out.println("\n--- Conversation History ---"); - for (Message msg : chat.getHistory()) { - System.out.println( - msg.getRole() - + ": " - + msg.getText().substring(0, Math.min(50, msg.getText().length())) - + "..."); - } - System.out.println(); - } - - /** Demonstrates session with custom state management. */ - private static void sessionWithStateExample(Genkit genkit) throws Exception { - System.out.println("--- Example 2: Session with Custom State ---\n"); - - // Create session with initial state - Session session = - genkit.createSession( - SessionOptions.builder().initialState(new UserState("Alice")).build()); - - System.out.println("Created session for user: " + session.getState().getUserName()); - - // Create chat - Chat chat = - session.chat( - ChatOptions.builder() - .model("openai/gpt-4o-mini") - .system( - "You are a helpful assistant. The user's name is " - + session.getState().getUserName() - + ".") - .build()); - - // Send message and update state - ModelResponse response = chat.send("Hello! Can you remember my name?"); - System.out.println("Assistant: " + response.getText()); - - // Update session state - UserState state = session.getState(); - state.incrementMessageCount(); - session.updateState(state).join(); - - System.out.println("Message count: " + session.getState().getMessageCount()); - System.out.println(); - } - - /** Demonstrates multiple conversation threads within a session. */ - private static void multiThreadExample(Genkit genkit) throws Exception { - System.out.println("--- Example 3: Multiple Conversation Threads ---\n"); - - Session session = genkit.createSession(); - - // Create chat for general conversation - Chat generalChat = - session.chat( - "general", - ChatOptions.builder() - .model("openai/gpt-4o-mini") - .system("You are a helpful general assistant.") - .build()); - - // Create chat for coding help - Chat codingChat = - session.chat( - "coding", - ChatOptions.builder() - .model("openai/gpt-4o-mini") - .system("You are an expert programmer. Provide concise code examples.") - .build()); - - // Use different threads for different topics - System.out.println("General thread:"); - ModelResponse generalResponse = generalChat.send("What's a good recipe for pasta?"); - System.out.println( - "Response: " - + generalResponse - .getText() - .substring(0, Math.min(100, generalResponse.getText().length())) - + "...\n"); - - System.out.println("Coding thread:"); - ModelResponse codingResponse = codingChat.send("How do I reverse a string in Java?"); - System.out.println( - "Response: " - + codingResponse - .getText() - .substring(0, Math.min(100, codingResponse.getText().length())) - + "...\n"); - - // Continue in general thread - context is preserved per thread - System.out.println("Back to general thread:"); - ModelResponse followUp = generalChat.send("What ingredients do I need for that?"); - System.out.println( - "Response: " - + followUp.getText().substring(0, Math.min(100, followUp.getText().length())) - + "...\n"); - } - - /** Demonstrates using tools within session-based chats. */ - private static void sessionWithToolsExample(Genkit genkit, Tool reminderTool) - throws Exception { - System.out.println("--- Example 4: Session with Tools ---\n"); - - Session session = genkit.createSession(); - - @SuppressWarnings("unchecked") - Chat chat = - session.chat( - ChatOptions.builder() - .model("openai/gpt-4o-mini") - .system("You are a helpful assistant that can set reminders for users.") - .tools(List.of((Tool) reminderTool)) - .build()); - - System.out.println("User: Remind me to buy groceries in 1 hour"); - ModelResponse response = chat.send("Remind me to buy groceries in 1 hour"); - System.out.println("Assistant: " + response.getText()); - System.out.println(); - } - - /** Demonstrates session persistence - saving and loading sessions. */ - private static void sessionPersistenceExample(Genkit genkit) throws Exception { - System.out.println("--- Example 5: Session Persistence ---\n"); - - // Create a custom session store (using in-memory for this example) - InMemorySessionStore store = new InMemorySessionStore<>(); - - // Create session with the store - Session session = - genkit.createSession( - SessionOptions.builder() - .store(store) - .sessionId("persistent-session-001") - .initialState(new UserState("Bob")) - .build()); - - // Have a conversation - Chat chat = - session.chat( - ChatOptions.builder() - .model("openai/gpt-4o-mini") - .system("You are a helpful assistant.") - .build()); - - chat.send("Hello, I'm learning about AI"); - chat.send("What's machine learning?"); - - System.out.println("Original session ID: " + session.getId()); - System.out.println("Messages in session: " + chat.getHistory().size()); - - // Load the session later (simulating app restart) - Session loadedSession = - genkit - .loadSession( - "persistent-session-001", SessionOptions.builder().store(store).build()) - .get(); - - if (loadedSession != null) { - System.out.println("\nLoaded session ID: " + loadedSession.getId()); - System.out.println("User name from state: " + loadedSession.getState().getUserName()); - System.out.println("Messages preserved: " + loadedSession.getMessages().size()); - - // Continue the conversation - Chat continuedChat = - loadedSession.chat( - ChatOptions.builder() - .model("openai/gpt-4o-mini") - .system("You are a helpful assistant.") - .build()); - - System.out.println("\nContinuing conversation..."); - ModelResponse response = continuedChat.send("Can you summarize what we discussed?"); - System.out.println("Assistant: " + response.getText()); - } - System.out.println(); - } -}