Removed requrement for unique node names for Peristency feature (#1288)

<!--
Thank you for opening a pull request!

Please add a brief description of the proposed change here.
Also, please tick the appropriate points in the checklist below.
-->

## Motivation and Context
Aim of the PR is to remove requirement (and check) for unique graph node
names and migrate to node path usage.

## Breaking Changes
<!-- Will users need to update their code or configurations? -->
Yes
---

#### Type of the changes
- [ ] New feature (non-breaking change which adds functionality)
- [ ] Bug fix (non-breaking change which fixes an issue)
- [x] 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
- [ ] Tests for the changes have been added
- [x] 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
- [ ] Docs have been added / updated

Author: Rizzen

agents/agents-core/src/commonMain/kotlin/ai/koog/agents/core/agent/context/AgentContextData.kt

@@ -9,7 +9,7 @@ import kotlinx.serialization.json.JsonElement
 @InternalAgentsApi
 public class AgentContextData(
     internal val messageHistory: List<Message>,
-    internal val nodeId: String,
+    internal val nodePath: String,
     internal val lastInput: JsonElement,
     internal val rollbackStrategy: RollbackStrategy,
     internal val additionalRollbackActions: suspend (AIAgentContext) -> Unit = {}

agents/agents-core/src/commonMain/kotlin/ai/koog/agents/core/agent/entity/AIAgentGraphStrategy.kt

@@ -7,6 +7,7 @@ import ai.koog.agents.core.agent.context.RollbackStrategy
 import ai.koog.agents.core.agent.context.getAgentContextData
 import ai.koog.agents.core.agent.context.removeAgentContextData
 import ai.koog.agents.core.agent.context.with
+import ai.koog.agents.core.agent.execution.DEFAULT_AGENT_PATH_SEPARATOR
 import ai.koog.agents.core.annotation.InternalAgentsApi
 import ai.koog.agents.core.utils.runCatchingCancellable
 import io.github.oshai.kotlinlogging.KotlinLogging
@@ -97,13 +98,13 @@ public class AIAgentGraphStrategy<TInput, TOutput>(
 
     @OptIn(InternalAgentsApi::class)
     private suspend fun restoreDefault(agentContext: AIAgentContext, data: AgentContextData) {
-        val nodeId = data.nodeId
+        val nodePath = data.nodePath
 
         // Perform additional cleanup (ex: rollback tools):
         data.additionalRollbackActions(agentContext)
 
         // Set current graph node:
-        setExecutionPoint(nodeId, data.lastInput)
+        setExecutionPoint(nodePath, data.lastInput)
 
         // Reset the message history:
         agentContext.llm.withPrompt {
@@ -114,17 +115,15 @@ public class AIAgentGraphStrategy<TInput, TOutput>(
     /**
      * Finds and sets the node for the strategy based on the provided context.
      */
-    public fun setExecutionPoint(nodeId: String, input: JsonElement) {
-        val fullPath = metadata.nodesMap.keys.firstOrNull {
-            val segments = it.split(":")
-            segments.last() == nodeId
-        } ?: throw IllegalArgumentException("Node $nodeId not found")
+    public fun setExecutionPoint(nodePath: String, input: JsonElement) {
+        // we drop first because it's agent's id, we don't need it here
+        val segments = nodePath.split(DEFAULT_AGENT_PATH_SEPARATOR).drop(1)
 
-        val segments = fullPath.split(":")
         if (segments.isEmpty()) {
-            throw IllegalArgumentException("Invalid node path: $fullPath")
+            throw IllegalArgumentException("Invalid node path: $nodePath")
         }
 
+        val actualPath = segments.joinToString(DEFAULT_AGENT_PATH_SEPARATOR)
         val strategyName = segments.firstOrNull() ?: return
 
         // getting the very first segment (it should be a root strategy node)
@@ -134,25 +133,25 @@ public class AIAgentGraphStrategy<TInput, TOutput>(
         // restoring the current node for each subgraph including strategy
         val segmentsInbetween = segments.drop(1).dropLast(1)
         for (segment in segmentsInbetween) {
-            currentNode as? ExecutionPointNode
-                ?: throw IllegalStateException("Node ${currentNode?.name} does not have subnodes")
+            val currNode = currentNode as? ExecutionPointNode
+                ?: throw IllegalStateException("Restore for path $nodePath failed: one of middle segments is not a subgraph")
 
-            currentPath = "$currentPath:$segment"
+            currentPath = "$currentPath${DEFAULT_AGENT_PATH_SEPARATOR}$segment"
             val nextNode = metadata.nodesMap[currentPath]
             if (nextNode is ExecutionPointNode) {
-                currentNode.enforceExecutionPoint(nextNode, input)
+                currNode.enforceExecutionPoint(nextNode, input)
                 currentNode = nextNode
             }
         }
 
         // forcing the very last segment to the latest pre-leaf node to complete the chain
-        val leaf = metadata.nodesMap[fullPath] ?: throw IllegalStateException("Node ${segments.last()} not found")
+        val leaf = metadata.nodesMap[actualPath] ?: throw IllegalStateException("Node $actualPath not found")
         val inputType = leaf.inputType
 
         val actualInput = serializer.decodeFromJsonElement(serializer.serializersModule.serializer(inputType), input)
         leaf.let {
             currentNode as? ExecutionPointNode
-                ?: throw IllegalStateException("Node ${currentNode?.name} does not have subnodes")
+                ?: throw IllegalStateException("Node ${currentNode?.name} is not a valid leaf node")
             currentNode.enforceExecutionPoint(it, actualInput)
         }
     }

agents/agents-core/src/commonMain/kotlin/ai/koog/agents/core/agent/execution/AgentExecutionInfo.kt

@@ -16,18 +16,6 @@ public data class AgentExecutionInfo(
     public val parent: AgentExecutionInfo?,
     public val partName: String
 ) {
-
-    /**
-     * A companion object for the `AgentExecutionInfo` class that provides utility constants.
-     */
-    public companion object {
-
-        /**
-         * The default separator used to join parts of a path.
-         */
-        public val defaultPathSeparator: CharSequence = "/"
-    }
-
     /**
      * Constructs a path string representing the sequence of `partName` values from the current
      * `AgentExecutionInfo` instance to the top-most parent, joined by the specified separator.
@@ -36,7 +24,7 @@ public data class AgentExecutionInfo(
      * @return A string representing the path constructed from `partName` values.
      */
     public fun path(separator: String? = null): String {
-        val separator = separator ?: defaultPathSeparator
+        val separator = separator ?: DEFAULT_AGENT_PATH_SEPARATOR
 
         return buildList {
             var current: AgentExecutionInfo? = this@AgentExecutionInfo

agents/agents-core/src/commonMain/kotlin/ai/koog/agents/core/agent/execution/AgentNodePath.kt

@@ -0,0 +1,17 @@
+package ai.koog.agents.core.agent.execution
+
+/**
+ * The default separator used to join parts of an agent's execution path.
+ */
+public const val DEFAULT_AGENT_PATH_SEPARATOR: String = "/"
+
+/**
+ * Joins the given parts into a single path string using the specified separator.
+ *
+ * @param parts The parts to join into a path.
+ * @param separator The separator to use between parts. Defaults to [DEFAULT_AGENT_PATH_SEPARATOR].
+ * @return A string representing the joined path.
+ */
+public fun path(vararg parts: String, separator: String = DEFAULT_AGENT_PATH_SEPARATOR): String {
+    return parts.joinToString(separator)
+}

agents/agents-core/src/commonMain/kotlin/ai/koog/agents/core/dsl/builder/AIAgentSubgraphBuilder.kt

@@ -12,6 +12,7 @@ import ai.koog.agents.core.agent.entity.FinishNode
 import ai.koog.agents.core.agent.entity.StartNode
 import ai.koog.agents.core.agent.entity.SubgraphMetadata
 import ai.koog.agents.core.agent.entity.ToolSelectionStrategy
+import ai.koog.agents.core.agent.execution.DEFAULT_AGENT_PATH_SEPARATOR
 import ai.koog.agents.core.annotation.InternalAgentsApi
 import ai.koog.agents.core.tools.Tool
 import ai.koog.prompt.llm.LLModel
@@ -229,7 +230,7 @@ public abstract class AIAgentSubgraphBuilderBase<Input, Output> {
     }
 
     private fun getNodePath(node: AIAgentNodeBase<*, *>, parentPath: String): String {
-        return "$parentPath:${node.id}"
+        return "$parentPath${DEFAULT_AGENT_PATH_SEPARATOR}${node.id}"
     }
 
     internal fun buildSubgraphMetadata(
@@ -248,12 +249,11 @@ public abstract class AIAgentSubgraphBuilderBase<Input, Output> {
         }
 
         // Validate that all nodes have unique names within the subgraph
-        val names = subgraphNodes.keys.map { it.split(":").last() }
-        val uniqueNames = names.toSet().size == names.size
+        val names = subgraphNodes.keys
 
         return SubgraphMetadata(
             nodesMap = subgraphNodes,
-            uniqueNames = uniqueNames
+            uniqueNames = names.toSet().size == names.size
         )
     }
 

agents/agents-core/src/commonTest/kotlin/ai/koog/agents/core/agent/execution/AgentExecutionInfoTest.kt

@@ -48,7 +48,7 @@ class AgentExecutionInfoTest {
 
     @Test
     fun testDefaultPathSeparator() {
-        assertEquals("/", AgentExecutionInfo.defaultPathSeparator)
+        assertEquals("/", DEFAULT_AGENT_PATH_SEPARATOR)
     }
 
     @Test

agents/agents-core/src/commonTest/kotlin/ai/koog/agents/core/feature/AIAgentPipelineTest.kt

@@ -8,6 +8,7 @@ import ai.koog.agents.core.agent.config.AIAgentConfig
 import ai.koog.agents.core.agent.entity.AIAgentGraphStrategy
 import ai.koog.agents.core.agent.entity.AIAgentSubgraph.Companion.FINISH_NODE_PREFIX
 import ai.koog.agents.core.agent.entity.AIAgentSubgraph.Companion.START_NODE_PREFIX
+import ai.koog.agents.core.agent.execution.DEFAULT_AGENT_PATH_SEPARATOR
 import ai.koog.agents.core.dsl.builder.forwardTo
 import ai.koog.agents.core.dsl.builder.strategy
 import ai.koog.agents.core.dsl.extension.nodeDoNothing
@@ -32,7 +33,6 @@ import ai.koog.agents.core.feature.handler.AgentLifecycleEventType.ToolCallFaile
 import ai.koog.agents.core.feature.handler.AgentLifecycleEventType.ToolCallStarting
 import ai.koog.agents.core.feature.handler.AgentLifecycleEventType.ToolValidationFailed
 import ai.koog.agents.core.tools.ToolRegistry
-import ai.koog.agents.testing.agent.agentExecutionPath
 import ai.koog.agents.testing.tools.DummyTool
 import ai.koog.agents.testing.tools.getMockExecutor
 import ai.koog.prompt.dsl.prompt
@@ -817,5 +817,7 @@ class AIAgentPipelineTest {
         )
     }
 
+    private fun agentExecutionPath(vararg parts: String) = parts.joinToString(DEFAULT_AGENT_PATH_SEPARATOR)
+
     //endregion Private Methods
 }

agents/agents-features/agents-features-snapshot/src/commonMain/kotlin/ai/koog/agents/snapshot/feature/AgentCheckpointData.kt

@@ -21,15 +21,17 @@ import kotlin.uuid.Uuid
  *
  * @property checkpointId The unique identifier of the checkpoint. This allows tracking and restoring the agent's session to a specific state.
  * @property messageHistory A list of messages exchanged in the session up to the checkpoint. Messages include interactions between the user, system, assistant, and tools.
- * @property nodeId The identifier of the node where the checkpoint was created.
- * @property lastInput Serialized input received for node with [nodeId]
+ * @property nodePath The identifier of the node where the checkpoint was created.
+ * @property lastInput Serialized input received for node with [nodePath]
  * @property properties Additional data associated with the checkpoint. This can be used to store additional information about the agent's state.
+ * @property createdAt The timestamp when the checkpoint was created.
+ * @property version The version of the checkpoint data structure
  */
 @Serializable
 public data class AgentCheckpointData(
     val checkpointId: String,
     val createdAt: Instant,
-    val nodeId: String,
+    val nodePath: String,
     val lastInput: JsonElement,
     val messageHistory: List<Message>,
     val version: Long,
@@ -48,7 +50,7 @@ public fun tombstoneCheckpoint(time: Instant, version: Long): AgentCheckpointDat
     return AgentCheckpointData(
         checkpointId = Uuid.random().toString(),
         createdAt = time,
-        nodeId = PersistenceUtils.TOMBSTONE_CHECKPOINT_NAME,
+        nodePath = PersistenceUtils.TOMBSTONE_CHECKPOINT_NAME,
         lastInput = JsonNull,
         messageHistory = emptyList(),
         properties = mapOf(PersistenceUtils.TOMBSTONE_CHECKPOINT_NAME to JsonPrimitive(true)),
@@ -67,11 +69,12 @@ public fun tombstoneCheckpoint(time: Instant, version: Long): AgentCheckpointDat
  */
 public fun AgentCheckpointData.toAgentContextData(
     rollbackStrategy: RollbackStrategy,
+    agentId: String,
     additionalRollbackAction: suspend (AIAgentContext) -> Unit = {}
 ): AgentContextData {
     return AgentContextData(
         messageHistory = messageHistory,
-        nodeId = nodeId,
+        nodePath = nodePath,
         lastInput = lastInput,
         rollbackStrategy,
         additionalRollbackAction

…g/agents/snapshot/feature/Persistency.kt …g/agents/snapshot/feature/Persistence.ktagents/agents-features/agents-features-snapshot/src/commonMain/kotlin/ai/koog/agents/snapshot/feature/Persistency.kt renamed to agents/agents-features/agents-features-snapshot/src/commonMain/kotlin/ai/koog/agents/snapshot/feature/Persistence.kt agents/agents-features/agents-features-snapshot/src/commonMain/kotlin/ai/koog/agents/snapshot/feature/Persistency.kt renamed to agents/agents-features/agents-features-snapshot/src/commonMain/kotlin/ai/koog/agents/snapshot/feature/Persistence.kt

@@ -10,6 +10,7 @@ import ai.koog.agents.core.agent.context.store
 import ai.koog.agents.core.agent.entity.AIAgentGraphStrategy
 import ai.koog.agents.core.agent.entity.AIAgentStorageKey
 import ai.koog.agents.core.agent.entity.AIAgentSubgraph
+import ai.koog.agents.core.agent.execution.DEFAULT_AGENT_PATH_SEPARATOR
 import ai.koog.agents.core.agent.featureOrThrow
 import ai.koog.agents.core.annotation.InternalAgentsApi
 import ai.koog.agents.core.feature.AIAgentGraphFeature
@@ -80,20 +81,6 @@ public class Persistence(
      */
     public var rollbackToolRegistry: RollbackToolRegistry = RollbackToolRegistry {}
 
-    /**
-     * Represents the identifier of the current node being executed within the agent pipeline.
-     *
-     * This property is used to track the state of the agent's execution and is updated whenever
-     * the agent begins processing a new node.
-     * It plays a crucial role in maintaining the agent's
-     * state across checkpoints and ensuring accurate state restoration during rollbacks.
-     *
-     * The value is nullable, indicating that there might be no current node under execution
-     * (e.g., when the pipeline is idle or has not started).
-     */
-    public var currentNodeId: String? = null
-        private set
-
     /**
      * Companion object implementing agent feature, handling [Persistence] creation and installation.
      */
@@ -122,7 +109,7 @@ public class Persistence(
                 val checkpoint = persistence.rollbackToLatestCheckpoint(ctx.context)
 
                 if (checkpoint != null) {
-                    logger.info { "Restoring checkpoint: ${checkpoint.checkpointId} to node ${checkpoint.nodeId}" }
+                    logger.info { "Restoring checkpoint: ${checkpoint.checkpointId} to node ${checkpoint.nodePath}" }
                 } else {
                     logger.info { "No non-tombstone checkpoint found, starting from the beginning" }
                 }
@@ -137,18 +124,14 @@ public class Persistence(
                     val parent = persistence.getLatestCheckpoint(eventCtx.context.agentId)
                     persistence.createCheckpoint(
                         agentContext = eventCtx.context,
-                        nodeId = eventCtx.node.id,
+                        nodePath = eventCtx.context.executionInfo.path(),
                         lastInput = eventCtx.input,
                         lastInputType = eventCtx.inputType,
                         version = parent?.version?.plus(1) ?: 0L,
                     )
                 }
             }
 
-            pipeline.interceptNodeExecutionStarting(this) { eventCtx ->
-                persistence.currentNodeId = eventCtx.node.id
-            }
-
             pipeline.interceptStrategyCompleted(this) { ctx ->
                 if (config.enableAutomaticPersistence && config.rollbackStrategy == RollbackStrategy.Default) {
                     val parent = persistence.getLatestCheckpoint(ctx.context.agentId)
@@ -182,7 +165,7 @@ public class Persistence(
      */
     public suspend fun createCheckpoint(
         agentContext: AIAgentContext,
-        nodeId: String,
+        nodePath: String,
         lastInput: Any?,
         lastInputType: KType,
         version: Long,
@@ -192,7 +175,7 @@ public class Persistence(
 
         if (inputJson == null) {
             logger.warn {
-                "Failed to serialize input of type $lastInputType for checkpoint creation for $nodeId, skipping..."
+                "Failed to serialize input of type $lastInputType for checkpoint creation for $nodePath, skipping..."
             }
             return null
         }
@@ -201,7 +184,7 @@ public class Persistence(
             return@readSession AgentCheckpointData(
                 checkpointId = checkpointId ?: Uuid.random().toString(),
                 messageHistory = prompt.messages,
-                nodeId = nodeId,
+                nodePath = agentContext.executionInfo.path(),
                 lastInput = inputJson,
                 createdAt = Clock.System.now(),
                 version = version,
@@ -263,17 +246,17 @@ public class Persistence(
      * with the specified message history and input data.
      *
      * @param agentContext The context of the agent to modify
-     * @param nodeId The ID of the node to set as the current execution point
+     * @param nodePath The path to the node inside the agent's graph
      * @param messageHistory The message history to set for the agent
      * @param input The input data to set for the agent
      */
     public fun setExecutionPoint(
         agentContext: AIAgentContext,
-        nodeId: String,
+        nodePath: String,
         messageHistory: List<Message>,
         input: JsonElement
     ) {
-        agentContext.store(AgentContextData(messageHistory, nodeId, input, rollbackStrategy))
+        agentContext.store(AgentContextData(messageHistory, agentContext.agentId + DEFAULT_AGENT_PATH_SEPARATOR + nodePath, input, rollbackStrategy))
     }
 
     /**
@@ -297,7 +280,7 @@ public class Persistence(
         val checkpoint: AgentCheckpointData? = getCheckpointById(agentContext.agentId, checkpointId)
         if (checkpoint != null) {
             agentContext.store(
-                checkpoint.toAgentContextData(rollbackStrategy) { context ->
+                checkpoint.toAgentContextData(rollbackStrategy, agentContext.agentId) { context ->
                     messageHistoryDiff(
                         currentMessages = context.llm.prompt.messages,
                         checkpointMessages = checkpoint.messageHistory
@@ -356,7 +339,7 @@ public class Persistence(
             return null
         }
 
-        agentContext.store(checkpoint.toAgentContextData(rollbackStrategy))
+        agentContext.store(checkpoint.toAgentContextData(rollbackStrategy, agentContext.agentId))
         return checkpoint
     }
 }