diff --git a/core/src/commonJvmAndroidTest/kotlin/com/google/adk/kt/skills/NewFileSystemSourceTest.kt b/core/src/commonJvmAndroidTest/kotlin/com/google/adk/kt/skills/NewFileSystemSourceTest.kt index 71b87302..2a71ec93 100644 --- a/core/src/commonJvmAndroidTest/kotlin/com/google/adk/kt/skills/NewFileSystemSourceTest.kt +++ b/core/src/commonJvmAndroidTest/kotlin/com/google/adk/kt/skills/NewFileSystemSourceTest.kt @@ -203,6 +203,35 @@ class NewFileSystemSourceTest { assertThat(result.getOrThrow().name).isEqualTo("skill1") } + @Test + fun loadFrontmatter_adkAdditionalToolsMetadata_returnsFrontmatterWithMetadata() = runTest { + writeRawSkillMd( + dirName = "test-skill", + content = + """ + --- + name: test-skill + description: Description + metadata: + adk_additional_tools: + - extra_tool_a + - extra_tool_b + owner: adk + --- + Instructions. + """ + .trimIndent(), + ) + + val result = source.loadFrontmatter("test-skill") + + assertThat(result.isSuccess).isTrue() + val frontmatter = result.getOrThrow() + assertThat(frontmatter.metadata["adk_additional_tools"]) + .isEqualTo(listOf("extra_tool_a", "extra_tool_b")) + assertThat(frontmatter.metadata["owner"]).isEqualTo("adk") + } + @Test fun loadFrontmatter_nonexistentSkill_returnsFailure() = runTest { val result = source.loadFrontmatter("nonexistent") diff --git a/core/src/commonMain/kotlin/com/google/adk/kt/agents/LlmAgentTurn.kt b/core/src/commonMain/kotlin/com/google/adk/kt/agents/LlmAgentTurn.kt index 6051cb97..331e4516 100644 --- a/core/src/commonMain/kotlin/com/google/adk/kt/agents/LlmAgentTurn.kt +++ b/core/src/commonMain/kotlin/com/google/adk/kt/agents/LlmAgentTurn.kt @@ -23,6 +23,7 @@ import com.google.adk.kt.callbacks.runOnModelErrorCallbacksPipeline import com.google.adk.kt.events.Event import com.google.adk.kt.events.getLongRunningFunctionIds import com.google.adk.kt.ids.Uuid +import com.google.adk.kt.logging.LoggerFactory import com.google.adk.kt.models.LlmRequest import com.google.adk.kt.models.LlmResponse import com.google.adk.kt.models.toTracePayload @@ -53,6 +54,12 @@ import kotlinx.coroutines.flow.emitAll import kotlinx.coroutines.flow.flow import kotlinx.coroutines.flow.flowOn +private fun Iterable.associateByNameFirstWins(): Map = buildMap { + for (tool in this@associateByNameFirstWins) { + if (tool.name !in this) put(tool.name, tool) + } +} + /** * Encapsulates the logic for a single turn of an [LlmAgent]. * @@ -70,6 +77,7 @@ internal class LlmAgentTurn( private val requestProcessors: List, private val responseProcessors: List, ) { + private val logger = LoggerFactory.getLogger(LlmAgentTurn::class) /** * Executes the turn logic and returns a flow of events. @@ -151,16 +159,39 @@ internal class LlmAgentTurn( } val toolContext = ToolContext(invocationContext = context) + val registeredToolsByName = + processedRequest.toolsDict.associateByNameFirstWins().toMutableMap() + + suspend fun registerTool(req: LlmRequest, tool: BaseTool): LlmRequest { + val existing = registeredToolsByName[tool.name] + if (existing != null) { + if (existing !== tool) { + logger.warn { + "Duplicate tool name '${tool.name}'; keeping the first registered tool and skipping " + + "the later one." + } + } + return req + } + registeredToolsByName[tool.name] = tool + return tool.processLlmRequest(toolContext, req) + } + val reqAfterCodeExecutor = canonicalDirectTools.fold(processedRequest) { req, tool -> - tool.processLlmRequest(toolContext, req) + registerTool(req, tool) } val finalRequest = agent.toolsets.fold(reqAfterCodeExecutor) { req, toolset -> val setReq = toolset.processLlmRequest(toolContext, req) + for (tool in setReq.toolsDict) { + if (tool.name !in registeredToolsByName) { + registeredToolsByName[tool.name] = tool + } + } toolset.getTools(context.toReadonlyContext()).fold(setReq) { r, tool -> - tool.processLlmRequest(toolContext, r) + registerTool(r, tool) } } return RequestOrResponse.Request(finalRequest) @@ -168,8 +199,11 @@ internal class LlmAgentTurn( private suspend fun getToolMap(request: LlmRequest?): Map { val readonlyCtx = context.toReadonlyContext() - val allTools = canonicalDirectTools + agent.toolsets.flatMap { it.getTools(readonlyCtx) } - return (allTools + request?.toolsDict.orEmpty()).associateBy { it.name } + val toolsInRegistrationOrder = + request?.toolsDict.orEmpty() + + canonicalDirectTools + + agent.toolsets.flatMap { it.getTools(readonlyCtx) } + return toolsInRegistrationOrder.associateByNameFirstWins() } /** diff --git a/core/src/commonMain/kotlin/com/google/adk/kt/annotations/ExperimentalSkillTools.kt b/core/src/commonMain/kotlin/com/google/adk/kt/annotations/ExperimentalSkillTools.kt new file mode 100644 index 00000000..29a13ddb --- /dev/null +++ b/core/src/commonMain/kotlin/com/google/adk/kt/annotations/ExperimentalSkillTools.kt @@ -0,0 +1,30 @@ +/* + * Copyright 2026 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.google.adk.kt.annotations + +/** + * Marks APIs for exposing skill-provided additional tools as experimental. + * + * These APIs may change or be removed as skill behavior is aligned across ADK implementations. + */ +@MustBeDocumented +@Retention(AnnotationRetention.BINARY) +@RequiresOptIn( + level = RequiresOptIn.Level.ERROR, + message = "Skill-provided additional tools are experimental and may change or be removed.", +) +annotation class ExperimentalSkillTools diff --git a/core/src/commonMain/kotlin/com/google/adk/kt/tools/SkillToolset.kt b/core/src/commonMain/kotlin/com/google/adk/kt/tools/SkillToolset.kt index c91053ba..56f87c38 100644 --- a/core/src/commonMain/kotlin/com/google/adk/kt/tools/SkillToolset.kt +++ b/core/src/commonMain/kotlin/com/google/adk/kt/tools/SkillToolset.kt @@ -17,6 +17,7 @@ package com.google.adk.kt.tools import com.google.adk.kt.agents.ReadonlyContext +import com.google.adk.kt.annotations.ExperimentalSkillTools import com.google.adk.kt.logging.Logger import com.google.adk.kt.logging.LoggerFactory import com.google.adk.kt.models.LlmRequest @@ -27,6 +28,16 @@ import com.google.adk.kt.types.FunctionDeclaration import com.google.adk.kt.types.Schema as GenaiSchema import com.google.adk.kt.types.Type +private inline fun Iterable.forEachDistinctReference(action: (T) -> Unit) { + val seen = mutableListOf() + for (value in this) { + if (seen.none { it === value }) { + seen += value + action(value) + } + } +} + /** Builds the standard `{error}` response map used by all skill tools. */ private fun errorResponse(message: String): Map = mapOf(SkillToolset.KEY_ERROR to message) @@ -77,7 +88,7 @@ internal class ListSkillsTool(private val toolset: SkillToolset) : } } -private fun Frontmatter.frontmatterDsl() = +private fun Frontmatter.frontmatterDsl(): Map = mapOf( "name" to name, "description" to description, @@ -87,6 +98,21 @@ private fun Frontmatter.frontmatterDsl() = "metadata" to metadata, ) +private fun Frontmatter.additionalToolNames(): List { + val value = + metadata[SkillToolset.ADK_ADDITIONAL_TOOLS_METADATA_KEY] as? List<*> + ?: return emptyList() + return value.filterIsInstance() +} + +private fun additionalToolNamesBySkill(value: Any?): Map> = buildMap { + for ((skillName, toolNames) in (value as? Map<*, *>).orEmpty()) { + val validSkillName = skillName as? String ?: continue + val validToolNames = (toolNames as? List<*>)?.filterIsInstance() ?: continue + put(validSkillName, validToolNames) + } +} + /** BaseTool responsible for loading the instructions for a specific skill. */ internal class LoadSkillTool(private val toolset: SkillToolset) : BaseTool( @@ -127,12 +153,70 @@ internal class LoadSkillTool(private val toolset: SkillToolset) : return e.toSkillSourceErrorResponse(logger) } + recordSkillActivation(context, skillName, frontmatter.additionalToolNames()) + return mapOf( SkillToolset.PARAM_SKILL_NAME to skillName, SkillToolset.KEY_INSTRUCTIONS to instructions, SkillToolset.KEY_FRONTMATTER to frontmatter.frontmatterDsl(), ) } + + /** + * Records [skillName] and its [additionalToolNames] in the per-agent activation state. + * + * The activation state is split across two stores in adk-kotlin: the pending + * [ToolContext.actions.stateDelta] (not yet applied to the session) and the already-applied + * [ToolContext.context.state]. To preserve other skills that have been activated earlier in the + * same invocation, both stores are merged before writing the updated values back into + * `stateDelta`. Activated skill names and their additional-tool snapshots use separate keys so + * subsequent model steps can expose the tools without loading the skill frontmatter again. A + * later load of the same skill replaces that skill's tool-name snapshot. + * + * Note: when the LLM emits multiple `load_skill` calls in the same step, ADK runs them in + * parallel and the per-call deltas are merged with last-write-wins semantics + * ([EventActions.mergeWith]). Only the activation state recorded by the final merge survives the + * step. This matches adk-python's behavior and is a known limitation; downstream tool exposure + * still works for serial `load_skill` calls across steps. + */ + private fun recordSkillActivation( + context: ToolContext, + skillName: String, + additionalToolNames: List, + ) { + appendUniqueStateValues( + context, + SkillToolset.activatedSkillStateKey(context.context.agentName), + listOf(skillName), + ) + replaceAdditionalToolNames(context, skillName, additionalToolNames) + } + + private fun appendUniqueStateValues( + context: ToolContext, + stateKey: String, + newValues: List, + ) { + val pending = (context.actions.stateDelta[stateKey] as? List<*>).orEmpty() + val applied = (context.context.state[stateKey] as? List<*>).orEmpty() + context.actions.stateDelta[stateKey] = + (applied + pending + newValues).filterIsInstance().distinct() + } + + private fun replaceAdditionalToolNames( + context: ToolContext, + skillName: String, + additionalToolNames: List, + ) { + val stateKey = SkillToolset.activatedSkillToolsStateKey(context.context.agentName) + val merged = + LinkedHashMap>().apply { + putAll(additionalToolNamesBySkill(context.context.state[stateKey])) + putAll(additionalToolNamesBySkill(context.actions.stateDelta[stateKey])) + put(skillName, additionalToolNames.toList()) + } + context.actions.stateDelta[stateKey] = merged + } } /** BaseTool responsible for loading resources (references/assets/scripts) from a specific skill. */ @@ -206,8 +290,37 @@ internal class LoadSkillResourceTool(private val toolset: SkillToolset) : } } -/** Toolset that manages and provides access to a collection of [Skill]s. */ -class SkillToolset(internal val source: SkillSource) : Toolset { +/** + * Toolset that manages and provides access to a collection of skills. + * + * The constructors accepting [additionalTools] or [additionalToolsets] transfer ownership of those + * instances to this toolset. [close] closes every directly provided tool and toolset once by + * reference. Callers must not provide the same owned resource through multiple containers. + * + * @param source Source from which skills are loaded. + * @param additionalTools Tools hidden until an activated skill names them in + * `metadata.adk_additional_tools`. + * @param additionalToolsets Toolsets that can supply hidden tools after activation. + */ +class SkillToolset +@ExperimentalSkillTools +constructor( + internal val source: SkillSource, + additionalTools: List, + additionalToolsets: List, +) : Toolset { + + @OptIn(ExperimentalSkillTools::class) + constructor(source: SkillSource) : this(source, emptyList(), emptyList()) + + /** + * Creates a skill toolset that owns [additionalTools] and closes them when [close] is called. + */ + @ExperimentalSkillTools + constructor( + source: SkillSource, + additionalTools: List, + ) : this(source, additionalTools, emptyList()) companion object { /** The name of the tool used to list available skills. */ @@ -235,14 +348,109 @@ class SkillToolset(internal val source: SkillSource) : Toolset { /** Message indicating that a loaded resource is a binary file. */ const val MSG_BINARY_FILE = "Binary file detected. Content not shown." + + /** Metadata key used by adk-python skills to declare tools activated by `load_skill`. */ + internal const val ADK_ADDITIONAL_TOOLS_METADATA_KEY = "adk_additional_tools" + + /** + * Prefix for the per-agent state key under which the SkillToolset records which skills the LLM + * has activated by calling `load_skill`. Mirrors adk-python's `_adk_activated_skill_` prefix. + */ + internal const val STATE_KEY_PREFIX_ACTIVATED_SKILL = "_adk_activated_skill_" + + /** Prefix for the per-agent state containing tools enabled by activated skills. */ + internal const val STATE_KEY_PREFIX_ACTIVATED_SKILL_TOOLS = + "_adk_activated_skill_tools_" + + /** Builds the activation-state key for the given agent name. */ + internal fun activatedSkillStateKey(agentName: String): String = + STATE_KEY_PREFIX_ACTIVATED_SKILL + agentName + + /** Builds the additional-tool activation-state key for the given agent name. */ + internal fun activatedSkillToolsStateKey(agentName: String): String = + STATE_KEY_PREFIX_ACTIVATED_SKILL_TOOLS + agentName } private val logger = LoggerFactory.getLogger(SkillToolset::class) + /** + * Additional tools that are hidden from the LLM until a skill declares them in its frontmatter's + * `adk_additional_tools` and that skill has been loaded via `load_skill`. Keyed by tool name; + * duplicates are dropped with a warning (last one wins). + */ + private val providedTools = additionalTools.toList() + + internal val providedToolsByName: Map = + providedTools.associateBy { it.name }.also { + providedTools.groupingBy { it.name }.eachCount().forEach { (name, count) -> + if (count > 1) { + logger.warn { "Duplicate additional tool name '$name'; last one wins." } + } + } + } + + /** + * Additional toolsets that can contribute tools after activation. This mirrors adk-python's + * ability to resolve additional tools from both tools and toolsets, while keeping the Kotlin API + * strongly typed. + */ + private val providedToolsets: List = additionalToolsets.toList() + private val tools: List = listOf(ListSkillsTool(this), LoadSkillTool(this), LoadSkillResourceTool(this)) - override suspend fun getTools(readonlyContext: ReadonlyContext?): List = tools + private val forbiddenToolNames: Set + get() = tools.mapTo(mutableSetOf()) { it.name } + + override suspend fun getTools(readonlyContext: ReadonlyContext?): List = + tools + resolveAdditionalToolsFromState(readonlyContext) + + /** + * Resolves the additional tools that should be exposed for this invocation, based on which skills + * have been activated via `load_skill`. + * + * Additional tool names are read from the per-skill snapshots in [ReadonlyContext.state] under + * the agent-specific key built by [activatedSkillToolsStateKey]. Tool names with no candidate + * match are skipped silently. Names that collide with the core skill tools are dropped with a + * warning. + */ + private suspend fun resolveAdditionalToolsFromState( + readonlyContext: ReadonlyContext?, + ): List { + if (readonlyContext == null) return emptyList() + if (providedToolsByName.isEmpty() && providedToolsets.isEmpty()) return emptyList() + + val additionalToolNames = + additionalToolNamesBySkill( + readonlyContext.state[activatedSkillToolsStateKey(readonlyContext.agentName)] + ) + .values + .flatten() + .distinct() + if (additionalToolNames.isEmpty()) return emptyList() + + val candidateTools = getCandidateTools(readonlyContext) + val resolved = mutableListOf() + for (toolName in additionalToolNames) { + if (toolName in forbiddenToolNames) { + logger.warn { + "Tool name collision: additional tool '$toolName' shadows a core skill tool; skipping." + } + continue + } + candidateTools[toolName]?.let { resolved += it } + } + return resolved + } + + private suspend fun getCandidateTools(readonlyContext: ReadonlyContext): Map = + providedToolsByName + + providedToolsets.flatMap { it.getTools(readonlyContext) }.associateBy { it.name } + + override fun close() { + (tools + providedTools).forEachDistinctReference { it.close() } + providedToolsets.forEachDistinctReference { it.close() } + } override suspend fun processLlmRequest( toolContext: ToolContext, diff --git a/core/src/commonTest/kotlin/com/google/adk/kt/agents/LlmAgentTurnTest.kt b/core/src/commonTest/kotlin/com/google/adk/kt/agents/LlmAgentTurnTest.kt index 03f6042b..8940a2f7 100644 --- a/core/src/commonTest/kotlin/com/google/adk/kt/agents/LlmAgentTurnTest.kt +++ b/core/src/commonTest/kotlin/com/google/adk/kt/agents/LlmAgentTurnTest.kt @@ -16,6 +16,7 @@ package com.google.adk.kt.agents +import com.google.adk.kt.models.LlmRequest import com.google.adk.kt.models.LlmResponse import com.google.adk.kt.runners.InMemoryRunner import com.google.adk.kt.testing.DummyModel @@ -27,10 +28,17 @@ import com.google.adk.kt.testing.modelTransferToAgentResponse import com.google.adk.kt.testing.simplifyEvents import com.google.adk.kt.testing.transferToAgentCallPart import com.google.adk.kt.testing.userMessage +import com.google.adk.kt.tools.BaseTool +import com.google.adk.kt.tools.ToolContext +import com.google.adk.kt.tools.Toolset import com.google.adk.kt.types.FunctionCall +import com.google.adk.kt.types.FunctionDeclaration import com.google.adk.kt.types.FunctionResponse import com.google.adk.kt.types.Part import com.google.adk.kt.types.Role +import com.google.adk.kt.types.Schema +import com.google.adk.kt.types.Tool +import com.google.adk.kt.types.Type import kotlinx.coroutines.flow.flow import kotlinx.coroutines.flow.toList import kotlinx.coroutines.test.runTest @@ -39,6 +47,97 @@ import org.junit.Test class LlmAgentTurnTest { + @Test + fun runAsync_toolsetHookAddsTool_preservesHookTool() = runTest { + val capturedRequests = mutableListOf() + val hookTool = RunTrackingDeclaredTool("runtime_tool") + val toolset = + object : Toolset { + override suspend fun processLlmRequest( + toolContext: ToolContext, + llmRequest: LlmRequest, + ): LlmRequest = llmRequest.appendTools(listOf(hookTool)) + + override suspend fun getTools(readonlyContext: ReadonlyContext?): List = + emptyList() + } + val model = + DummyModel("hook-tool-model") { request -> + capturedRequests += request + flow { + val hasFunctionResponse = + request.contents.lastOrNull()?.parts?.any { it.functionResponse != null } == true + if (hasFunctionResponse) { + emit(LlmResponse(content = modelMessage("done"))) + } else { + emit(modelFunctionCallResponse("runtime_tool", id = "runtime-tool-call")) + } + } + } + val agent = LlmAgent(name = "test-agent", model = model, toolsets = listOf(toolset)) + val runner = InMemoryRunner(agent) + + runner + .runAsync(userId = "user", sessionId = "session", newMessage = userMessage("run tool")) + .toList() + + assertEquals(2, capturedRequests.size) + capturedRequests.forEach { request -> + val declarationNames = + request.config.tools.orEmpty().flatMap { it.functionDeclarations.orEmpty() }.map { it.name } + assertEquals(listOf("runtime_tool"), declarationNames) + assertEquals(1, request.toolsDict.count { it.name == "runtime_tool" }) + } + assertEquals(1, hookTool.runCalls) + } + + @Test + fun runAsync_directToolNotInToolsDictCollidesWithToolset_keepsDirectTool() = runTest { + val capturedRequests = mutableListOf() + val directTool = ConfigOnlyRunTrackingDeclaredTool("shared_tool") + val additionalTool = RunTrackingDeclaredTool("shared_tool") + val toolset = + object : Toolset { + override suspend fun getTools(readonlyContext: ReadonlyContext?): List = + listOf(additionalTool) + } + val model = + DummyModel("direct-tool-collision-model") { request -> + capturedRequests += request + flow { + val hasFunctionResponse = + request.contents.lastOrNull()?.parts?.any { it.functionResponse != null } == true + if (hasFunctionResponse) { + emit(LlmResponse(content = modelMessage("done"))) + } else { + emit(modelFunctionCallResponse("shared_tool", id = "shared-tool-call")) + } + } + } + val agent = + LlmAgent( + name = "test-agent", + model = model, + tools = listOf(directTool), + toolsets = listOf(toolset), + ) + val runner = InMemoryRunner(agent) + + runner + .runAsync(userId = "user", sessionId = "session", newMessage = userMessage("run tool")) + .toList() + + assertEquals(2, capturedRequests.size) + capturedRequests.forEach { request -> + val declarationNames = + request.config.tools.orEmpty().flatMap { it.functionDeclarations.orEmpty() }.map { it.name } + assertEquals(listOf("shared_tool"), declarationNames) + assertEquals(0, request.toolsDict.count { it.name == "shared_tool" }) + } + assertEquals(1, directTool.runCalls) + assertEquals(0, additionalTool.runCalls) + } + /** * Root `MissionControl` transfers to sub-agent `HeartOfGold`, which calls a tool and then * answers. Asserts the exact 5-event sequence (transfer call → transfer response → tool call → @@ -122,3 +221,31 @@ class LlmAgentTurnTest { ) } } + +private open class RunTrackingDeclaredTool(name: String) : + BaseTool(name = name, description = "Run-tracking tool $name") { + var runCalls = 0 + + override fun declaration() = + FunctionDeclaration( + name = name, + description = description, + parameters = Schema(type = Type.OBJECT, properties = emptyMap()), + ) + + override suspend fun run(context: ToolContext, args: Map): Any { + runCalls++ + return emptyMap() + } +} + +private class ConfigOnlyRunTrackingDeclaredTool(name: String) : RunTrackingDeclaredTool(name) { + override suspend fun processLlmRequest( + toolContext: ToolContext, + llmRequest: LlmRequest, + ): LlmRequest { + val existingTools = llmRequest.config.tools?.toMutableList() ?: mutableListOf() + existingTools += Tool(functionDeclarations = listOf(declaration())) + return llmRequest.copy(config = llmRequest.config.copy(tools = existingTools)) + } +} diff --git a/core/src/commonTest/kotlin/com/google/adk/kt/tools/SkillToolsetTest.kt b/core/src/commonTest/kotlin/com/google/adk/kt/tools/SkillToolsetTest.kt index bbdd4032..7f4c72c5 100644 --- a/core/src/commonTest/kotlin/com/google/adk/kt/tools/SkillToolsetTest.kt +++ b/core/src/commonTest/kotlin/com/google/adk/kt/tools/SkillToolsetTest.kt @@ -14,17 +14,40 @@ * limitations under the License. */ +@file:OptIn(ExperimentalSkillTools::class) + package com.google.adk.kt.tools +import com.google.adk.kt.agents.LlmAgent +import com.google.adk.kt.agents.ReadonlyContext +import com.google.adk.kt.agents.toReadonlyContext +import com.google.adk.kt.annotations.ExperimentalSkillTools +import com.google.adk.kt.events.EventActions +import com.google.adk.kt.models.LlmRequest +import com.google.adk.kt.models.LlmResponse +import com.google.adk.kt.models.Model +import com.google.adk.kt.runners.InMemoryRunner +import com.google.adk.kt.sessions.SessionKey import com.google.adk.kt.skills.Frontmatter import com.google.adk.kt.skills.SkillSource import com.google.adk.kt.skills.SkillSourceException +import com.google.adk.kt.testing.modelFunctionCallResponse +import com.google.adk.kt.testing.modelMessage +import com.google.adk.kt.testing.testInvocationContext +import com.google.adk.kt.testing.testSession import com.google.adk.kt.testing.testToolContext +import com.google.adk.kt.testing.userMessage +import com.google.adk.kt.types.FunctionDeclaration +import com.google.adk.kt.types.Schema as GenaiSchema +import com.google.adk.kt.types.Type import kotlin.test.Test import kotlin.test.assertEquals import kotlin.test.assertFailsWith import kotlin.test.assertNotNull import kotlin.test.assertTrue +import kotlinx.coroutines.flow.Flow +import kotlinx.coroutines.flow.collect +import kotlinx.coroutines.flow.flow import kotlinx.coroutines.test.runTest /** Unit tests for [SkillToolset]. */ @@ -34,8 +57,22 @@ class SkillToolsetTest { listOf( Frontmatter(name = "skill1", description = "Description 1"), Frontmatter(name = "skill2", description = "Description 2"), + Frontmatter( + name = "skill3", + description = "Description 3", + metadata = + mapOf( + SkillToolset.ADK_ADDITIONAL_TOOLS_METADATA_KEY to + listOf("extra_tool_a", "extra_tool_b"), + ), + ), + ) + private val mockInstructions = + mapOf( + "skill1" to "Instructions 1", + "skill2" to "Instructions 2", + "skill3" to "Instructions 3", ) - private val mockInstructions = mapOf("skill1" to "Instructions 1", "skill2" to "Instructions 2") private val mockSource = object : SkillSource { @@ -106,7 +143,7 @@ class SkillToolsetTest { val result = tool.run(testToolContext(), emptyMap()) as Map<*, *> val skillsList = result["skills"] as? List> assertNotNull(skillsList) - assertEquals(2, skillsList.size) + assertEquals(3, skillsList.size) assertEquals("skill1", skillsList[0]["name"]) assertEquals("Description 1", skillsList[0]["description"]) assertEquals("skill2", skillsList[1]["name"]) @@ -129,6 +166,23 @@ class SkillToolsetTest { assertEquals("Description 1", frontmatter["description"]) } + @Test + fun loadSkillTool_run_returnsAdkAdditionalToolsInFrontmatterMetadata() = runTest { + val tools = skillToolset.getTools(null) + val loadSkillTool = tools.first { it.name == SkillToolset.TOOL_NAME_LOAD_SKILL } + + val result = + loadSkillTool.run(testToolContext(), mapOf(SkillToolset.PARAM_SKILL_NAME to "skill3")) + as Map<*, *> + + val frontmatter = result[SkillToolset.KEY_FRONTMATTER] as Map<*, *> + val metadata = frontmatter["metadata"] as Map<*, *> + assertEquals( + listOf("extra_tool_a", "extra_tool_b"), + metadata[SkillToolset.ADK_ADDITIONAL_TOOLS_METADATA_KEY], + ) + } + @Test fun loadSkillTool_run_requiresName() = runTest { val tools = skillToolset.getTools(null) @@ -433,4 +487,572 @@ class SkillToolsetTest { kotlin.test.assertNull(instruction) } + + @Test + fun getTools_returnsOnlyCoreTools_whenAdditionalToolsAndNoActivation() = runTest { + val toolset = + SkillToolset( + source = mockSource, + additionalTools = + listOf(makeAdditionalTool("extra_tool_a"), makeAdditionalTool("extra_tool_b")), + ) + val tools = toolset.getTools(null) + assertEquals(3, tools.size) + assertTrue(tools.none { it.name == "extra_tool_a" || it.name == "extra_tool_b" }) + } + + @Test + fun getTools_exposesAdditionalTools_afterActivation() = runTest { + val toolset = + SkillToolset( + source = mockSource, + additionalTools = + listOf(makeAdditionalTool("extra_tool_a"), makeAdditionalTool("extra_tool_b")), + ) + val session = testSession() + session.state[SkillToolset.activatedSkillToolsStateKey("test-agent")] = + mapOf("skill3" to listOf("extra_tool_a", "extra_tool_b")) + val readonlyCtx = testInvocationContext(session = session).toReadonlyContext() + val tools = toolset.getTools(readonlyCtx) + assertEquals(5, tools.size) + assertTrue(tools.any { it.name == "extra_tool_a" }) + assertTrue(tools.any { it.name == "extra_tool_b" }) + } + + @Test + fun getTools_usesStoredToolNames_withoutReloadingFrontmatter() = runTest { + var loadFrontmatterCalls = 0 + val source = + object : SkillSource by mockSource { + override suspend fun loadFrontmatter(skillName: String): Result { + loadFrontmatterCalls++ + return mockSource.loadFrontmatter(skillName) + } + } + val toolset = + SkillToolset(source = source, additionalTools = listOf(makeAdditionalTool("extra_tool_a"))) + val session = testSession() + session.state[SkillToolset.activatedSkillToolsStateKey("test-agent")] = + mapOf("skill3" to listOf("extra_tool_a")) + val readonlyCtx = testInvocationContext(session = session).toReadonlyContext() + + val tools = toolset.getTools(readonlyCtx) + + assertEquals(4, tools.size) + assertTrue(tools.any { it.name == "extra_tool_a" }) + assertEquals(0, loadFrontmatterCalls) + } + + @Test + fun getTools_skipsToolNameNotInProvidedTools() = runTest { + // skill3 declares both extra_tool_a and extra_tool_b, but only extra_tool_a is provided. + val toolset = + SkillToolset( + source = mockSource, + additionalTools = listOf(makeAdditionalTool("extra_tool_a")), + ) + val session = testSession() + session.state[SkillToolset.activatedSkillToolsStateKey("test-agent")] = + mapOf("skill3" to listOf("extra_tool_a", "extra_tool_b")) + val readonlyCtx = testInvocationContext(session = session).toReadonlyContext() + val tools = toolset.getTools(readonlyCtx) + assertEquals(4, tools.size) + assertTrue(tools.any { it.name == "extra_tool_a" }) + assertTrue(tools.none { it.name == "extra_tool_b" }) + } + + @Test + fun getTools_resolvesAdditionalToolsFromProvidedToolsets_afterActivation() = runTest { + val toolset = + SkillToolset( + source = mockSource, + additionalTools = emptyList(), + additionalToolsets = + listOf( + makeAdditionalToolset( + makeAdditionalTool("extra_tool_a"), + makeAdditionalTool("extra_tool_b"), + ) + ), + ) + val session = testSession() + session.state[SkillToolset.activatedSkillToolsStateKey("test-agent")] = + mapOf("skill3" to listOf("extra_tool_a", "extra_tool_b")) + val readonlyCtx = testInvocationContext(session = session).toReadonlyContext() + + val tools = toolset.getTools(readonlyCtx) + + assertEquals(5, tools.size) + assertTrue(tools.any { it.name == "extra_tool_a" }) + assertTrue(tools.any { it.name == "extra_tool_b" }) + } + + @Test + fun loadSkillTool_run_writesSkillAndAdditionalToolStateDelta() = runTest { + val actions = EventActions() + val ctx = testToolContext(actions = actions) + val toolset = SkillToolset(source = mockSource) + val loadSkillTool = + toolset.getTools(null).first { it.name == SkillToolset.TOOL_NAME_LOAD_SKILL } + + loadSkillTool.run(ctx, mapOf(SkillToolset.PARAM_SKILL_NAME to "skill3")) + + assertEquals( + listOf("skill3"), + actions.stateDelta[SkillToolset.activatedSkillStateKey("test-agent")], + ) + assertEquals( + mapOf("skill3" to listOf("extra_tool_a", "extra_tool_b")), + actions.stateDelta[SkillToolset.activatedSkillToolsStateKey("test-agent")], + ) + } + + @Test + fun loadSkillTool_run_preservesPreviouslyActivatedSkills() = runTest { + // Session state already has skill1 activated; loading skill3 must preserve both the prior + // activation and its associated tool names. + val session = testSession() + session.state[SkillToolset.activatedSkillStateKey("test-agent")] = listOf("skill1") + session.state[SkillToolset.activatedSkillToolsStateKey("test-agent")] = + mapOf("skill1" to listOf("previous_tool")) + val actions = EventActions() + val ctx = + testToolContext( + invocationContext = testInvocationContext(session = session), + actions = actions, + ) + val toolset = SkillToolset(source = mockSource) + val loadSkillTool = + toolset.getTools(null).first { it.name == SkillToolset.TOOL_NAME_LOAD_SKILL } + + loadSkillTool.run(ctx, mapOf(SkillToolset.PARAM_SKILL_NAME to "skill3")) + + assertEquals( + listOf("skill1", "skill3"), + actions.stateDelta[SkillToolset.activatedSkillStateKey("test-agent")], + ) + assertEquals( + mapOf( + "skill1" to listOf("previous_tool"), + "skill3" to listOf("extra_tool_a", "extra_tool_b"), + ), + actions.stateDelta[SkillToolset.activatedSkillToolsStateKey("test-agent")], + ) + } + + @Test + fun loadSkillTool_run_preservesBlankAdditionalToolNames() = runTest { + val source = + object : SkillSource by mockSource { + override suspend fun loadFrontmatter(skillName: String): Result = + Result.success( + Frontmatter( + name = skillName, + description = "Description", + metadata = + mapOf( + SkillToolset.ADK_ADDITIONAL_TOOLS_METADATA_KEY to + listOf("", " ", "extra_tool_a") + ), + ) + ) + } + val actions = EventActions() + val toolset = SkillToolset(source) + val loadSkillTool = toolset.getTools().first { it.name == SkillToolset.TOOL_NAME_LOAD_SKILL } + + loadSkillTool.run( + testToolContext(actions = actions), + mapOf(SkillToolset.PARAM_SKILL_NAME to "skill3"), + ) + + assertEquals( + mapOf("skill3" to listOf("", " ", "extra_tool_a")), + actions.stateDelta[SkillToolset.activatedSkillToolsStateKey("test-agent")], + ) + } + + @Test + fun loadSkillTool_run_replacesAdditionalToolsWhenSkillIsReloaded() = runTest { + var declaredToolNames = listOf("old_tool") + val source = + object : SkillSource by mockSource { + override suspend fun loadFrontmatter(skillName: String): Result = + Result.success( + Frontmatter( + name = skillName, + description = "Description", + metadata = + mapOf(SkillToolset.ADK_ADDITIONAL_TOOLS_METADATA_KEY to declaredToolNames), + ) + ) + } + val actions = EventActions() + val context = testToolContext(actions = actions) + val toolset = SkillToolset(source) + val loadSkillTool = toolset.getTools().first { it.name == SkillToolset.TOOL_NAME_LOAD_SKILL } + + loadSkillTool.run(context, mapOf(SkillToolset.PARAM_SKILL_NAME to "skill3")) + declaredToolNames = listOf("new_tool") + loadSkillTool.run(context, mapOf(SkillToolset.PARAM_SKILL_NAME to "skill3")) + + assertEquals( + mapOf("skill3" to listOf("new_tool")), + actions.stateDelta[SkillToolset.activatedSkillToolsStateKey("test-agent")], + ) + } + + @Test + fun getTools_duplicateActivatedToolName_returnsToolOnce() = runTest { + val toolset = + SkillToolset(mockSource, additionalTools = listOf(makeAdditionalTool("extra_tool_a"))) + val session = testSession() + session.state[SkillToolset.activatedSkillToolsStateKey("test-agent")] = + mapOf("skill3" to listOf("extra_tool_a", "extra_tool_a")) + + val tools = toolset.getTools(testInvocationContext(session = session).toReadonlyContext()) + + assertEquals(1, tools.count { it.name == "extra_tool_a" }) + } + + @Test + fun getTools_additionalToolCollidesWithCoreTool_skipsAdditionalTool() = runTest { + val toolset = + SkillToolset( + mockSource, + additionalTools = listOf(makeAdditionalTool(SkillToolset.TOOL_NAME_LOAD_SKILL)), + ) + val session = testSession() + session.state[SkillToolset.activatedSkillToolsStateKey("test-agent")] = + mapOf("skill3" to listOf(SkillToolset.TOOL_NAME_LOAD_SKILL)) + + val tools = toolset.getTools(testInvocationContext(session = session).toReadonlyContext()) + + assertEquals(1, tools.count { it.name == SkillToolset.TOOL_NAME_LOAD_SKILL }) + } + + @Test + fun close_closesOwnedResourcesOnce_withoutClosingToolsOwnedByProvidedToolsets() { + val shadowedTool = CloseTrackingTool("duplicate") + val selectedTool = CloseTrackingTool("duplicate") + val repeatedTool = CloseTrackingTool("repeated") + val toolsetOwnedTool = CloseTrackingTool("toolset_owned") + val providedToolset = CloseTrackingToolset(listOf(toolsetOwnedTool)) + val toolset = + SkillToolset( + source = mockSource, + additionalTools = + listOf(shadowedTool, selectedTool, repeatedTool, repeatedTool), + additionalToolsets = listOf(providedToolset, providedToolset), + ) + + toolset.close() + + assertEquals(1, shadowedTool.closeCalls) + assertEquals(1, selectedTool.closeCalls) + assertEquals(1, repeatedTool.closeCalls) + assertEquals(1, providedToolset.closeCalls) + assertEquals(0, toolsetOwnedTool.closeCalls) + } + + @Test + fun constructor_snapshotsOwnedInputLists() { + val originalTool = CloseTrackingTool("original_tool") + val laterTool = CloseTrackingTool("later_tool") + val originalToolset = CloseTrackingToolset(emptyList()) + val laterToolset = CloseTrackingToolset(emptyList()) + val additionalTools = mutableListOf(originalTool) + val additionalToolsets = mutableListOf(originalToolset) + val toolset = + SkillToolset( + source = mockSource, + additionalTools = additionalTools, + additionalToolsets = additionalToolsets, + ) + + additionalTools.apply { + clear() + add(laterTool) + } + additionalToolsets.apply { + clear() + add(laterToolset) + } + toolset.close() + + assertEquals(1, originalTool.closeCalls) + assertEquals(0, laterTool.closeCalls) + assertEquals(1, originalToolset.closeCalls) + assertEquals(0, laterToolset.closeCalls) + } + + @Test + fun runAsync_loadSkill_exposesStoredAdditionalToolsOnNextModelStep() = runTest { + val capturedRequests = mutableListOf() + val responses = + listOf( + modelFunctionCallResponse( + SkillToolset.TOOL_NAME_LOAD_SKILL, + args = mapOf(SkillToolset.PARAM_SKILL_NAME to "skill3"), + id = "load-skill-call", + ), + LlmResponse(content = modelMessage("done")), + ) + var responseIndex = 0 + val model = + object : Model { + override val name = "skill-activation-model" + + override fun generateContent(request: LlmRequest, stream: Boolean): Flow = + flow { + capturedRequests += request + emit(responses[responseIndex++]) + } + } + val skillToolset = + SkillToolset( + source = mockSource, + additionalTools = + listOf(makeAdditionalTool("extra_tool_a"), makeAdditionalTool("extra_tool_b")), + ) + val agent = LlmAgent(name = "test-agent", model = model, toolsets = listOf(skillToolset)) + val runner = InMemoryRunner(agent) + + runner + .runAsync(userId = "user", sessionId = "session", newMessage = userMessage("use skill3")) + .collect {} + + val session = + runner.sessionService.getSession(SessionKey(runner.appName, "user", "session")) + assertNotNull(session) + assertEquals( + mapOf("skill3" to listOf("extra_tool_a", "extra_tool_b")), + session.state[SkillToolset.activatedSkillToolsStateKey("test-agent")], + ) + assertEquals(2, capturedRequests.size) + assertTrue(capturedRequests[0].toolsDict.none { it.name.startsWith("extra_tool_") }) + assertTrue( + capturedRequests[0].functionDeclarationNames().none { it.startsWith("extra_tool_") } + ) + assertEquals( + setOf("extra_tool_a", "extra_tool_b"), + capturedRequests[1].toolsDict + .filter { it.name.startsWith("extra_tool_") } + .mapTo(mutableSetOf()) { it.name }, + ) + assertEquals( + setOf("extra_tool_a", "extra_tool_b"), + capturedRequests[1] + .functionDeclarationNames() + .filterTo(mutableSetOf()) { it.startsWith("extra_tool_") }, + ) + } + + @Test + fun runAsync_additionalToolCollidesWithDirectTool_keepsDirectTool() = runTest { + val capturedRequests = mutableListOf() + val responses = + listOf( + modelFunctionCallResponse( + SkillToolset.TOOL_NAME_LOAD_SKILL, + args = mapOf(SkillToolset.PARAM_SKILL_NAME to "skill3"), + id = "load-skill-call", + ), + modelFunctionCallResponse("extra_tool_a", id = "extra-tool-call"), + LlmResponse(content = modelMessage("done")), + ) + var responseIndex = 0 + val model = + object : Model { + override val name = "tool-collision-model" + + override fun generateContent(request: LlmRequest, stream: Boolean): Flow = + flow { + capturedRequests += request + emit(responses[responseIndex++]) + } + } + val directTool = RunTrackingTool("extra_tool_a") + val skillTool = RunTrackingTool("extra_tool_a") + val skillToolset = SkillToolset(mockSource, additionalTools = listOf(skillTool)) + val agent = + LlmAgent( + name = "test-agent", + model = model, + tools = listOf(directTool), + toolsets = listOf(skillToolset), + ) + + InMemoryRunner(agent) + .runAsync(userId = "user", sessionId = "session", newMessage = userMessage("use skill3")) + .collect {} + + assertTrue(capturedRequests.size >= 2) + assertEquals(1, capturedRequests[1].functionDeclarationNames().count { it == "extra_tool_a" }) + assertEquals(1, directTool.runCalls) + assertEquals(0, skillTool.runCalls) + } + + @Test + fun runAsync_additionalToolCollidesWithEarlierToolset_keepsEarlierTool() = runTest { + val capturedRequests = mutableListOf() + val responses = + listOf( + modelFunctionCallResponse( + SkillToolset.TOOL_NAME_LOAD_SKILL, + args = mapOf(SkillToolset.PARAM_SKILL_NAME to "skill3"), + id = "load-skill-call", + ), + modelFunctionCallResponse("extra_tool_a", id = "extra-tool-call"), + LlmResponse(content = modelMessage("done")), + ) + var responseIndex = 0 + val model = + object : Model { + override val name = "toolset-collision-model" + + override fun generateContent(request: LlmRequest, stream: Boolean): Flow = + flow { + capturedRequests += request + emit(responses[responseIndex++]) + } + } + val earlierTool = RunTrackingTool("extra_tool_a") + val skillTool = RunTrackingTool("extra_tool_a") + val skillToolset = SkillToolset(mockSource, additionalTools = listOf(skillTool)) + val agent = + LlmAgent( + name = "test-agent", + model = model, + toolsets = listOf(makeAdditionalToolset(earlierTool), skillToolset), + ) + + InMemoryRunner(agent) + .runAsync(userId = "user", sessionId = "session", newMessage = userMessage("use skill3")) + .collect {} + + assertTrue(capturedRequests.size >= 2) + assertEquals(1, capturedRequests[1].functionDeclarationNames().count { it == "extra_tool_a" }) + assertEquals(1, earlierTool.runCalls) + assertEquals(0, skillTool.runCalls) + } + + @Test + fun runAsync_additionalToolCollidesWithEarlierToolsetHook_keepsHookTool() = runTest { + val capturedRequests = mutableListOf() + val responses = + listOf( + modelFunctionCallResponse( + SkillToolset.TOOL_NAME_LOAD_SKILL, + args = mapOf(SkillToolset.PARAM_SKILL_NAME to "skill3"), + id = "load-skill-call", + ), + modelFunctionCallResponse("extra_tool_a", id = "extra-tool-call"), + LlmResponse(content = modelMessage("done")), + ) + var responseIndex = 0 + val model = + object : Model { + override val name = "toolset-hook-collision-model" + + override fun generateContent(request: LlmRequest, stream: Boolean): Flow = + flow { + capturedRequests += request + emit(responses[responseIndex++]) + } + } + val hookTool = RunTrackingTool("extra_tool_a") + val hookToolset = + object : Toolset { + override suspend fun getTools(readonlyContext: ReadonlyContext?): List = + emptyList() + + override suspend fun processLlmRequest( + toolContext: ToolContext, + llmRequest: LlmRequest, + ): LlmRequest = llmRequest.appendTools(listOf(hookTool)) + } + val skillTool = RunTrackingTool("extra_tool_a") + val skillToolset = SkillToolset(mockSource, additionalTools = listOf(skillTool)) + val agent = + LlmAgent( + name = "test-agent", + model = model, + toolsets = listOf(hookToolset, skillToolset), + ) + + InMemoryRunner(agent) + .runAsync(userId = "user", sessionId = "session", newMessage = userMessage("use skill3")) + .collect {} + + assertTrue(capturedRequests.size >= 2) + assertEquals(1, capturedRequests[1].functionDeclarationNames().count { it == "extra_tool_a" }) + assertEquals(1, capturedRequests[1].toolsDict.count { it.name == "extra_tool_a" }) + assertEquals(1, hookTool.runCalls) + assertEquals(0, skillTool.runCalls) + } +} + +/** Minimal [BaseTool] implementation for SkillToolset activation tests. */ +private fun makeAdditionalTool(name: String): BaseTool = + object : + BaseTool(name = name, description = "Additional tool $name for tests.") { + override fun declaration() = + FunctionDeclaration( + name = name, + description = description, + parameters = GenaiSchema(type = Type.OBJECT, properties = emptyMap()), + ) + + override suspend fun run(context: ToolContext, args: Map): Any = + emptyMap() + } + +private fun makeAdditionalToolset(vararg tools: BaseTool): Toolset = + object : Toolset { + override suspend fun getTools(readonlyContext: ReadonlyContext?) = tools.toList() + } + +private class CloseTrackingTool(name: String) : + BaseTool(name = name, description = "Close-tracking tool $name") { + var closeCalls = 0 + + override fun declaration(): FunctionDeclaration? = null + + override suspend fun run(context: ToolContext, args: Map): Any = + emptyMap() + + override fun close() { + closeCalls++ + } +} + +private class CloseTrackingToolset(private val tools: List) : Toolset { + var closeCalls = 0 + + override suspend fun getTools(readonlyContext: ReadonlyContext?): List = tools + + override fun close() { + closeCalls++ + } } + +private class RunTrackingTool(name: String) : + BaseTool(name = name, description = "Run-tracking tool $name") { + var runCalls = 0 + + override fun declaration() = + FunctionDeclaration( + name = name, + description = description, + parameters = GenaiSchema(type = Type.OBJECT, properties = emptyMap()), + ) + + override suspend fun run(context: ToolContext, args: Map): Any { + runCalls++ + return mapOf("source" to "direct") + } +} + +private fun LlmRequest.functionDeclarationNames(): List = + config.tools.orEmpty().flatMap { it.functionDeclarations.orEmpty() }.map { it.name }