-
Notifications
You must be signed in to change notification settings - Fork 240
Add RollbackToolRegistry
in the Persistency
feature in order to roll back tool calls with side effects when checkpointing. Make AIAgent
state-manageable and introduce AIAgentService
to manage multiple uniform running agents. Deprecate concurrent unsafe AIAgent.asTool
in favor of AIAgentService.createAgentTool
#873
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Merged
Ololoshechkin
merged 34 commits into
develop
from
vbr/rollback-sideeffect-tools-in-persistency
Sep 30, 2025
Merged
Add RollbackToolRegistry
in the Persistency
feature in order to roll back tool calls with side effects when checkpointing. Make AIAgent
state-manageable and introduce AIAgentService
to manage multiple uniform running agents. Deprecate concurrent unsafe AIAgent.asTool
in favor of AIAgentService.createAgentTool
#873
Ololoshechkin
merged 34 commits into
develop
from
vbr/rollback-sideeffect-tools-in-persistency
Sep 30, 2025
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Qodana for JVM862 new problems were found
@@ Code coverage @@
+ 70% total lines covered
12679 lines analyzed, 8901 lines covered
# Calculated according to the filters of your coverage tool ☁️ View the detailed Qodana report Contact Qodana teamContact us at qodana-support@jetbrains.com
|
AIAgent
manageable, and introduce AIAgentService
to manage uniform running agents
AIAgent
manageable, and introduce AIAgentService
to manage uniform running agentsRollbackToolRegistry
in the Persistency
feature in order to roll back tool calls with side effects when checkpointing. Also make AIAgent
manageable, and introduce AIAgentService
to manage uniform running agents
RollbackToolRegistry
in the Persistency
feature in order to roll back tool calls with side effects when checkpointing. Also make AIAgent
manageable, and introduce AIAgentService
to manage uniform running agentsRollbackToolRegistry
in the Persistency
feature in order to roll back tool calls with side effects when checkpointing. Also make AIAgent
manageable, and introduce AIAgentService
to manage multiple uniform running agents
RollbackToolRegistry
in the Persistency
feature in order to roll back tool calls with side effects when checkpointing. Also make AIAgent
manageable, and introduce AIAgentService
to manage multiple uniform running agentsRollbackToolRegistry
in the Persistency
feature in order to roll back tool calls with side effects when checkpointing. Also make AIAgent
manageable, and introduce AIAgentService
to manage multiple uniform running agents. Other than that, deprecate concurrent unsafe AIAgent.asTool
in favor of AIAgentService.createAgentTool
RollbackToolRegistry
in the Persistency
feature in order to roll back tool calls with side effects when checkpointing. Also make AIAgent
manageable, and introduce AIAgentService
to manage multiple uniform running agents. Other than that, deprecate concurrent unsafe AIAgent.asTool
in favor of AIAgentService.createAgentTool
RollbackToolRegistry
in the Persistency
feature in order to roll back tool calls with side effects when checkpointing. Make AIAgent
manageable, and introduce AIAgentService
to manage multiple uniform running agents. Deprecate concurrent unsafe AIAgent.asTool
in favor of AIAgentService.createAgentTool
RollbackToolRegistry
in the Persistency
feature in order to roll back tool calls with side effects when checkpointing. Make AIAgent
manageable, and introduce AIAgentService
to manage multiple uniform running agents. Deprecate concurrent unsafe AIAgent.asTool
in favor of AIAgentService.createAgentTool
RollbackToolRegistry
in the Persistency
feature in order to roll back tool calls with side effects when checkpointing. Make AIAgent
state-manageable and introduce AIAgentService
to manage multiple uniform running agents. Deprecate concurrent unsafe AIAgent.asTool
in favor of AIAgentService.createAgentTool
15 tasks
…back tool calls with side effects when checkpointing.
…DO: fix structured concurrency
…ool` because it's not working with parallel tools anyway
…Tool via it, fix example in docs
Rizzen
approved these changes
Sep 30, 2025
8a3a423
to
e80d688
Compare
valery1707
added a commit
to valery1707/koog
that referenced
this pull request
Oct 17, 2025
…ice.createAgentTool` Refs JetBrains#873
15 tasks
karloti
pushed a commit
to karloti/koog
that referenced
this pull request
Oct 21, 2025
…oll back tool calls with side effects when checkpointing. Make `AIAgent` state-manageable and introduce `AIAgentService` to manage multiple uniform running agents. Deprecate concurrent unsafe `AIAgent.asTool` in favor of `AIAgentService.createAgentTool` (JetBrains#873) 1. Add `RollbackToolRegistry` in the `Persistency` feature in order to roll back tool calls with side effects when checkpointing: ```kotlin val agent = AIAgent( toolRegistry = ToolRegistry { tool(::createUser) tool(::sendMessage) tool(::inviteMember) }, ... ) { install(Persistency) { storage = myStateStorageImpl rollbackToolRegistry = RollbackToolRegistry { // TOOL -> "REVERSE" TOOL registerRollback(::createUser, ::deleteUser) registerRollback(::sendMessage, ::undoMessage) registerRollback(::inviteMember, ::revokeInvitation) } } } ``` 2. Now you can also manage (create, run, find, access running state) multiple AI Agents using `AIAgentService` : ```kotlin val agentService = AIAgentService( toolRegistry = ToolRegistry { tool(::createUser) tool(::sendMessage) tool(::inviteMember) }, ... ) { install(Persistency) { storage = myStateStorageImpl rollbackToolRegistry = RollbackToolRegistry { registerRollback(::createUser, ::deleteUser) registerRollback(::sendMessage, ::undoMessage) registerRollback(::inviteMember, ::revokeInvitation) } } } ``` And then you can create instances of AIAgent, and manage their running state. This is particularly useful to rollback long-running operations if user realizes that agent took the wrong direction: ```kotlin // user creates new agent: post("/agent") { val input = call.receive<String>() val agent = agentService.createAgent() launch { agent.run(input } call.respond(agent.id) } // user checks the agent's state: get("/agent") { val id = call.receive<String>() val agent = agentService.agentById(id) if (!agent.finished()) call.respondText("Agent is still running...") else call.respond(agent.resultIfReady()!!) } // user asks running agent to rollback: data class RollbackRequest(val agentId: String, val checkpoint: String) put("/agent/rollback") { val userRequest = call.receive<RollbackRequest>() val agent = agentService.agentById(userRequest.agentId) if (agent.finished()) { call.respondText("Agent has already finished!") } else { // Rolling back agent to a checkpoint agent.withRunningContext { withPersistency(this) { ctx -> rollbackToCheckpoint(userRequest.checkpoint, ctx) } } call.respond(HttpStatusCode.OK) } } ``` 3. Make `AIAgent` explicitly single-run. Previous semantic was: `AIAgent.run` can be called multiple times, but if called in parallel -- throws exception that agent is currently running. This was creating a hard to manage contract for `run()` and was very error-prone (see example of such errors below in `4`) 4. Fix `Agent.asTool` for parallel use JetBrains#864 . Previously -- `AIAgentTool` was holding an instance of `AIAgent` and running it. When LLM decided to call tools in parallel, because of the previous contract of `AIAgent` it was failing in runtime with exception. Hence, `AIAgent.asTool` was broken. Now -- intended usage is `AIAgentService.createAgentTool()`. `AIAgent.asTool` is now deprecated and is currently working correctly via `AIAgentService.fromAgent(this).createAgentTool()` (i.e. it creates an instance of `AIAgentService` from the current `AIAgent`, then creates a tool that holds a new copy of `AIAgent`). ## Motivation and Context <!-- Why is this change needed? What problem does it solve? --> ## Breaking Changes <!-- Will users need to update their code or configurations? --> --- #### Type of the changes - [x] New feature (non-breaking change which adds functionality) - [ ] Bug fix (non-breaking change which fixes an issue) - [ ] Breaking change (fix or feature that would cause existing functionality to change) - [ ] Documentation update - [ ] Tests improvement - [ ] Refactoring #### Checklist - [x] The pull request has a description of the proposed change - [x] I read the [Contributing Guidelines](https://github.com/JetBrains/koog/blob/main/CONTRIBUTING.md) before opening the pull request - [x] The pull request uses **`develop`** as the base branch - [x] Tests for the changes have been added - [ ] All new and existing tests passed ##### Additional steps for pull requests adding a new feature - [ ] An issue describing the proposed change exists - [ ] The pull request includes a link to the issue - [ ] The change was discussed and approved in the issue - [x] Docs have been added / updated
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
RollbackToolRegistry
in thePersistency
feature in order to roll back tool calls with side effects when checkpointing:AIAgentService
:And then you can create instances of AIAgent, and manage their running state. This is particularly useful to rollback long-running operations if user realizes that agent took the wrong direction:
Make
AIAgent
explicitly single-run. Previous semantic was:AIAgent.run
can be called multiple times, but if called in parallel -- throws exception that agent is currently running.This was creating a hard to manage contract for
run()
and was very error-prone (see example of such errors below in4
)Fix
Agent.asTool
for parallel useAIAgent.asTool()
fails withparallelTools
due to agent is already running. #864 . Previously --AIAgentTool
was holding an instance ofAIAgent
and running it. When LLM decided to call tools in parallel, because of the previous contract ofAIAgent
it was failing in runtime with exception. Hence,AIAgent.asTool
was broken.Now -- intended usage is
AIAgentService.createAgentTool()
.AIAgent.asTool
is now deprecated and is currently working correctly viaAIAgentService.fromAgent(this).createAgentTool()
(i.e. it creates an instance ofAIAgentService
from the currentAIAgent
, then creates a tool that holds a new copy ofAIAgent
).Motivation and Context
Breaking Changes
Type of the changes
Checklist
develop
as the base branchAdditional steps for pull requests adding a new feature