Introduced AIAgentRunSession

agents/agents-core/src/androidMain/kotlin/ai/koog/agents/core/agent/AIAgent.kt

@@ -3,7 +3,9 @@
 package ai.koog.agents.core.agent
 
 import ai.koog.agents.core.agent.config.AIAgentConfig
+import ai.koog.agents.core.agent.context.AIAgentContext
 import ai.koog.agents.core.agent.entity.AIAgentGraphStrategy
+import ai.koog.agents.core.agent.session.AIAgentRunSession
 import ai.koog.agents.core.tools.ToolRegistry
 import ai.koog.prompt.executor.model.PromptExecutor
 import ai.koog.prompt.llm.LLModel
@@ -17,11 +19,9 @@ public actual abstract class AIAgent<Input, Output> : Closeable {
     public actual abstract val id: String
     public actual abstract val agentConfig: AIAgentConfig
 
-    public actual abstract suspend fun getState(): AIAgentState<Output>
+    public actual abstract suspend fun run(agentInput: Input, sessionId: String?): Output
 
-    public actual open suspend fun result(): Output = AIAgentHelper.result(this)
-
-    public actual abstract suspend fun run(agentInput: Input): Output
+    public actual abstract fun createSession(sessionId: String?): AIAgentRunSession<Input, Output, out AIAgentContext>
 
     public actual companion object {
         @OptIn(markerClass = [ExperimentalUuidApi::class])

agents/agents-core/src/androidMain/kotlin/ai/koog/agents/core/agent/AIAgentService.kt

@@ -33,11 +33,6 @@ public actual abstract class AIAgentService<Input, Output, TAgent : AIAgent<Inpu
     public actual abstract suspend fun removeAgent(agent: TAgent): Boolean
     public actual abstract suspend fun removeAgentWithId(id: String): Boolean
     public actual abstract suspend fun agentById(id: String): TAgent?
-    public actual abstract suspend fun listAllAgents(): List<TAgent>
-    public actual abstract suspend fun listActiveAgents(): List<TAgent>
-    public actual abstract suspend fun listInactiveAgents(): List<TAgent>
-    public actual abstract suspend fun listFinishedAgents(): List<TAgent>
-    public actual abstract suspend fun closeAll()
 
     @Suppress("ACTUAL_ANNOTATIONS_NOT_MATCH_EXPECT")
     public actual companion object {

agents/agents-core/src/appleMain/kotlin/ai/koog/agents/core/agent/AIAgent.kt

@@ -3,7 +3,9 @@
 package ai.koog.agents.core.agent
 
 import ai.koog.agents.core.agent.config.AIAgentConfig
+import ai.koog.agents.core.agent.context.AIAgentContext
 import ai.koog.agents.core.agent.entity.AIAgentGraphStrategy
+import ai.koog.agents.core.agent.session.AIAgentRunSession
 import ai.koog.agents.core.tools.ToolRegistry
 import ai.koog.prompt.executor.model.PromptExecutor
 import ai.koog.prompt.llm.LLModel
@@ -17,11 +19,9 @@ public actual abstract class AIAgent<Input, Output> : Closeable {
     public actual abstract val id: String
     public actual abstract val agentConfig: AIAgentConfig
 
-    public actual abstract suspend fun getState(): AIAgentState<Output>
+    public actual abstract suspend fun run(agentInput: Input, sessionId: String?): Output
 
-    public actual open suspend fun result(): Output = AIAgentHelper.result(this)
-
-    public actual abstract suspend fun run(agentInput: Input): Output
+    public actual abstract fun createSession(sessionId: String?): AIAgentRunSession<Input, Output, out AIAgentContext>
 
     public actual companion object {
         @OptIn(markerClass = [ExperimentalUuidApi::class])

agents/agents-core/src/appleMain/kotlin/ai/koog/agents/core/agent/AIAgentService.kt

@@ -33,11 +33,6 @@ public actual abstract class AIAgentService<Input, Output, TAgent : AIAgent<Inpu
     public actual abstract suspend fun removeAgent(agent: TAgent): Boolean
     public actual abstract suspend fun removeAgentWithId(id: String): Boolean
     public actual abstract suspend fun agentById(id: String): TAgent?
-    public actual abstract suspend fun listAllAgents(): List<TAgent>
-    public actual abstract suspend fun listActiveAgents(): List<TAgent>
-    public actual abstract suspend fun listInactiveAgents(): List<TAgent>
-    public actual abstract suspend fun listFinishedAgents(): List<TAgent>
-    public actual abstract suspend fun closeAll()
 
     @Suppress("ACTUAL_ANNOTATIONS_NOT_MATCH_EXPECT")
     public actual companion object {

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

@@ -2,11 +2,11 @@
 
 package ai.koog.agents.core.agent
 
-import ai.koog.agents.core.agent.AIAgentState.Finished
-import ai.koog.agents.core.agent.AIAgentState.Running
 import ai.koog.agents.core.agent.GraphAIAgent.FeatureContext
 import ai.koog.agents.core.agent.config.AIAgentConfig
+import ai.koog.agents.core.agent.context.AIAgentContext
 import ai.koog.agents.core.agent.entity.AIAgentGraphStrategy
+import ai.koog.agents.core.agent.session.AIAgentRunSession
 import ai.koog.agents.core.tools.ToolRegistry
 import ai.koog.prompt.executor.model.PromptExecutor
 import ai.koog.prompt.llm.LLModel
@@ -32,31 +32,23 @@ public expect abstract class AIAgent<Input, Output> constructor() : Closeable {
     public abstract val agentConfig: AIAgentConfig
 
     /**
-     * Retrieves the current state of the AI agent during its lifecycle.
-     *
-     * This method provides the current [AIAgentState] of the agent, which can
-     * be one of the defined states: [AIAgentState.NotStarted], [AIAgentState.Running], [AIAgentState.Finished], or [AIAgentState.Failed].
+     * Executes the AI agent with the given input and retrieves the resulting output.
      *
-     * @return The current state of the AI agent.
+     * @param agentInput The input for the agent.
+     * @return The output produced by the agent.
      */
-    public abstract suspend fun getState(): AIAgentState<Output>
+    public abstract suspend fun run(agentInput: Input, sessionId: String? = null): Output
 
     /**
-     * Retrieves the result of the operation if the current state is [AIAgentState.Finished].
-     * Throws an `IllegalStateException` if the operation is not in a finished state.
+     * Creates a new session for executing the agent with the given input.
      *
-     * @return The result of type `Output` when the operation is completed successfully.
-     * @throws IllegalStateException if the operation's state is not [AIAgentState.Finished].
-     */
-    public open suspend fun result(): Output
-
-    /**
-     * Executes the AI agent with the given input and retrieves the resulting output.
+     * This method provides a way to get a session object that can be used to execute
+     * the agent independently. The session manages the complete execution lifecycle, including
+     * state tracking, pipeline coordination, and strategy execution.
      *
-     * @param agentInput The input for the agent.
-     * @return The output produced by the agent.
+     * @return A session instance that can be used to run the agent with specific input and context.
      */
-    public abstract suspend fun run(agentInput: Input): Output
+    public abstract fun createSession(sessionId: String? = null): AIAgentRunSession<Input, Output, out AIAgentContext>
 
     /**
      * The companion object for the AIAgent class, providing functionality to instantiate an AI agent
@@ -243,17 +235,3 @@ public expect abstract class AIAgent<Input, Output> constructor() : Closeable {
         ): AIAgent<Input, Output>
     }
 }
-
-/**
- * Checks whether the AI agent is currently in a running state.
- *
- * @return `true` if the AI agent's state is `Running`, otherwise `false`.
- */
-public suspend fun AIAgent<*, *>.isRunning(): Boolean = this.getState() is Running
-
-/**
- * Checks whether the AI agent has reached a finished state.
- *
- * @return true if the current state of the AI agent is of type `Finished`, false otherwise.
- */
-public suspend fun AIAgent<*, *>.isFinished(): Boolean = this.getState() is Finished

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

@@ -0,0 +1,115 @@
+@file:Suppress("EXPECT_ACTUAL_CLASSIFIERS_ARE_IN_BETA_WARNING")
+@file:OptIn(InternalAgentsApi::class)
+
+package ai.koog.agents.core.agent
+
+import ai.koog.agents.core.agent.context.AIAgentContext
+import ai.koog.agents.core.agent.entity.AIAgentStrategy
+import ai.koog.agents.core.agent.session.AIAgentRunSession
+import ai.koog.agents.core.annotation.InternalAgentsApi
+import ai.koog.agents.core.feature.pipeline.AIAgentPipeline
+import io.github.oshai.kotlinlogging.KLogger
+import kotlin.uuid.ExperimentalUuidApi
+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.
+ *
+ * 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
+ * thread-safe operations via a mutex.
+ *
+ * @param Input the type of the input accepted by the agent.
+ * @param Output the type of the output produced by the agent.
+ * @param TContext the type of the context used during the agent's execution, extending [AIAgentContext].
+ * @property logger the logger used for logging execution details and errors.
+ * @param id the unique identifier for the agent. Random UUID will be generated if set to null.
+ */
+public abstract class AIAgentBase<Input, Output, TContext : AIAgentContext> constructor(
+    logger: KLogger,
+    id: String? = null,
+) : AIAgent<Input, Output>() {
+    /**
+     * Logger instance used for logging messages and events specific to this agent.
+     */
+    internal open val logger: KLogger = logger
+
+    @OptIn(ExperimentalUuidApi::class)
+    final override val id: String by lazy { id ?: Uuid.random().toString() }
+
+    /**
+     * The execution strategy defining how the agent processes input and produces output.
+     */
+    public abstract val strategy: AIAgentStrategy<Input, Output, TContext>
+
+    /**
+     * Represents the pipeline used by the AI agent for processing tasks or data.
+     *
+     * This abstract property defines the structure or sequence of operations
+     * within the AI agent's pipeline. It serves as the core mechanism for
+     * executing workflows, handling inputs, and generating outputs in the
+     * AI agent's functionality.
+     */
+    @InternalAgentsApi
+    public abstract val pipeline: AIAgentPipeline
+
+    /**
+     * Executes the agent's main functionality, coordinating with various components
+     * such as pipelines and strategies. Ensures the agent is run in a thread-safe
+     * manner using locking mechanisms.
+     *
+     * @param agentInput The input required to execute the agent's strategy.
+     *                   This includes any data necessary for processing.
+     * @return The output generated by the agent's execution, produced as a
+     *         result of applying the strategy to the provided input.
+     * @throws IllegalStateException if the agent was already started.
+     * @throws Throwable if any exception occurs during the execution process.
+     */
+    @OptIn(ExperimentalUuidApi::class)
+    override suspend fun run(agentInput: Input, sessionId: String?): Output {
+        val runId = sessionId ?: Uuid.random().toString()
+        val session = AIAgentRunSessionImpl(runId, logger, this, strategy, pipeline, ::prepareContext)
+        return session.run(agentInput)
+    }
+
+    /**
+     * 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
+     * that any associated features or stream providers are properly closed.
+     *
+     * Overrides the `close` method to implement agent-specific shutdown logic.
+     */
+    override suspend fun close() {
+        // TODO: Delete Closeable interface from [AIAgent] declaration.
+    }
+
+    /**
+     * Creates a new [AIAgentRunSession] for executing the agent with the given input.
+     *
+     * This method provides a way to obtain a session object that can be used to execute
+     * the agent independently. The session manages the complete execution lifecycle including
+     * state tracking, pipeline coordination, and strategy execution.
+     *
+     * @return A new [AIAgentRunSession] instance configured with this agent's logger, strategy,
+     *         and identifier. The session can be used to run the agent with specific input and context.
+     */
+    @OptIn(ExperimentalUuidApi::class)
+    override fun createSession(sessionId: String?): AIAgentRunSession<Input, Output, TContext> {
+        val runId = sessionId ?: Uuid.random().toString()
+        return AIAgentRunSessionImpl(runId, logger, this, strategy, pipeline, ::prepareContext)
+    }
+
+    /**
+     * Prepares and initializes the agent context required to handle the given input and run ID.
+     *
+     * @param agentInput the input provided to the agent for processing.
+     * @param runId a unique identifier representing the current execution or operation run.
+     * @param eventId a unique identifier for agent-related events.
+     * @return the initialized context specific to the agent setup for the provided input and run ID.
+     */
+    public abstract suspend fun prepareContext(agentInput: Input, runId: String, eventId: String): TContext
+}

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

@@ -1,6 +1,5 @@
 package ai.koog.agents.core.agent
 
-import ai.koog.agents.core.agent.AIAgentState.Finished
 import ai.koog.agents.core.agent.GraphAIAgent.FeatureContext
 import ai.koog.agents.core.agent.config.AIAgentConfig
 import ai.koog.agents.core.agent.entity.AIAgentGraphStrategy
@@ -17,10 +16,6 @@ import kotlin.uuid.ExperimentalUuidApi
 
 @PublishedApi
 internal object AIAgentHelper {
-    suspend fun <Output> result(agent: AIAgent<*, Output>): Output = when (val state = agent.getState()) {
-        is Finished<Output> -> state.result
-        else -> throw IllegalStateException("Output is not ready, agent's state is: $state")
-    }
 
     /**
      * Creates and returns a new instance of the `Builder` class to configure and construct an AI agent.