Fix AIAgentFunctionalContext delegations when doing copy()

Ololoshechkin · Ololoshechkin · commit d7ea7a43eb96 · 2026-01-08T19:10:13.000+01:00
diff --git a/agents/agents-core/src/androidMain/kotlin/ai/koog/agents/core/agent/context/AIAgentFunctionalContext.kt b/agents/agents-core/src/androidMain/kotlin/ai/koog/agents/core/agent/context/AIAgentFunctionalContext.kt
@@ -17,18 +17,6 @@ import ai.koog.prompt.structure.StructuredResponse
 
 @Suppress("MissingKDocForPublicAPI")
 public actual class AIAgentFunctionalContext internal actual constructor(
-    environment: AIAgentEnvironment,
-    agentId: String,
-    runId: String,
-    agentInput: Any?,
-    config: AIAgentConfig,
-    llm: AIAgentLLMContext,
-    stateManager: AIAgentStateManager,
-    storage: AIAgentStorage,
-    strategyName: String,
-    pipeline: AIAgentFunctionalPipeline,
-    executionInfo: AgentExecutionInfo,
-    parentContext: AIAgentContext?,
     @PublishedApi internal actual val delegate: AIAgentFunctionalContextImpl
 ) : AIAgentFunctionalContextAPI by delegate {
 
@@ -46,17 +34,6 @@ public actual class AIAgentFunctionalContext internal actual constructor(
         executionInfo: AgentExecutionInfo,
         parentContext: AIAgentContext?
     ) : this(
-        environment = environment,
-        agentId = agentId,
-        runId = runId,
-        agentInput = agentInput,
-        config = config,
-        llm = llm,
-        stateManager = stateManager,
-        storage = storage,
-        strategyName = strategyName,
-        pipeline = pipeline,
-        executionInfo = executionInfo,
         delegate = AIAgentFunctionalContextImpl(
             environment = environment,
             agentId = agentId,
@@ -87,7 +64,7 @@ public actual class AIAgentFunctionalContext internal actual constructor(
         llmParams: LLMParams?,
         runMode: ToolCalls,
         assistantResponseRepeatMax: Int?,
-    ): Output = delegate.subtask(
+    ): Output = delegate.subtaskImpl(
         taskDescription,
         input,
         tools,
@@ -111,18 +88,20 @@ public actual class AIAgentFunctionalContext internal actual constructor(
         executionInfo: AgentExecutionInfo,
         parentRootContext: AIAgentContext?
     ): AIAgentFunctionalContext = AIAgentFunctionalContext(
-        environment = environment,
-        agentId = agentId,
-        runId = runId,
-        agentInput = agentInput,
-        config = config,
-        llm = llm,
-        stateManager = stateManager,
-        storage = storage,
-        strategyName = strategyName,
-        pipeline = pipeline,
-        executionInfo = executionInfo,
-        parentContext = parentRootContext,
-        delegate = delegate,
+        delegate = AIAgentFunctionalContextImpl(
+            environment = environment,
+            agentId = agentId,
+            pipeline = pipeline,
+            runId = runId,
+            agentInput = agentInput,
+            config = config,
+            llm = llm,
+            stateManager = stateManager,
+            storage = storage,
+            strategyName = strategyName,
+            parentContext = parentContext,
+            executionInfo = executionInfo,
+            storeMap = delegate.storeMap
+        )
     )
 }
diff --git a/agents/agents-core/src/appleMain/kotlin/ai/koog/agents/core/agent/context/AIAgentFunctionalContext.kt b/agents/agents-core/src/appleMain/kotlin/ai/koog/agents/core/agent/context/AIAgentFunctionalContext.kt
@@ -17,18 +17,6 @@ import ai.koog.prompt.structure.StructuredResponse
 
 @Suppress("MissingKDocForPublicAPI")
 public actual class AIAgentFunctionalContext internal actual constructor(
-    environment: AIAgentEnvironment,
-    agentId: String,
-    runId: String,
-    agentInput: Any?,
-    config: AIAgentConfig,
-    llm: AIAgentLLMContext,
-    stateManager: AIAgentStateManager,
-    storage: AIAgentStorage,
-    strategyName: String,
-    pipeline: AIAgentFunctionalPipeline,
-    executionInfo: AgentExecutionInfo,
-    parentContext: AIAgentContext?,
     @PublishedApi internal actual val delegate: AIAgentFunctionalContextImpl
 ) : AIAgentFunctionalContextAPI by delegate {
 
@@ -46,17 +34,6 @@ public actual class AIAgentFunctionalContext internal actual constructor(
         executionInfo: AgentExecutionInfo,
         parentContext: AIAgentContext?
     ) : this(
-        environment = environment,
-        agentId = agentId,
-        runId = runId,
-        agentInput = agentInput,
-        config = config,
-        llm = llm,
-        stateManager = stateManager,
-        storage = storage,
-        strategyName = strategyName,
-        pipeline = pipeline,
-        executionInfo = executionInfo,
         delegate = AIAgentFunctionalContextImpl(
             environment = environment,
             agentId = agentId,
@@ -87,7 +64,7 @@ public actual class AIAgentFunctionalContext internal actual constructor(
         llmParams: LLMParams?,
         runMode: ToolCalls,
         assistantResponseRepeatMax: Int?,
-    ): Output = delegate.subtask(
+    ): Output = delegate.subtaskImpl(
         taskDescription,
         input,
         tools,
@@ -111,18 +88,20 @@ public actual class AIAgentFunctionalContext internal actual constructor(
         executionInfo: AgentExecutionInfo,
         parentRootContext: AIAgentContext?
     ): AIAgentFunctionalContext = AIAgentFunctionalContext(
-        environment = environment,
-        agentId = agentId,
-        runId = runId,
-        agentInput = agentInput,
-        config = config,
-        llm = llm,
-        stateManager = stateManager,
-        storage = storage,
-        strategyName = strategyName,
-        pipeline = pipeline,
-        executionInfo = executionInfo,
-        parentContext = parentRootContext,
-        delegate = delegate,
+        delegate = AIAgentFunctionalContextImpl(
+            environment = environment,
+            agentId = agentId,
+            pipeline = pipeline,
+            runId = runId,
+            agentInput = agentInput,
+            config = config,
+            llm = llm,
+            stateManager = stateManager,
+            storage = storage,
+            strategyName = strategyName,
+            parentContext = parentContext,
+            executionInfo = executionInfo,
+            storeMap = delegate.storeMap
+        )
     )
 }
diff --git a/agents/agents-core/src/commonMain/kotlin/ai/koog/agents/core/agent/AIAgentFunctionalStrategy.kt b/agents/agents-core/src/commonMain/kotlin/ai/koog/agents/core/agent/AIAgentFunctionalStrategy.kt
@@ -44,5 +44,7 @@ public fun <Input, Output> functionalStrategy(
     func: suspend AIAgentFunctionalContext.(input: Input) -> Output
 ): AIAgentFunctionalStrategy<Input, Output> = object : AIAgentFunctionalStrategy<Input, Output> {
     override val name: String = name
-    override suspend fun execute(context: AIAgentFunctionalContext, input: Input): Output = context.func(input)
+    override suspend fun execute(context: AIAgentFunctionalContext, input: Input): Output {
+        return context.func(input)
+    }
 }
diff --git a/agents/agents-core/src/commonMain/kotlin/ai/koog/agents/core/agent/context/AIAgentFunctionalContext.kt b/agents/agents-core/src/commonMain/kotlin/ai/koog/agents/core/agent/context/AIAgentFunctionalContext.kt
@@ -51,18 +51,6 @@ import kotlin.reflect.KClass
 @OptIn(InternalAgentsApi::class)
 @Suppress("UNCHECKED_CAST", "MissingKDocForPublicAPI")
 public expect class AIAgentFunctionalContext internal constructor(
-    environment: AIAgentEnvironment,
-    agentId: String,
-    runId: String,
-    agentInput: Any?,
-    config: AIAgentConfig,
-    llm: AIAgentLLMContext,
-    stateManager: AIAgentStateManager,
-    storage: AIAgentStorage,
-    strategyName: String,
-    pipeline: AIAgentFunctionalPipeline,
-    executionInfo: AgentExecutionInfo,
-    parentContext: AIAgentContext? = null,
     delegate: AIAgentFunctionalContextImpl
 ) : AIAgentFunctionalContextAPI {
     public constructor(
diff --git a/agents/agents-core/src/commonMain/kotlin/ai/koog/agents/core/agent/context/AIAgentFunctionalContextImpl.kt b/agents/agents-core/src/commonMain/kotlin/ai/koog/agents/core/agent/context/AIAgentFunctionalContextImpl.kt
@@ -53,11 +53,10 @@ internal class AIAgentFunctionalContextImpl(
     override val strategyName: String,
     override val pipeline: AIAgentFunctionalPipeline,
     override var executionInfo: AgentExecutionInfo,
+    internal val storeMap: MutableMap<AIAgentStorageKey<*>, Any> = mutableMapOf(),
     override val parentContext: AIAgentContext? = null
 ) : AIAgentFunctionalContextAPI {
 
-    private val storeMap: MutableMap<AIAgentStorageKey<*>, Any> = mutableMapOf()
-
     override fun store(key: AIAgentStorageKey<*>, value: Any) {
         storeMap[key] = value
     }
diff --git a/agents/agents-core/src/jsMain/kotlin/ai/koog/agents/core/agent/context/AIAgentFunctionalContext.kt b/agents/agents-core/src/jsMain/kotlin/ai/koog/agents/core/agent/context/AIAgentFunctionalContext.kt
@@ -17,20 +17,9 @@ import ai.koog.prompt.structure.StructuredResponse
 
 @Suppress("MissingKDocForPublicAPI")
 public actual class AIAgentFunctionalContext internal actual constructor(
-    environment: AIAgentEnvironment,
-    agentId: String,
-    runId: String,
-    agentInput: Any?,
-    config: AIAgentConfig,
-    llm: AIAgentLLMContext,
-    stateManager: AIAgentStateManager,
-    storage: AIAgentStorage,
-    strategyName: String,
-    pipeline: AIAgentFunctionalPipeline,
-    executionInfo: AgentExecutionInfo,
-    parentContext: AIAgentContext?,
     @PublishedApi internal actual val delegate: AIAgentFunctionalContextImpl
 ) : AIAgentFunctionalContextAPI by delegate {
+
     public actual constructor(
         environment: AIAgentEnvironment,
         agentId: String,
@@ -45,17 +34,6 @@ public actual class AIAgentFunctionalContext internal actual constructor(
         executionInfo: AgentExecutionInfo,
         parentContext: AIAgentContext?
     ) : this(
-        environment = environment,
-        agentId = agentId,
-        runId = runId,
-        agentInput = agentInput,
-        config = config,
-        llm = llm,
-        stateManager = stateManager,
-        storage = storage,
-        strategyName = strategyName,
-        pipeline = pipeline,
-        executionInfo = executionInfo,
         delegate = AIAgentFunctionalContextImpl(
             environment = environment,
             agentId = agentId,
@@ -86,7 +64,7 @@ public actual class AIAgentFunctionalContext internal actual constructor(
         llmParams: LLMParams?,
         runMode: ToolCalls,
         assistantResponseRepeatMax: Int?,
-    ): Output = delegate.subtask(
+    ): Output = delegate.subtaskImpl(
         taskDescription,
         input,
         tools,
@@ -110,18 +88,20 @@ public actual class AIAgentFunctionalContext internal actual constructor(
         executionInfo: AgentExecutionInfo,
         parentRootContext: AIAgentContext?
     ): AIAgentFunctionalContext = AIAgentFunctionalContext(
-        environment = environment,
-        agentId = agentId,
-        runId = runId,
-        agentInput = agentInput,
-        config = config,
-        llm = llm,
-        stateManager = stateManager,
-        storage = storage,
-        strategyName = strategyName,
-        pipeline = pipeline,
-        executionInfo = executionInfo,
-        parentContext = parentRootContext,
-        delegate = delegate,
+        delegate = AIAgentFunctionalContextImpl(
+            environment = environment,
+            agentId = agentId,
+            pipeline = pipeline,
+            runId = runId,
+            agentInput = agentInput,
+            config = config,
+            llm = llm,
+            stateManager = stateManager,
+            storage = storage,
+            strategyName = strategyName,
+            parentContext = parentContext,
+            executionInfo = executionInfo,
+            storeMap = delegate.storeMap
+        )
     )
 }
diff --git a/agents/agents-core/src/jvmMain/kotlin/ai/koog/agents/core/agent/context/AIAgentFunctionalContext.kt b/agents/agents-core/src/jvmMain/kotlin/ai/koog/agents/core/agent/context/AIAgentFunctionalContext.kt
@@ -38,20 +38,9 @@ import kotlin.reflect.KClass
 
 @Suppress("MissingKDocForPublicAPI")
 public actual class AIAgentFunctionalContext internal actual constructor(
-    environment: AIAgentEnvironment,
-    agentId: String,
-    runId: String,
-    agentInput: Any?,
-    config: AIAgentConfig,
-    llm: AIAgentLLMContext,
-    stateManager: AIAgentStateManager,
-    storage: AIAgentStorage,
-    strategyName: String,
-    pipeline: AIAgentFunctionalPipeline,
-    executionInfo: AgentExecutionInfo,
-    parentContext: AIAgentContext?,
     @PublishedApi internal actual val delegate: AIAgentFunctionalContextImpl
 ) : AIAgentFunctionalContextAPI by delegate {
+    @JvmOverloads
     public actual constructor(
         environment: AIAgentEnvironment,
         agentId: String,
@@ -66,17 +55,6 @@ public actual class AIAgentFunctionalContext internal actual constructor(
         executionInfo: AgentExecutionInfo,
         parentContext: AIAgentContext?
     ) : this(
-        environment = environment,
-        agentId = agentId,
-        runId = runId,
-        agentInput = agentInput,
-        config = config,
-        llm = llm,
-        stateManager = stateManager,
-        storage = storage,
-        strategyName = strategyName,
-        pipeline = pipeline,
-        executionInfo = executionInfo,
         delegate = AIAgentFunctionalContextImpl(
             environment = environment,
             agentId = agentId,
@@ -93,12 +71,14 @@ public actual class AIAgentFunctionalContext internal actual constructor(
         )
     )
 
+    @JvmOverloads
     public actual suspend inline fun <reified T> requestLLMStructured(
         message: String,
         examples: List<T>,
         fixingParser: StructureFixingParser?
     ): Result<StructuredResponse<T>> = delegate.requestLLMStructured(message, examples, fixingParser)
 
+    @JvmOverloads
     public actual suspend inline fun <Input, reified Output> subtask(
         taskDescription: String,
         input: Input,
@@ -132,19 +112,21 @@ public actual class AIAgentFunctionalContext internal actual constructor(
         executionInfo: AgentExecutionInfo,
         parentRootContext: AIAgentContext?
     ): AIAgentFunctionalContext = AIAgentFunctionalContext(
-        environment = environment,
-        agentId = agentId,
-        runId = runId,
-        agentInput = agentInput,
-        config = config,
-        llm = llm,
-        stateManager = stateManager,
-        storage = storage,
-        strategyName = strategyName,
-        pipeline = pipeline,
-        executionInfo = executionInfo,
-        parentContext = parentRootContext,
-        delegate = delegate,
+        delegate = AIAgentFunctionalContextImpl(
+            environment = environment,
+            agentId = agentId,
+            pipeline = pipeline,
+            runId = runId,
+            agentInput = agentInput,
+            config = config,
+            llm = llm,
+            stateManager = stateManager,
+            storage = storage,
+            strategyName = strategyName,
+            parentContext = parentContext,
+            executionInfo = executionInfo,
+            storeMap = delegate.storeMap
+        )
     )
 
     /**
diff --git a/agents/agents-core/src/wasmJsMain/kotlin/ai/koog/agents/core/agent/context/AIAgentFunctionalContext.kt b/agents/agents-core/src/wasmJsMain/kotlin/ai/koog/agents/core/agent/context/AIAgentFunctionalContext.kt
diff --git a/agents/agents-tools/src/jvmTest/java/ai/koog/agents/tools/JavaMethodToolsTest.java b/agents/agents-tools/src/jvmTest/java/ai/koog/agents/tools/JavaMethodToolsTest.java

Original file line number	Diff line number	Diff line change
`@@ -44,5 +44,7 @@ public fun <Input, Output> functionalStrategy(`
`44`	`44`	`func: suspend AIAgentFunctionalContext.(input: Input) -> Output`
`45`	`45`	`): AIAgentFunctionalStrategy<Input, Output> = object : AIAgentFunctionalStrategy<Input, Output> {`
`46`	`46`	`override val name: String = name`
`47`		`- override suspend fun execute(context: AIAgentFunctionalContext, input: Input): Output = context.func(input)`
	`47`	`+ override suspend fun execute(context: AIAgentFunctionalContext, input: Input): Output {`
	`48`	`+ return context.func(input)`
	`49`	`+ }`
`48`	`50`	`}`