KG-576. Finalize pipeline features after agent run for the StatefulSingleUseAIAgent

- Finalize the feature processors after an agent run. After an agent run, it does not make sense to keep running feature processors. We can safely close them and re-open later on the second run if needed;
- Add tests to verify graph and functional agents close processors for features.

Author: sdubov

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

@@ -16,7 +16,6 @@ import ai.koog.agents.core.feature.config.FeatureConfig
 import ai.koog.agents.core.feature.pipeline.AIAgentGraphPipeline
 import ai.koog.agents.core.tools.ToolRegistry
 import ai.koog.prompt.executor.model.PromptExecutor
-import ai.koog.utils.io.Closeable
 import io.github.oshai.kotlinlogging.KotlinLogging
 import kotlinx.datetime.Clock
 import kotlin.reflect.KType
@@ -59,7 +58,7 @@ public open class GraphAIAgent<Input, Output>(
 ) : StatefulSingleUseAIAgent<Input, Output, AIAgentGraphContextBase>(
     logger = logger,
     id = id,
-), Closeable {
+) {
 
     private companion object {
         private val logger = KotlinLogging.logger {}

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

@@ -18,7 +18,7 @@ import kotlin.uuid.Uuid
 /**
  * Abstract base class representing a single-use AI agent with state.
  *
- * This AI agent is designed to execute a specific long-running strategy only once, and provides API to monitor and manage it's state.
+ * This AI agent is designed to execute a specific long-running strategy only once and provides an API to monitor and manage its state.
  *
  * It maintains internal states including its running status, whether it was started, its result (if available), and
  * the root context associated with its execution. The class enforces safe state transitions and provides
@@ -86,87 +86,85 @@ public abstract class StatefulSingleUseAIAgent<Input, Output, TContext : AIAgent
         }
 
         val runId = Uuid.random().toString()
-
-        pipeline.prepareFeatures()
+        val context = prepareContext(agentInput, runId)
 
         return withContext(
             AgentRunInfoContextElement(
-                agentId = this@StatefulSingleUseAIAgent.id,
+                agentId = id,
                 runId = runId,
                 agentConfig = agentConfig,
                 strategyName = strategy.name
             )
         ) {
-            val context = prepareContext(agentInput, runId)
+            pipeline.withPreparedPipeline {
+                agentStateMutex.withLock {
+                    state = State.Running(context)
+                }
 
-            agentStateMutex.withLock {
-                state = State.Running(context)
-            }
+                logger.debug {
+                    formatLog(
+                        agentId = id,
+                        runId = runId,
+                        message = "Starting agent execution"
+                    )
+                }
 
-            logger.debug {
-                formatLog(
-                    agentId = this@StatefulSingleUseAIAgent.id,
+                pipeline.onAgentStarting<Input, Output>(
                     runId = runId,
-                    message = "Starting agent execution"
+                    agent = this@StatefulSingleUseAIAgent,
+                    context = context
                 )
-            }
-
-            pipeline.onAgentStarting<Input, Output>(
-                runId = runId,
-                agent = this@StatefulSingleUseAIAgent,
-                context = context
-            )
 
-            val result = try {
-                strategy.execute(context = context, input = agentInput)
-            } catch (e: Throwable) {
-                logger.error(e) { "Execution exception reported by server!" }
-                pipeline.onAgentExecutionFailed(
-                    agentId = this@StatefulSingleUseAIAgent.id,
-                    runId = runId,
-                    throwable = e
-                )
-                agentStateMutex.withLock { state = State.Failed(e) }
-                throw e
-            }
+                val result = try {
+                    strategy.execute(context = context, input = agentInput)
+                } catch (e: Throwable) {
+                    logger.error(e) { "Execution exception reported by server!" }
+                    pipeline.onAgentExecutionFailed(
+                        agentId = id,
+                        runId = runId,
+                        throwable = e
+                    )
+                    agentStateMutex.withLock { state = State.Failed(e) }
+                    throw e
+                }
 
-            logger.debug {
-                formatLog(
-                    agentId = this@StatefulSingleUseAIAgent.id,
+                logger.debug {
+                    formatLog(
+                        agentId = id,
+                        runId = runId,
+                        message = "Finished agent execution"
+                    )
+                }
+                pipeline.onAgentCompleted(
+                    agentId = id,
                     runId = runId,
-                    message = "Finished agent execution"
+                    result = result
                 )
-            }
-            pipeline.onAgentCompleted(
-                agentId = this@StatefulSingleUseAIAgent.id,
-                runId = runId,
-                result = result
-            )
 
-            agentStateMutex.withLock {
-                state = if (result != null) {
-                    State.Finished(result)
-                } else {
-                    State.Failed(Exception("result is null"))
+                agentStateMutex.withLock {
+                    state = if (result != null) {
+                        State.Finished(result)
+                    } else {
+                        State.Failed(Exception("result is null"))
+                    }
                 }
-            }
 
-            return@withContext result ?: error("result is null")
+                result ?: error("result is null")
+            }
         }
     }
 
     /**
      * Closes the AI Agent and performs necessary cleanup operations.
      *
      * This method is a suspending function that ensures that the AI Agent's resources are released
-     * when it is no longer needed. It notifies the pipeline of the agent's closure and ensures
+     * when it is no longer required. It notifies the pipeline of the agent's closure and ensures
      * that any associated features or stream providers are properly closed.
      *
      * Overrides the `close` method to implement agent-specific shutdown logic.
      */
     final override suspend fun close() {
-        pipeline.onAgentClosing(agentId = this@StatefulSingleUseAIAgent.id)
-        pipeline.closeFeaturesStreamProviders()
+        // TODO: Delete Closeable interface from [AIAgent] declaration.
     }
 
     /**
@@ -202,6 +200,22 @@ public abstract class StatefulSingleUseAIAgent<Input, Output, TContext : AIAgent
      */
     protected fun formatLog(agentId: String, runId: String, message: String): String =
         "[agent id: $agentId, run id: $runId] $message"
+
+    /**
+     * Executes the provided block within a prepared pipeline context.
+     * Ensures that the necessary feature resources are initialized before the block is invoked
+     * and properly cleaned up after finishing the block.
+     *
+     * @return The result of the block's execution.
+     */
+    private suspend fun <T> AIAgentPipeline.withPreparedPipeline(block: suspend () -> T): T =
+        try {
+            prepareAllFeatures()
+            block.invoke()
+        } finally {
+            onAgentClosing(agentId = id)
+            closeAllFeaturesMessageProcessors()
+        }
 }
 
 /**

agents/agents-core/src/commonMain/kotlin/ai/koog/agents/core/feature/pipeline/AIAgentFunctionalPipeline.kt

@@ -33,6 +33,6 @@ public class AIAgentFunctionalPipeline(clock: Clock = Clock.System) : AIAgentPip
             pipeline = this,
         )
 
-        registeredFeatures[feature.key] = RegisteredFeature(featureImpl, featureConfig)
+        super.install(feature.key, featureConfig, featureImpl)
     }
 }

agents/agents-core/src/commonMain/kotlin/ai/koog/agents/core/feature/pipeline/AIAgentGraphPipeline.kt

@@ -48,12 +48,12 @@ public class AIAgentGraphPipeline(clock: Clock = Clock.System) : AIAgentPipeline
      * The feature's message processors are initialized during installation.
      *
      * @param TConfig The type of the feature configuration
-     * @param TFeature The type of the feature being installed
+     * @param TFeatureImpl The type of the feature being installed
      * @param feature The feature implementation to be installed
      * @param configure A lambda to customize the feature configuration
      */
-    public fun <TConfig : FeatureConfig, TFeature : Any> install(
-        feature: AIAgentGraphFeature<TConfig, TFeature>,
+    public fun <TConfig : FeatureConfig, TFeatureImpl : Any> install(
+        feature: AIAgentGraphFeature<TConfig, TFeatureImpl>,
         configure: TConfig.() -> Unit,
     ) {
         val featureConfig = feature.createInitialConfig().apply { configure() }
@@ -62,7 +62,7 @@ public class AIAgentGraphPipeline(clock: Clock = Clock.System) : AIAgentPipeline
             pipeline = this,
         )
 
-        registeredFeatures[feature.key] = RegisteredFeature(featureImpl, featureConfig)
+        super.install(feature.key, featureConfig, featureImpl)
     }
 
     //region Trigger Node Handlers