fixes

Author: Rizzen

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

@@ -95,7 +95,6 @@ public fun tombstoneCheckpoint(time: Instant, version: Long): AgentCheckpointDat
  */
 public fun AgentCheckpointData.toAgentContextData(
     rollbackStrategy: RollbackStrategy,
-    agentId: String,
     additionalRollbackAction: suspend (AIAgentContext) -> Unit = {}
 ): AgentContextData {
     @Suppress("DEPRECATION")

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

@@ -359,7 +359,7 @@ public class Persistence(
         val checkpoint: AgentCheckpointData? = getCheckpointById(agentContext.agentId, checkpointId)
         if (checkpoint != null) {
             agentContext.store(
-                checkpoint.toAgentContextData(rollbackStrategy, agentContext.agentId) { context ->
+                checkpoint.toAgentContextData(rollbackStrategy) { context ->
                     messageHistoryDiff(
                         currentMessages = context.llm.prompt.messages,
                         checkpointMessages = checkpoint.messageHistory
@@ -375,7 +375,6 @@ public class Persistence(
                                 } catch (_: Exception) {
                                     null
                                 }
-
                                 rollbackTool.executeUnsafe(toolArgs)
                             }
                         }
@@ -390,8 +389,8 @@ public class Persistence(
      * Returns the difference only.
      * ex: current messages: [1, 2, 3, 4, 5, 6, 7], checkpoint messages: [1, 2, 3, 4, 5] -> diff messages: 6, 7
      *
-     * Only works for the scenario when current chat histor is AHEAD of the checkpoint (i.e. we are restoring BACKWARDS in time),
-     * otherwise will return an empty list!
+     * Only works for the scenario when the current chat history is AHEAD of the checkpoint (i.e. we are restoring BACKWARDS in time),
+     *  otherwise it will return an empty list!
      * */
     private fun messageHistoryDiff(currentMessages: List<Message>, checkpointMessages: List<Message>): List<Message> {
         if (checkpointMessages.size > currentMessages.size) {
@@ -424,7 +423,7 @@ public class Persistence(
             return null
         }
 
-        agentContext.store(checkpoint.toAgentContextData(rollbackStrategy, agentContext.agentId))
+        agentContext.store(checkpoint.toAgentContextData(rollbackStrategy))
         return checkpoint
     }
 }