override fun install()

in agents/agents-features/agents-features-snapshot/src/commonMain/kotlin/ai/koog/agents/snapshot/feature/Persistence.kt [94:147]
46 lines of code
6 McCabe index (conditional complexity)

        override fun install(
            config: PersistenceFeatureConfig,
            pipeline: AIAgentGraphPipeline,
        ): Persistence {
            val persistence = Persistence(config.storage)
            persistence.rollbackStrategy = config.rollbackStrategy
            persistence.rollbackToolRegistry = config.rollbackToolRegistry

            pipeline.interceptStrategyStarting(this) { ctx ->
                val strategy = ctx.strategy as AIAgentGraphStrategy<*, *>

                require(strategy.metadata.uniqueNames) {
                    "Checkpoint feature requires unique node names in the strategy metadata"
                }

                val checkpoint = persistence.rollbackToLatestCheckpoint(ctx.context)

                if (checkpoint != null) {
                    logger.info { "Restoring checkpoint: ${checkpoint.checkpointId} to node ${checkpoint.nodePath}" }
                } else {
                    logger.info { "No non-tombstone checkpoint found, starting from the beginning" }
                }
            }

            pipeline.interceptNodeExecutionCompleted(this) { eventCtx ->
                if (persistence.isTechnicalNode(eventCtx.node.id)) {
                    return@interceptNodeExecutionCompleted
                }

                if (config.enableAutomaticPersistence) {
                    val parent = persistence.getLatestCheckpoint(eventCtx.context.agentId)
                    persistence.createCheckpointAfterNode(
                        agentContext = eventCtx.context,
                        nodePath = eventCtx.context.executionInfo.path(),
                        lastOutput = eventCtx.output,
                        lastOutputType = eventCtx.outputType,
                        version = parent?.version?.plus(1) ?: 0L,
                    )
                }
            }

            pipeline.interceptStrategyCompleted(this) { ctx ->
                if (config.enableAutomaticPersistence && config.rollbackStrategy == RollbackStrategy.Default) {
                    val parent = persistence.getLatestCheckpoint(ctx.context.agentId)
                    persistence.createTombstoneCheckpoint(
                        ctx.context.agentId,
                        persistence.clock.now(),
                        parent?.version?.plus(1) ?: 0L
                    )
                }
            }

            return persistence
        }