From b7fa8fdf48a588f6ef9eb83759e5980249e605a8 Mon Sep 17 00:00:00 2001 From: Damian Momot Date: Wed, 9 Sep 2026 23:18:59 -0700 Subject: [PATCH] feat: add Plugin onRunError notification callback Adds Plugin.onRunError(invocationContext, error), a notification-only hook the runner invokes when a run fails, before re-raising the error. Mirrors ADK Python and ADK Java. The Java-friendly BaseFuturePlugin base gains a matching onRunErrorAsync hook. PiperOrigin-RevId: 978974181 --- .../adk/tokt/adapters/JavaPluginToKt.kt | 13 +- .../google/adk/tokt/adapters/KtToolToJava.kt | 67 +++++ .../google/adk/tokt/KtRunnerInteropTest.kt | 281 ++++++++++++++++++ 3 files changed, 356 insertions(+), 5 deletions(-) create mode 100644 tokt/src/main/kotlin/com/google/adk/tokt/adapters/KtToolToJava.kt diff --git a/tokt/src/main/kotlin/com/google/adk/tokt/adapters/JavaPluginToKt.kt b/tokt/src/main/kotlin/com/google/adk/tokt/adapters/JavaPluginToKt.kt index bbf000652..0aa50be66 100644 --- a/tokt/src/main/kotlin/com/google/adk/tokt/adapters/JavaPluginToKt.kt +++ b/tokt/src/main/kotlin/com/google/adk/tokt/adapters/JavaPluginToKt.kt @@ -49,7 +49,9 @@ import kotlinx.coroutines.withContext * Exposes an ADK Java [JavaPlugin] as a Kotlin [KtPlugin] so a Java app's plugins run on the Kotlin * runner. Each callback converts the Kotlin context to its Java view, invokes the plugin off the * engine dispatcher, and reconciles the actions it wrote back onto the Kotlin side. Tool-level - * callbacks apply only to Kt-backed Java tools ([JavaToolToKt]). + * callbacks fire for every tool: an adapted Java tool ([JavaToolToKt]) unwraps to its original + * instance, and a native Kotlin tool is presented through an inspection-only [KtToolToJava] view + * (see [ktToolAsJava]) -- so the plugin can read the tool and write actions, but must not run it. */ internal class JavaPluginToKt(internal val plugin: JavaPlugin) : KtPlugin { @@ -189,14 +191,15 @@ internal class JavaPluginToKt(internal val plugin: JavaPlugin) : KtPlugin { else CallbackChoice.Continue(Unit) } - // Tool-level callbacks (only Kt-backed Java tools). + // Tool-level callbacks. An adapted Java tool fires natively; a native Kotlin tool is passed + // through an inspection-only Java view ([KtToolToJava]). override suspend fun beforeTool( context: KtToolContext, tool: KtBaseTool, args: Map, ): CallbackChoice, Map> { - val javaTool = (tool as? JavaToolToKt)?.javaTool ?: return CallbackChoice.Continue(args) + val javaTool = ktToolAsJava(tool) val javaContext = ktToolContextToJava(context) val mutableArgs = args.toMutableMap() val override = onIo { plugin.beforeToolCallback(javaTool, mutableArgs, javaContext) } @@ -211,7 +214,7 @@ internal class JavaPluginToKt(internal val plugin: JavaPlugin) : KtPlugin { args: Map, result: Map, ): Map { - val javaTool = (tool as? JavaToolToKt)?.javaTool ?: return result + val javaTool = ktToolAsJava(tool) val javaContext = ktToolContextToJava(context) val override = onIo { plugin.afterToolCallback(javaTool, args, javaContext, result) } reconcileActionsToKt(javaContext.actions(), context.actions) @@ -224,7 +227,7 @@ internal class JavaPluginToKt(internal val plugin: JavaPlugin) : KtPlugin { args: Map, error: Throwable, ): CallbackChoice> { - val javaTool = (tool as? JavaToolToKt)?.javaTool ?: return CallbackChoice.Continue(Unit) + val javaTool = ktToolAsJava(tool) val javaContext = ktToolContextToJava(context) val fallback = onIo { plugin.onToolErrorCallback(javaTool, args, javaContext, error) } reconcileActionsToKt(javaContext.actions(), context.actions) diff --git a/tokt/src/main/kotlin/com/google/adk/tokt/adapters/KtToolToJava.kt b/tokt/src/main/kotlin/com/google/adk/tokt/adapters/KtToolToJava.kt new file mode 100644 index 000000000..92b64e545 --- /dev/null +++ b/tokt/src/main/kotlin/com/google/adk/tokt/adapters/KtToolToJava.kt @@ -0,0 +1,67 @@ +/* + * 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.tokt.adapters + +import com.google.adk.kt.tools.BaseTool as KtBaseTool +import com.google.adk.tokt.codecs.FunctionDeclarationCodec +import com.google.adk.tools.BaseTool as JavaBaseTool +import com.google.adk.tools.ToolContext as JavaToolContext +import com.google.genai.types.FunctionDeclaration as GenaiFunctionDeclaration +import io.reactivex.rxjava3.core.Single +import java.util.Optional + +/** + * An inspection-only Java view of a native Kotlin [KtBaseTool], so ADK Java plugin tool callbacks + * (`beforeToolCallback`, `afterToolCallback`, `onToolErrorCallback`) can read a native Kotlin tool + * as an ADK Java [JavaBaseTool]: name, description, long-running flag, declaration, and custom + * metadata. It is not meant to be executed -- the Kotlin runner is the sole driver of tool + * execution -- so [runAsync] fails loud with a descriptive error rather than doing anything. + * Mirrors [KtAgentToJava]. + */ +internal class KtToolToJava(internal val ktTool: KtBaseTool) : + JavaBaseTool(ktTool.name, ktTool.description, ktTool.isLongRunning) { + + init { + for ((key, value) in ktTool.customMetadata) { + setCustomMetadata(key, value) + } + } + + override fun declaration(): Optional = + Optional.ofNullable(ktTool.declaration()?.let { FunctionDeclarationCodec.toJava(it) }) + + @JvmSuppressWildcards + override fun runAsync( + args: Map, + toolContext: JavaToolContext, + ): Single> = + Single.error( + UnsupportedOperationException( + "This is an inspection-only view of a Kotlin engine tool, handed to Java plugin callbacks " + + "via ktToolAsJava; the Kotlin runner is the sole driver of tool execution, so it cannot " + + "be run from Java." + ) + ) +} + +/** + * Presents a Kotlin [tool] to a Java plugin callback as an ADK Java [JavaBaseTool]: a round-tripped + * adapted Java tool ([JavaToolToKt]) unwraps to its original instance (so it fires natively), and a + * native Kotlin tool gets an inspection-only [KtToolToJava] view. + */ +internal fun ktToolAsJava(tool: KtBaseTool): JavaBaseTool = + (tool as? JavaToolToKt)?.javaTool ?: KtToolToJava(tool) diff --git a/tokt/src/test/kotlin/com/google/adk/tokt/KtRunnerInteropTest.kt b/tokt/src/test/kotlin/com/google/adk/tokt/KtRunnerInteropTest.kt index 8c088000f..727e1e133 100644 --- a/tokt/src/test/kotlin/com/google/adk/tokt/KtRunnerInteropTest.kt +++ b/tokt/src/test/kotlin/com/google/adk/tokt/KtRunnerInteropTest.kt @@ -40,6 +40,8 @@ import com.google.adk.kt.runners.Runner as KtRunner import com.google.adk.kt.sessions.GetSessionConfig as KtGetSessionConfig import com.google.adk.kt.sessions.SessionKey as KtSessionKey import com.google.adk.kt.sessions.State as KtState +import com.google.adk.kt.tools.BaseTool as KtBaseTool +import com.google.adk.kt.tools.ToolContext as KtToolContext import com.google.adk.kt.types.Blob as KtBlob import com.google.adk.kt.types.Content as KtContent import com.google.adk.kt.types.FileData as KtFileData @@ -515,6 +517,101 @@ class KtRunnerInteropTest { } } + /** A native Kotlin tool (not an adapted Java tool), used to prove plugin tool callbacks fire. */ + private class NativeKtEchoTool : KtBaseTool("native_kt_echo", "native kotlin echo") { + override fun declaration(): com.google.adk.kt.types.FunctionDeclaration? = null + + override suspend fun run(context: KtToolContext, args: Map): Any = + mapOf("echoed" to (args["text"] ?: "")) + } + + /** A native Kotlin tool whose body always fails - used to drive the onToolError path. */ + private class NativeKtThrowingTool : + KtBaseTool("native_kt_throwing", "native kotlin failing tool") { + override fun declaration(): com.google.adk.kt.types.FunctionDeclaration? = null + + override suspend fun run(context: KtToolContext, args: Map): Any = + throw IllegalStateException("native tool boom") + } + + /** + * A Java plugin that records the tools its beforeTool callback is handed and denies each one - + * used to prove the callback fires, with the real tool name, even for a native Kotlin tool. + */ + private class ToolDenyingJavaPlugin : JavaBasePlugin("tool_denying_plugin") { + val seenToolNames = CopyOnWriteArrayList() + + @JvmSuppressWildcards + override fun beforeToolCallback( + tool: JavaBaseTool, + toolArgs: Map, + toolContext: JavaToolContext, + ): Maybe> { + seenToolNames.add(tool.name()) + return Maybe.just(mapOf("error" to "denied by policy: ${tool.name()}")) + } + } + + /** A Java plugin whose beforeTool tries to *run* the tool it is handed, capturing any error. */ + private class ToolRunningJavaPlugin : JavaBasePlugin("tool_running_plugin") { + var runError: Throwable? = null + + @JvmSuppressWildcards + override fun beforeToolCallback( + tool: JavaBaseTool, + toolArgs: Map, + toolContext: JavaToolContext, + ): Maybe> { + try { + val unused = tool.runAsync(toolArgs, toolContext).blockingGet() + } catch (e: Throwable) { + runError = e + } + return Maybe.empty() + } + } + + /** + * A Java plugin whose afterTool callback records the tools it is handed and replaces each + * result - used to prove the callback fires, with the real tool name, even for a native Kotlin + * tool. + */ + private class ToolResultOverridingJavaPlugin : JavaBasePlugin("tool_result_overriding_plugin") { + val seenToolNames = CopyOnWriteArrayList() + + @JvmSuppressWildcards + override fun afterToolCallback( + tool: JavaBaseTool, + toolArgs: Map, + toolContext: JavaToolContext, + result: Map, + ): Maybe> { + seenToolNames.add(tool.name()) + return Maybe.just(mapOf("overridden" to tool.name())) + } + } + + /** + * A Java plugin whose onToolError callback records the tool and error it is handed and returns a + * fallback - used to prove the callback fires for a native Kotlin tool and can swallow its error. + */ + private class ToolErrorSwallowingJavaPlugin : JavaBasePlugin("tool_error_swallowing_plugin") { + val seenToolNames = CopyOnWriteArrayList() + var seenError: Throwable? = null + + @JvmSuppressWildcards + override fun onToolErrorCallback( + tool: JavaBaseTool, + toolArgs: Map, + toolContext: JavaToolContext, + error: Throwable, + ): Maybe> { + seenToolNames.add(tool.name()) + seenError = error + return Maybe.just(mapOf("recovered" to "from ${tool.name()}")) + } + } + /** A Java model that replays a fixed sequence of responses, one per LLM step. */ private class SequentialJavaModel(private val turns: List) : JavaBaseLlm("java-model") { @@ -647,6 +744,190 @@ class KtRunnerInteropTest { ) } + @Test + fun ktRunner_pluginBeforeToolCallback_firesForNativeKotlinTool_andCanDenyIt() = runBlocking { + // A native Kotlin tool has no Java form, yet a Java plugin's beforeToolCallback must still + // reach + // it (via an inspection-only Java view) so it can read the tool by name and short-circuit the + // call - the mechanism tool-call policy enforcement relies on. + val plugin = ToolDenyingJavaPlugin() + val agent = + KtLlmAgent( + name = "a", + model = + JavaAdkToKt.asKtModel( + SequentialJavaModel( + listOf(modelFunctionCall("native_kt_echo", mapOf("text" to "hi")), modelText("done")) + ) + ), + tools = listOf(NativeKtEchoTool()), + ) + val runner = + KtInMemoryRunner( + app = + KtApp( + appName = "app", + rootAgent = agent, + plugins = listOf(JavaAdkToKt.asKtPlugin(plugin)), + ) + ) + + val events = runner.turn() + + // The callback fired for the native Kotlin tool, seeing its real name. + assertEquals(listOf("native_kt_echo"), plugin.seenToolNames.toList()) + + // The denial short-circuited execution: the function response is the injected error, not the + // tool's own "echoed" output. + val functionResponse = + events.flatMap { it.functionResponses() }.singleOrNull() + ?: fail("expected exactly one function response") + assertEquals("denied by policy: native_kt_echo", functionResponse.response["error"]) + assertTrue( + !functionResponse.response.containsKey("echoed"), + "the tool body must not have run after a beforeTool denial", + ) + } + + @Test + fun ktRunner_pluginAfterToolCallback_firesForNativeKotlinTool_andCanOverrideResult() = + runBlocking { + // afterToolCallback is wired through the same inspection-only Java view as beforeTool, so it + // must reach a native Kotlin tool too and be able to replace its result. + val plugin = ToolResultOverridingJavaPlugin() + val agent = + KtLlmAgent( + name = "a", + model = + JavaAdkToKt.asKtModel( + SequentialJavaModel( + listOf( + modelFunctionCall("native_kt_echo", mapOf("text" to "hi")), + modelText("done"), + ) + ) + ), + tools = listOf(NativeKtEchoTool()), + ) + val runner = + KtInMemoryRunner( + app = + KtApp( + appName = "app", + rootAgent = agent, + plugins = listOf(JavaAdkToKt.asKtPlugin(plugin)), + ) + ) + + val events = runner.turn() + + // The callback fired for the native Kotlin tool, seeing its real name. + assertEquals(listOf("native_kt_echo"), plugin.seenToolNames.toList()) + + // The override replaced the tool's own "echoed" output. + val functionResponse = + events.flatMap { it.functionResponses() }.singleOrNull() + ?: fail("expected exactly one function response") + assertEquals("native_kt_echo", functionResponse.response["overridden"]) + assertTrue( + !functionResponse.response.containsKey("echoed"), + "the afterTool override must replace the tool's own result", + ) + } + + @Test + fun ktRunner_pluginOnToolErrorCallback_firesForNativeKotlinTool_andCanSwallowTheError() = + runBlocking { + // onToolErrorCallback is wired through the same inspection-only Java view, so a Java plugin + // can + // observe and recover from a native Kotlin tool's failure, returning a fallback in its place. + val plugin = ToolErrorSwallowingJavaPlugin() + val agent = + KtLlmAgent( + name = "a", + model = + JavaAdkToKt.asKtModel( + SequentialJavaModel( + listOf( + modelFunctionCall("native_kt_throwing", mapOf("text" to "hi")), + modelText("done"), + ) + ) + ), + tools = listOf(NativeKtThrowingTool()), + ) + val runner = + KtInMemoryRunner( + app = + KtApp( + appName = "app", + rootAgent = agent, + plugins = listOf(JavaAdkToKt.asKtPlugin(plugin)), + ) + ) + + val events = runner.turn() + + // The callback fired for the native Kotlin tool, seeing its real name and the exact error the + // tool body threw (delivered unwrapped, so this is the tool's failure, not a framework one). + assertEquals(listOf("native_kt_throwing"), plugin.seenToolNames.toList()) + val seenError = plugin.seenError + assertTrue( + seenError is IllegalStateException, + "onToolError should receive the tool's thrown error, got $seenError", + ) + assertEquals("native tool boom", seenError?.message) + + // The fallback swallowed the error: the run produced the recovery result rather than failing. + val functionResponse = + events.flatMap { it.functionResponses() }.singleOrNull() + ?: fail("expected exactly one function response") + assertEquals("from native_kt_throwing", functionResponse.response["recovered"]) + } + + @Test + fun ktRunner_inspectionOnlyToolView_failsLoudIfAPluginRunsANativeKotlinTool() = runBlocking { + // The native Kotlin tool is handed to the plugin as an inspection-only Java view; running it + // must fail loud (the Kotlin runner is the sole driver of tool execution), mirroring the + // agent inspection-only view. + val plugin = ToolRunningJavaPlugin() + val agent = + KtLlmAgent( + name = "a", + model = + JavaAdkToKt.asKtModel( + SequentialJavaModel( + listOf(modelFunctionCall("native_kt_echo", mapOf("text" to "hi")), modelText("done")) + ) + ), + tools = listOf(NativeKtEchoTool()), + ) + val runner = + KtInMemoryRunner( + app = + KtApp( + appName = "app", + rootAgent = agent, + plugins = listOf(JavaAdkToKt.asKtPlugin(plugin)), + ) + ) + + runner.turn() + + val runError = plugin.runError + assertTrue( + runError is UnsupportedOperationException, + "running the inspection-only tool view should fail with UnsupportedOperationException, got " + + "$runError", + ) + // BaseTool.runAsync's default also throws UnsupportedOperationException, so assert the message. + assertTrue( + runError?.message?.contains("inspection-only") == true, + "the failure should be KtToolToJava's descriptive error mentioning \"inspection-only\", got " + + "${runError?.message}", + ) + } + @Test fun ktRunner_modelPartCarryingOnlyAnUnmappedKind_isDropped() = runBlocking { // executableCode has no Kotlin counterpart, so a part carrying only it must be dropped rather