From 29481a98afc0c919998294562b67631d404829b5 Mon Sep 17 00:00:00 2001 From: Damian Momot Date: Mon, 21 Sep 2026 03:22:12 -0700 Subject: [PATCH] feat: let the ADK Java <-> Kotlin interop run on a caller-supplied dispatcher JavaAdkToKt and KotlinAdkToJava now accept an optional CoroutineDispatcher (defaulting to Dispatchers.IO) when bridging blocking Java tools, toolsets, plugins, models, and services onto the Kotlin engine. The parameter is defaulted and exposed to Java via @JvmOverloads, so existing callers are unaffected. The supplied dispatcher must be able to service nested bridged calls concurrently, so a single-threaded or tightly bounded one can deadlock. PiperOrigin-RevId: 985155333 --- .../com/google/adk/tokt/InteropDispatcher.kt | 8 +- .../kotlin/com/google/adk/tokt/JavaAdkToKt.kt | 110 +++++++++++----- .../com/google/adk/tokt/KotlinAdkToJava.kt | 14 +- .../com/google/adk/tokt/KtRunnerToJava.kt | 13 +- .../google/adk/tokt/adapters/JavaModelToKt.kt | 17 ++- .../adk/tokt/adapters/JavaPluginToKt.kt | 75 ++++++----- .../google/adk/tokt/adapters/JavaToolToKt.kt | 23 ++-- .../adk/tokt/adapters/JavaToolsetToKt.kt | 17 ++- .../tokt/context/KtCallbackContextToJava.kt | 28 ++-- .../com/google/adk/tokt/context/KtToJava.kt | 33 +++-- .../tokt/services/JavaArtifactServiceToKt.kt | 22 ++-- .../tokt/services/JavaMemoryServiceToKt.kt | 16 ++- .../tokt/services/JavaSessionServiceToKt.kt | 30 ++--- .../tokt/services/KtArtifactServiceToJava.kt | 22 ++-- .../tokt/services/KtMemoryServiceToJava.kt | 16 +-- .../tokt/services/KtSessionServiceToJava.kt | 26 ++-- .../adk/tokt/services/ServiceAdapters.kt | 46 +++++-- .../google/adk/tokt/KtRunnerInteropTest.kt | 124 ++++++++++++++++++ 18 files changed, 441 insertions(+), 199 deletions(-) diff --git a/tokt/src/main/kotlin/com/google/adk/tokt/InteropDispatcher.kt b/tokt/src/main/kotlin/com/google/adk/tokt/InteropDispatcher.kt index bf738509d..bb300470d 100644 --- a/tokt/src/main/kotlin/com/google/adk/tokt/InteropDispatcher.kt +++ b/tokt/src/main/kotlin/com/google/adk/tokt/InteropDispatcher.kt @@ -20,10 +20,10 @@ import kotlinx.coroutines.CoroutineDispatcher import kotlinx.coroutines.Dispatchers /** - * The dispatcher every crossing in this module hops to. + * The default dispatcher a crossing in this module hops to when the caller supplies none. * - * ADK Java's SPI is RxJava, which is synchronous unless the implementation says otherwise, so a - * user-authored Java tool, plugin or service may block. Running it on the engine's dispatcher would - * stall the coroutine driving the agent loop, so each adapter moves the call here. + * ADK Java's SPI is synchronous RxJava, so a user-authored Java tool, plugin or service may block; + * running it on the coroutine driving the agent loop would stall it, so each adapter moves the call + * off. The [JavaAdkToKt] and [KotlinAdkToJava] entry points accept a `dispatcher` to override this. */ internal val InteropDispatcher: CoroutineDispatcher = Dispatchers.IO diff --git a/tokt/src/main/kotlin/com/google/adk/tokt/JavaAdkToKt.kt b/tokt/src/main/kotlin/com/google/adk/tokt/JavaAdkToKt.kt index 4ed0f166a..e75ca84e4 100644 --- a/tokt/src/main/kotlin/com/google/adk/tokt/JavaAdkToKt.kt +++ b/tokt/src/main/kotlin/com/google/adk/tokt/JavaAdkToKt.kt @@ -37,6 +37,7 @@ import com.google.adk.tokt.services.javaMemoryServiceAsKt import com.google.adk.tokt.services.javaSessionServiceAsKt import com.google.adk.tools.BaseTool as JavaBaseTool import com.google.adk.tools.BaseToolset as JavaBaseToolset +import kotlinx.coroutines.CoroutineDispatcher /** * Forward interop entry point: adapts ADK Java tools, toolsets, plugins, services, and models so @@ -46,7 +47,11 @@ import com.google.adk.tools.BaseToolset as JavaBaseToolset * An adapted component behaves as it does on ADK Java. It sees the session as it currently stands, * including events and state written earlier in the same turn, and its state, artifact and * control-flow writes reach the engine. Blocking work is fine: calls are dispatched off the thread - * driving the agent. + * driving the agent, onto the optional `dispatcher` (default `Dispatchers.IO`) each conversion + * accepts. That dispatcher must be able to run nested bridged calls concurrently -- a bridged tool + * or plugin that itself makes a blocking bridged call (e.g. one that blocks on a bridged service) + * holds its thread until that call returns, so a single-threaded or tightly bounded dispatcher can + * deadlock; the default `Dispatchers.IO` grows its pool and avoids this. * * A bridged plugin's error callbacks fire: `onRunErrorCallback` is notification-only -- the engine * re-raises the run's error to the caller afterwards regardless, so it cannot recover the run (it @@ -65,59 +70,104 @@ import com.google.adk.tools.BaseToolset as JavaBaseToolset */ object JavaAdkToKt { - /** Adapts an ADK Java tool. */ - @JvmStatic fun asKtTool(javaTool: JavaBaseTool): KtBaseTool = JavaToolToKt(javaTool) + /** Adapts an ADK Java tool, hopping to `dispatcher` for its (possibly blocking) calls. */ + @JvmStatic + @JvmOverloads + fun asKtTool( + javaTool: JavaBaseTool, + dispatcher: CoroutineDispatcher = InteropDispatcher, + ): KtBaseTool = JavaToolToKt(javaTool, dispatcher) /** - * Adapts a whole collection of ADK Java tools (e.g. an `LlmAgent`'s `tools`). Kept alongside - * [asKtTool] for Java callers, who would otherwise write `stream().map(...).toList()`. + * Adapts a whole collection of ADK Java tools (e.g. an `LlmAgent`'s `tools`), each on + * `dispatcher`. Kept alongside [asKtTool] for Java callers, who would otherwise write + * `stream().map(...).toList()`. */ @JvmStatic - fun asKtTools(javaTools: List): List = javaTools.map { asKtTool(it) } + @JvmOverloads + fun asKtTools( + javaTools: List, + dispatcher: CoroutineDispatcher = InteropDispatcher, + ): List = javaTools.map { asKtTool(it, dispatcher) } - /** Adapts an ADK Java toolset. */ - @JvmStatic fun asKtToolset(javaToolset: JavaBaseToolset): KtToolset = JavaToolsetToKt(javaToolset) + /** Adapts an ADK Java toolset, hopping to `dispatcher` for its (possibly blocking) calls. */ + @JvmStatic + @JvmOverloads + fun asKtToolset( + javaToolset: JavaBaseToolset, + dispatcher: CoroutineDispatcher = InteropDispatcher, + ): KtToolset = JavaToolsetToKt(javaToolset, dispatcher) - /** Adapts a whole collection of ADK Java toolsets. */ + /** Adapts a whole collection of ADK Java toolsets, each on `dispatcher`. */ @JvmStatic - fun asKtToolsets(javaToolsets: List): List = javaToolsets.map { - asKtToolset(it) - } + @JvmOverloads + fun asKtToolsets( + javaToolsets: List, + dispatcher: CoroutineDispatcher = InteropDispatcher, + ): List = javaToolsets.map { asKtToolset(it, dispatcher) } - /** Adapts an ADK Java plugin. */ - @JvmStatic fun asKtPlugin(javaPlugin: JavaPlugin): KtPlugin = JavaPluginToKt(javaPlugin) + /** Adapts an ADK Java plugin, hopping to `dispatcher` for its (possibly blocking) callbacks. */ + @JvmStatic + @JvmOverloads + fun asKtPlugin( + javaPlugin: JavaPlugin, + dispatcher: CoroutineDispatcher = InteropDispatcher, + ): KtPlugin = JavaPluginToKt(javaPlugin, dispatcher) - /** Adapts a whole collection of ADK Java plugins (e.g. a `Runner`'s `plugins`). */ + /** + * Adapts a whole collection of ADK Java plugins (e.g. a `Runner`'s `plugins`), each on + * `dispatcher`. + */ @JvmStatic - fun asKtPlugins(javaPlugins: List): List = javaPlugins.map { - asKtPlugin(it) - } + @JvmOverloads + fun asKtPlugins( + javaPlugins: List, + dispatcher: CoroutineDispatcher = InteropDispatcher, + ): List = javaPlugins.map { asKtPlugin(it, dispatcher) } - /** Adapts an ADK Java model so the Kotlin engine can call it. */ - @JvmStatic fun asKtModel(javaLlm: JavaBaseLlm): KtModel = JavaModelToKt(javaLlm) + /** + * Adapts an ADK Java model so the Kotlin engine can call it, running its generation on + * `dispatcher`. + */ + @JvmStatic + @JvmOverloads + fun asKtModel( + javaLlm: JavaBaseLlm, + dispatcher: CoroutineDispatcher = InteropDispatcher, + ): KtModel = JavaModelToKt(javaLlm, dispatcher) /** * Adapts an ADK Java session service for the Kotlin engine, unwrapping a round-tripped Kotlin one - * rather than stacking a second adapter. A `rewindBeforeInvocationId` does not survive, since ADK - * Java has no such field. + * rather than stacking a second adapter. Its calls run on `dispatcher`. A + * `rewindBeforeInvocationId` does not survive, since ADK Java has no such field. */ @JvmStatic - fun asKtSessionService(service: JavaSessionService): KtSessionService = - javaSessionServiceAsKt(service) + @JvmOverloads + fun asKtSessionService( + service: JavaSessionService, + dispatcher: CoroutineDispatcher = InteropDispatcher, + ): KtSessionService = javaSessionServiceAsKt(service, dispatcher) /** * Adapts an ADK Java artifact service for the Kotlin engine, unwrapping a round-tripped Kotlin - * one rather than stacking adapters. An empty or unmapped artifact part is rejected outright. + * one rather than stacking adapters. Its calls run on `dispatcher`. An empty or unmapped artifact + * part is rejected outright. */ @JvmStatic - fun asKtArtifactService(service: JavaArtifactService): KtArtifactService = - javaArtifactServiceAsKt(service) + @JvmOverloads + fun asKtArtifactService( + service: JavaArtifactService, + dispatcher: CoroutineDispatcher = InteropDispatcher, + ): KtArtifactService = javaArtifactServiceAsKt(service, dispatcher) /** * Adapts an ADK Java memory service for the Kotlin engine, unwrapping a round-tripped Kotlin one - * rather than stacking adapters. + * rather than stacking adapters. Its calls run on `dispatcher`. */ @JvmStatic - fun asKtMemoryService(service: JavaMemoryService): KtMemoryService = - javaMemoryServiceAsKt(service) + @JvmOverloads + fun asKtMemoryService( + service: JavaMemoryService, + dispatcher: CoroutineDispatcher = InteropDispatcher, + ): KtMemoryService = javaMemoryServiceAsKt(service, dispatcher) } diff --git a/tokt/src/main/kotlin/com/google/adk/tokt/KotlinAdkToJava.kt b/tokt/src/main/kotlin/com/google/adk/tokt/KotlinAdkToJava.kt index 3574da1d1..9f0740924 100644 --- a/tokt/src/main/kotlin/com/google/adk/tokt/KotlinAdkToJava.kt +++ b/tokt/src/main/kotlin/com/google/adk/tokt/KotlinAdkToJava.kt @@ -18,6 +18,7 @@ package com.google.adk.tokt import com.google.adk.kt.runners.Runner as KtRunner import com.google.adk.runner.Runner as JavaRunner +import kotlinx.coroutines.CoroutineDispatcher /** * Reverse interop entry point: exposes an ADK Kotlin-engine [KtRunner] through the ADK Java @@ -28,7 +29,16 @@ import com.google.adk.runner.Runner as JavaRunner */ object KotlinAdkToJava { - /** Exposes a Kotlin-engine [runner] as an ADK Java [JavaRunner]. */ + /** + * Exposes a Kotlin-engine [runner] as an ADK Java [JavaRunner]. Its reverse service adapters + * bridge the Java RxJava calls onto the Kotlin engine via `dispatcher` (default + * `Dispatchers.IO`), which must be able to run nested bridged calls concurrently, so a + * single-threaded or tightly bounded dispatcher can deadlock. + */ @JvmStatic - fun asJavaRunner(runner: KtRunner): JavaRunner = KtRunnerToJava(runner) + @JvmOverloads + fun asJavaRunner( + runner: KtRunner, + dispatcher: CoroutineDispatcher = InteropDispatcher, + ): JavaRunner = KtRunnerToJava(runner, dispatcher) } diff --git a/tokt/src/main/kotlin/com/google/adk/tokt/KtRunnerToJava.kt b/tokt/src/main/kotlin/com/google/adk/tokt/KtRunnerToJava.kt index c5c5afb6d..c141d184e 100644 --- a/tokt/src/main/kotlin/com/google/adk/tokt/KtRunnerToJava.kt +++ b/tokt/src/main/kotlin/com/google/adk/tokt/KtRunnerToJava.kt @@ -38,6 +38,7 @@ import com.google.genai.types.Content as GenaiContent import io.reactivex.rxjava3.core.Completable import io.reactivex.rxjava3.core.Flowable import java.util.Optional +import kotlinx.coroutines.CoroutineDispatcher import kotlinx.coroutines.flow.map import kotlinx.coroutines.rx3.asFlowable @@ -50,20 +51,22 @@ import kotlinx.coroutines.rx3.asFlowable * bridged, so [runLive] returns a failed stream. [agent], [sessionService], [memoryService], * [artifactService] and [pluginManager] read the Kotlin runner's own components back through the * reverse adapters ([memoryService] / [artifactService] return `null` when absent; [pluginManager] - * is read-only and throws on registration). + * is read-only and throws on registration). The reverse service adapters bridge Java RxJava calls + * onto the Kotlin engine via `dispatcher`. */ // Subclassing Runner via its @Deprecated 8-arg super-constructor is intended here. @Suppress("DEPRECATION") -internal class KtRunnerToJava(private val ktRunner: KtRunner) : +internal class KtRunnerToJava(private val ktRunner: KtRunner, dispatcher: CoroutineDispatcher) : JavaRunner( // A Java view of the Kotlin runner's agent so agent() reads back; never run (see runAsync). ktAgentAsJava(ktRunner.agent), ktRunner.appName, // Non-null for the Java Runner field; artifactService() returns this bridge when present and // null when the Kotlin runner has none (the run then uses no artifact service). - ktRunner.artifactService?.let { ktArtifactServiceAsJava(it) } ?: JavaInMemoryArtifactService(), - ktSessionServiceAsJava(ktRunner.sessionService), - ktRunner.memoryService?.let { ktMemoryServiceAsJava(it) }, + ktRunner.artifactService?.let { ktArtifactServiceAsJava(it, dispatcher) } + ?: JavaInMemoryArtifactService(), + ktSessionServiceAsJava(ktRunner.sessionService, dispatcher), + ktRunner.memoryService?.let { ktMemoryServiceAsJava(it, dispatcher) }, emptyList(), null, null, diff --git a/tokt/src/main/kotlin/com/google/adk/tokt/adapters/JavaModelToKt.kt b/tokt/src/main/kotlin/com/google/adk/tokt/adapters/JavaModelToKt.kt index 497eb7923..64977285a 100644 --- a/tokt/src/main/kotlin/com/google/adk/tokt/adapters/JavaModelToKt.kt +++ b/tokt/src/main/kotlin/com/google/adk/tokt/adapters/JavaModelToKt.kt @@ -20,9 +20,9 @@ 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.models.BaseLlm as JavaBaseLlm -import com.google.adk.tokt.InteropDispatcher import com.google.adk.tokt.codecs.LlmRequestCodec import com.google.adk.tokt.codecs.LlmResponseCodec +import kotlinx.coroutines.CoroutineDispatcher import kotlinx.coroutines.flow.Flow import kotlinx.coroutines.flow.emitAll import kotlinx.coroutines.flow.flow @@ -37,15 +37,18 @@ import kotlinx.coroutines.reactive.asFlow * * The Kotlin `LlmRequest` is converted to a Java `LlmRequest` ([LlmRequestCodec]), the Java model's * RxJava `Flowable` is consumed as a coroutine [Flow] (via - * kotlinx-coroutines-reactive) and each response is converted back ([LlmResponseCodec]). + * kotlinx-coroutines-reactive) on `dispatcher` and each response is converted back + * ([LlmResponseCodec]). */ -internal class JavaModelToKt(private val javaLlm: JavaBaseLlm) : Model { +internal class JavaModelToKt( + private val javaLlm: JavaBaseLlm, + private val dispatcher: CoroutineDispatcher, +) : Model { override val name: String = javaLlm.model() - // Deferred into `flow {}` and dispatched on IO: the Java model's request build + generation run - // off the engine dispatcher (RxJava is synchronous by default), and a synchronous throw is routed - // through the Flow's error channel rather than escaping at collection time. + // Deferred into `flow {}` and run on dispatcher so the Java model's synchronous RxJava generation + // (and any synchronous throw, routed through the Flow's error channel) stays off the engine loop. override fun generateContent(request: LlmRequest, stream: Boolean): Flow = flow { emitAll( @@ -54,5 +57,5 @@ internal class JavaModelToKt(private val javaLlm: JavaBaseLlm) : Model { } ) } - .flowOn(InteropDispatcher) + .flowOn(dispatcher) } 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 93af75452..5f1b09d95 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 @@ -28,7 +28,6 @@ import com.google.adk.kt.tools.BaseTool as KtBaseTool import com.google.adk.kt.tools.ToolContext as KtToolContext import com.google.adk.kt.types.Content as KtContent import com.google.adk.plugins.Plugin as JavaPlugin -import com.google.adk.tokt.InteropDispatcher import com.google.adk.tokt.codecs.ContentCodec import com.google.adk.tokt.codecs.EventCodec import com.google.adk.tokt.codecs.LlmRequestCodec @@ -39,6 +38,7 @@ import com.google.adk.tokt.context.ktCallbackContextToJava import com.google.adk.tokt.context.ktToolContextToJava import io.reactivex.rxjava3.core.Completable import io.reactivex.rxjava3.core.Maybe +import kotlinx.coroutines.CoroutineDispatcher import kotlinx.coroutines.rx3.await import kotlinx.coroutines.rx3.awaitSingleOrNull import kotlinx.coroutines.withContext @@ -51,22 +51,25 @@ import kotlinx.coroutines.withContext * 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 { +internal class JavaPluginToKt( + internal val plugin: JavaPlugin, + private val dispatcher: CoroutineDispatcher, +) : KtPlugin { override val name: String get() = plugin.name /** - * Runs a Java plugin callback off the engine dispatcher. Plugin callbacks may do blocking I/O - * (logging, metrics, network); RxJava is synchronous by default, so running them on the coroutine - * that drives the agent loop could stall it. + * Runs a Java plugin callback off the engine dispatcher (`dispatcher`). Plugin callbacks may do + * blocking I/O (logging, metrics, network); RxJava is synchronous by default, so running them on + * the coroutine that drives the agent loop could stall it. */ - private suspend fun onIo(source: () -> Maybe): T? = - withContext(InteropDispatcher) { source().awaitSingleOrNull() } + private suspend fun onDispatcher(source: () -> Maybe): T? = + withContext(dispatcher) { source().awaitSingleOrNull() } - /** As [onIo], for the two callbacks that report completion rather than a value. */ - private suspend fun completeOnIo(source: () -> Completable) { - withContext(InteropDispatcher) { source().await() } + /** As [onDispatcher], for the two callbacks that report completion rather than a value. */ + private suspend fun completeOnDispatcher(source: () -> Completable) { + withContext(dispatcher) { source().await() } } // Run-level callbacks. @@ -75,8 +78,8 @@ internal class JavaPluginToKt(internal val plugin: JavaPlugin) : KtPlugin { invocationContext: KtInvocationContext, userMessage: KtContent, ): KtContent { - val javaContext = KtInvocationContextToJavaView(invocationContext) - val replacement = onIo { + val javaContext = KtInvocationContextToJavaView(invocationContext, dispatcher) + val replacement = onDispatcher { plugin.onUserMessageCallback(javaContext, ContentCodec.toJava(userMessage)) } return replacement?.let { ContentCodec.fromJava(it) } ?: userMessage @@ -85,26 +88,26 @@ internal class JavaPluginToKt(internal val plugin: JavaPlugin) : KtPlugin { override suspend fun beforeRun( invocationContext: KtInvocationContext ): CallbackChoice { - val javaContext = KtInvocationContextToJavaView(invocationContext) - val halt = onIo { plugin.beforeRunCallback(javaContext) } + val javaContext = KtInvocationContextToJavaView(invocationContext, dispatcher) + val halt = onDispatcher { plugin.beforeRunCallback(javaContext) } return if (halt != null) CallbackChoice.Break(ContentCodec.fromJava(halt)) else CallbackChoice.Continue(Unit) } override suspend fun onEvent(invocationContext: KtInvocationContext, event: KtEvent): KtEvent { - val javaContext = KtInvocationContextToJavaView(invocationContext) - val replacement = onIo { plugin.onEventCallback(javaContext, EventCodec.toJava(event)) } + val javaContext = KtInvocationContextToJavaView(invocationContext, dispatcher) + val replacement = onDispatcher { plugin.onEventCallback(javaContext, EventCodec.toJava(event)) } return replacement?.let { EventCodec.fromJava(it) } ?: event } override suspend fun afterRun(invocationContext: KtInvocationContext) { - val javaContext = KtInvocationContextToJavaView(invocationContext) - completeOnIo { plugin.afterRunCallback(javaContext) } + val javaContext = KtInvocationContextToJavaView(invocationContext, dispatcher) + completeOnDispatcher { plugin.afterRunCallback(javaContext) } } override suspend fun onRunError(invocationContext: KtInvocationContext, error: Throwable) { - val javaContext = KtInvocationContextToJavaView(invocationContext) - completeOnIo { plugin.onRunErrorCallback(javaContext, error) } + val javaContext = KtInvocationContextToJavaView(invocationContext, dispatcher) + completeOnDispatcher { plugin.onRunErrorCallback(javaContext, error) } } // Agent-level callbacks. @@ -113,8 +116,8 @@ internal class JavaPluginToKt(internal val plugin: JavaPlugin) : KtPlugin { context: KtCallbackContext ): CallbackChoice { val javaAgent = javaAgentView(context) - val javaContext = ktCallbackContextToJava(context, javaAgent) - val override = onIo { plugin.beforeAgentCallback(javaAgent, javaContext) } + val javaContext = ktCallbackContextToJava(context, javaAgent, dispatcher) + val override = onDispatcher { plugin.beforeAgentCallback(javaAgent, javaContext) } reconcileActionsToKt(javaContext.eventActions(), context.eventActions) return if (override != null) CallbackChoice.Break(ContentCodec.fromJava(override)) else CallbackChoice.Continue(KtEventActions()) @@ -122,8 +125,8 @@ internal class JavaPluginToKt(internal val plugin: JavaPlugin) : KtPlugin { override suspend fun afterAgent(context: KtCallbackContext): CallbackChoice { val javaAgent = javaAgentView(context) - val javaContext = ktCallbackContextToJava(context, javaAgent) - val override = onIo { plugin.afterAgentCallback(javaAgent, javaContext) } + val javaContext = ktCallbackContextToJava(context, javaAgent, dispatcher) + val override = onDispatcher { plugin.afterAgentCallback(javaAgent, javaContext) } reconcileActionsToKt(javaContext.eventActions(), context.eventActions) return if (override != null) CallbackChoice.Break(ContentCodec.fromJava(override)) else CallbackChoice.Continue(Unit) @@ -135,9 +138,9 @@ internal class JavaPluginToKt(internal val plugin: JavaPlugin) : KtPlugin { context: KtCallbackContext, request: KtLlmRequest, ): CallbackChoice { - val javaContext = ktCallbackContextToJava(context, javaAgentView(context)) + val javaContext = ktCallbackContextToJava(context, javaAgentView(context), dispatcher) val builder = LlmRequestCodec.toJava(request).toBuilder() - val override = onIo { plugin.beforeModelCallback(javaContext, builder) } + val override = onDispatcher { plugin.beforeModelCallback(javaContext, builder) } reconcileActionsToKt(javaContext.eventActions(), context.eventActions) return if (override != null) CallbackChoice.Break(LlmResponseCodec.fromJava(override)) // Re-apply onto the original request to preserve Kotlin-only fields (toolsDict, cache config); @@ -149,8 +152,8 @@ internal class JavaPluginToKt(internal val plugin: JavaPlugin) : KtPlugin { context: KtCallbackContext, response: KtLlmResponse, ): KtLlmResponse { - val javaContext = ktCallbackContextToJava(context, javaAgentView(context)) - val override = onIo { + val javaContext = ktCallbackContextToJava(context, javaAgentView(context), dispatcher) + val override = onDispatcher { plugin.afterModelCallback(javaContext, LlmResponseCodec.toJava(response)) } reconcileActionsToKt(javaContext.eventActions(), context.eventActions) @@ -162,9 +165,9 @@ internal class JavaPluginToKt(internal val plugin: JavaPlugin) : KtPlugin { request: KtLlmRequest, error: Throwable, ): CallbackChoice { - val javaContext = ktCallbackContextToJava(context, javaAgentView(context)) + val javaContext = ktCallbackContextToJava(context, javaAgentView(context), dispatcher) val builder = LlmRequestCodec.toJava(request).toBuilder() - val fallback = onIo { plugin.onModelErrorCallback(javaContext, builder, error) } + val fallback = onDispatcher { plugin.onModelErrorCallback(javaContext, builder, error) } reconcileActionsToKt(javaContext.eventActions(), context.eventActions) return if (fallback != null) CallbackChoice.Break(LlmResponseCodec.fromJava(fallback)) else CallbackChoice.Continue(Unit) @@ -179,9 +182,9 @@ internal class JavaPluginToKt(internal val plugin: JavaPlugin) : KtPlugin { args: Map, ): CallbackChoice, Map> { val javaTool = ktToolAsJava(tool) - val javaContext = ktToolContextToJava(context) + val javaContext = ktToolContextToJava(context, dispatcher) val mutableArgs = args.toMutableMap() - val override = onIo { plugin.beforeToolCallback(javaTool, mutableArgs, javaContext) } + val override = onDispatcher { plugin.beforeToolCallback(javaTool, mutableArgs, javaContext) } reconcileActionsToKt(javaContext.actions(), context.actions) return if (override != null) CallbackChoice.Break(override) else CallbackChoice.Continue(mutableArgs) @@ -194,8 +197,8 @@ internal class JavaPluginToKt(internal val plugin: JavaPlugin) : KtPlugin { result: Map, ): Map { val javaTool = ktToolAsJava(tool) - val javaContext = ktToolContextToJava(context) - val override = onIo { plugin.afterToolCallback(javaTool, args, javaContext, result) } + val javaContext = ktToolContextToJava(context, dispatcher) + val override = onDispatcher { plugin.afterToolCallback(javaTool, args, javaContext, result) } reconcileActionsToKt(javaContext.actions(), context.actions) return override ?: result } @@ -207,8 +210,8 @@ internal class JavaPluginToKt(internal val plugin: JavaPlugin) : KtPlugin { error: Throwable, ): CallbackChoice> { val javaTool = ktToolAsJava(tool) - val javaContext = ktToolContextToJava(context) - val fallback = onIo { plugin.onToolErrorCallback(javaTool, args, javaContext, error) } + val javaContext = ktToolContextToJava(context, dispatcher) + val fallback = onDispatcher { plugin.onToolErrorCallback(javaTool, args, javaContext, error) } reconcileActionsToKt(javaContext.actions(), context.actions) return if (fallback != null) CallbackChoice.Break(fallback) else CallbackChoice.Continue(Unit) } diff --git a/tokt/src/main/kotlin/com/google/adk/tokt/adapters/JavaToolToKt.kt b/tokt/src/main/kotlin/com/google/adk/tokt/adapters/JavaToolToKt.kt index c4d2f1d22..e7bdc6edf 100644 --- a/tokt/src/main/kotlin/com/google/adk/tokt/adapters/JavaToolToKt.kt +++ b/tokt/src/main/kotlin/com/google/adk/tokt/adapters/JavaToolToKt.kt @@ -20,22 +20,25 @@ import com.google.adk.kt.models.LlmRequest as KtLlmRequest import com.google.adk.kt.tools.BaseTool import com.google.adk.kt.tools.ToolContext import com.google.adk.kt.types.FunctionDeclaration -import com.google.adk.tokt.InteropDispatcher import com.google.adk.tokt.codecs.FunctionDeclarationCodec import com.google.adk.tokt.context.ktToolContextToJava import com.google.adk.tools.BaseTool as JavaBaseTool import kotlin.jvm.optionals.getOrNull +import kotlinx.coroutines.CoroutineDispatcher import kotlinx.coroutines.reactive.awaitFirstOrNull import kotlinx.coroutines.reactive.awaitSingle import kotlinx.coroutines.withContext /** * Exposes a Java [JavaBaseTool] as a Kotlin [BaseTool] so the Kotlin engine can invoke it. [run] - * awaits the tool's RxJava result off the engine dispatcher, with actions delegating live to the - * Kotlin context. [processLlmRequest] re-applies a tool's request edits while preserving - * Kotlin-only fields (model, toolsDict). + * awaits the tool's RxJava result off the engine dispatcher (`dispatcher`), with actions delegating + * live to the Kotlin context. [processLlmRequest] re-applies a tool's request edits while + * preserving Kotlin-only fields (model, toolsDict). */ -internal class JavaToolToKt(internal val javaTool: JavaBaseTool) : +internal class JavaToolToKt( + internal val javaTool: JavaBaseTool, + private val dispatcher: CoroutineDispatcher, +) : BaseTool( javaTool.name(), javaTool.description(), @@ -47,13 +50,11 @@ internal class JavaToolToKt(internal val javaTool: JavaBaseTool) : javaTool.declaration().map { FunctionDeclarationCodec.fromJava(it) }.getOrNull() override suspend fun run(context: ToolContext, args: Map): Any { - val javaContext = ktToolContextToJava(context) + val javaContext = ktToolContextToJava(context, dispatcher) // Run off the engine dispatcher so synchronous RxJava or blocking tool I/O cannot stall the // agent loop; the live actions view is backed by concurrent maps, so another thread is safe. val result = - withContext(InteropDispatcher) { - javaTool.runAsync(args, javaContext).toFlowable().awaitSingle() - } + withContext(dispatcher) { javaTool.runAsync(args, javaContext).toFlowable().awaitSingle() } // Carry back what the live view cannot (a setActions replacement, confirmations, skip-summary). reconcileActionsToKt(javaContext.actions(), context.actions) return result @@ -63,9 +64,9 @@ internal class JavaToolToKt(internal val javaTool: JavaBaseTool) : toolContext: ToolContext, llmRequest: KtLlmRequest, ): KtLlmRequest { - val javaToolContext = ktToolContextToJava(toolContext) + val javaToolContext = ktToolContextToJava(toolContext, dispatcher) return bridgeProcessLlmRequest(llmRequest) { builder -> - withContext(InteropDispatcher) { + withContext(dispatcher) { javaTool.processLlmRequest(builder, javaToolContext).toFlowable().awaitFirstOrNull() } } diff --git a/tokt/src/main/kotlin/com/google/adk/tokt/adapters/JavaToolsetToKt.kt b/tokt/src/main/kotlin/com/google/adk/tokt/adapters/JavaToolsetToKt.kt index 7d25c915d..c7d58e96d 100644 --- a/tokt/src/main/kotlin/com/google/adk/tokt/adapters/JavaToolsetToKt.kt +++ b/tokt/src/main/kotlin/com/google/adk/tokt/adapters/JavaToolsetToKt.kt @@ -21,10 +21,10 @@ import com.google.adk.kt.models.LlmRequest as KtLlmRequest 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.tokt.InteropDispatcher import com.google.adk.tokt.context.KtReadonlyContextToJavaView import com.google.adk.tokt.context.ktToolContextToJava import com.google.adk.tools.BaseToolset as JavaBaseToolset +import kotlinx.coroutines.CoroutineDispatcher import kotlinx.coroutines.reactive.awaitFirstOrNull import kotlinx.coroutines.reactive.awaitSingle import kotlinx.coroutines.withContext @@ -33,13 +33,16 @@ import kotlinx.coroutines.withContext * Java -> Kotlin adapter: exposes a Java [JavaBaseToolset] as a Kotlin [Toolset] so the ADK Kotlin * can pull tools from a user-authored Java toolset at request time, with the live [ReadonlyContext] * ([KtReadonlyContextToJavaView]). Each Java tool the toolset returns is wrapped in a - * [JavaToolToKt]. + * [JavaToolToKt] on the same `dispatcher`. * * Both `getTools` (tool provisioning) and `processLlmRequest` are bridged; the latter lets the Java * toolset mutate the request's contents/config, which are re-applied to the Kotlin request while * preserving Kotlin-only fields (model, toolsDict). */ -internal class JavaToolsetToKt(internal val javaToolset: JavaBaseToolset) : Toolset { +internal class JavaToolsetToKt( + internal val javaToolset: JavaBaseToolset, + private val dispatcher: CoroutineDispatcher, +) : Toolset { override suspend fun getTools(readonlyContext: ReadonlyContext?): List { // ADK Java's BaseToolset.getTools requires a non-null ReadonlyContext (and the engine always @@ -53,19 +56,19 @@ internal class JavaToolsetToKt(internal val javaToolset: JavaBaseToolset) : Tool // Off the engine dispatcher: toolset discovery (getTools) often does blocking I/O. // RxJava Flowable -> Single -> Flowable -> awaitSingle(). val javaTools = - withContext(InteropDispatcher) { + withContext(dispatcher) { javaToolset.getTools(javaContext).toList().toFlowable().awaitSingle() } - return javaTools.map { JavaToolToKt(it) } + return javaTools.map { JavaToolToKt(it, dispatcher) } } override suspend fun processLlmRequest( toolContext: ToolContext, llmRequest: KtLlmRequest, ): KtLlmRequest { - val javaToolContext = ktToolContextToJava(toolContext) + val javaToolContext = ktToolContextToJava(toolContext, dispatcher) return bridgeProcessLlmRequest(llmRequest) { builder -> - withContext(InteropDispatcher) { + withContext(dispatcher) { javaToolset.processLlmRequest(builder, javaToolContext).toFlowable().awaitFirstOrNull() } } diff --git a/tokt/src/main/kotlin/com/google/adk/tokt/context/KtCallbackContextToJava.kt b/tokt/src/main/kotlin/com/google/adk/tokt/context/KtCallbackContextToJava.kt index f13f39fad..23beb628d 100644 --- a/tokt/src/main/kotlin/com/google/adk/tokt/context/KtCallbackContextToJava.kt +++ b/tokt/src/main/kotlin/com/google/adk/tokt/context/KtCallbackContextToJava.kt @@ -30,6 +30,7 @@ import com.google.adk.tokt.codecs.ktSessionToJavaLive import com.google.adk.tokt.services.ktArtifactServiceAsJava import com.google.adk.tokt.services.ktMemoryServiceAsJava import java.util.Optional +import kotlinx.coroutines.CoroutineDispatcher /** * A Java [JavaInvocationContext] backed by a Kotlin [KtCallbackContext]'s readonly view. The @@ -40,7 +41,8 @@ import java.util.Optional private class KtCallbackContextToJavaView( private val context: KtCallbackContext, private val agent: JavaBaseAgent, -) : JavaInvocationContext(callbackContextJavaBuilder(context)) { + dispatcher: CoroutineDispatcher, +) : JavaInvocationContext(callbackContextJavaBuilder(context, dispatcher)) { override fun invocationId(): String = context.invocationId @@ -80,12 +82,16 @@ private class KtCallbackContextToJavaView( /** * Builds the base builder for [KtCallbackContextToJavaView], wiring the artifact / memory services - * (unwrapped where possible) and user content the readonly [KtCallbackContext] exposes. A Kotlin - * callback context does not expose its session service, so the view throws for it rather than - * leaving the field null: this builder bypasses `build()`, so `validate()` never runs and the unset - * field would otherwise surface as an NPE at the plugin's own call site. + * (unwrapped where possible, bridged on `dispatcher`) and user content the readonly + * [KtCallbackContext] exposes. A Kotlin callback context does not expose its session service, so + * the view throws for it rather than leaving the field null: this builder bypasses `build()`, so + * `validate()` never runs and the unset field would otherwise surface as an NPE at the plugin's own + * call site. */ -private fun callbackContextJavaBuilder(context: KtCallbackContext): JavaInvocationContext.Builder { +private fun callbackContextJavaBuilder( + context: KtCallbackContext, + dispatcher: CoroutineDispatcher, +): JavaInvocationContext.Builder { val builder = JavaInvocationContext.builder() .invocationId(context.invocationId) @@ -93,8 +99,8 @@ private fun callbackContextJavaBuilder(context: KtCallbackContext): JavaInvocati .session(ktSessionToJavaLive(context.session)) // Carry the run's RunConfig so a callback reading it does not see a default. context.runConfig?.let { builder.runConfig(RunConfigCodec.toJava(it)) } - context.artifactService?.let { builder.artifactService(ktArtifactServiceAsJava(it)) } - context.memoryService?.let { builder.memoryService(ktMemoryServiceAsJava(it)) } + context.artifactService?.let { builder.artifactService(ktArtifactServiceAsJava(it, dispatcher)) } + context.memoryService?.let { builder.memoryService(ktMemoryServiceAsJava(it, dispatcher)) } context.userContent?.let { builder.userContent(ContentCodec.toJava(it)) } return builder } @@ -102,13 +108,15 @@ private fun callbackContextJavaBuilder(context: KtCallbackContext): JavaInvocati /** * Presents the Kotlin [context] as a Java [JavaCallbackContext] (reusing [KtEventActionsToJavaView] * so a callback's state / artifact deltas write straight through to the Kotlin context). [agent] is - * the Java agent the callback belongs to, so a callback sees the correct `context.agent()`. + * the Java agent the callback belongs to, so a callback sees the correct `context.agent()`. Bridged + * service calls run on `dispatcher`. */ internal fun ktCallbackContextToJava( context: KtCallbackContext, agent: JavaBaseAgent, + dispatcher: CoroutineDispatcher, ): JavaCallbackContext = JavaCallbackContext( - KtCallbackContextToJavaView(context, agent), + KtCallbackContextToJavaView(context, agent, dispatcher), KtEventActionsToJavaView(context.eventActions), ) diff --git a/tokt/src/main/kotlin/com/google/adk/tokt/context/KtToJava.kt b/tokt/src/main/kotlin/com/google/adk/tokt/context/KtToJava.kt index 12c1859c6..7b9718e5c 100644 --- a/tokt/src/main/kotlin/com/google/adk/tokt/context/KtToJava.kt +++ b/tokt/src/main/kotlin/com/google/adk/tokt/context/KtToJava.kt @@ -44,6 +44,7 @@ import com.google.genai.types.Content as GenaiContent import java.util.Collections import java.util.Optional import java.util.concurrent.ConcurrentMap +import kotlinx.coroutines.CoroutineDispatcher /** * Kt -> Java context views used when the ADK Kotlin calls back into ADK Java code (a Java tool) or @@ -119,10 +120,14 @@ internal fun javaAgentView(ic: KtInvocationContext): JavaBaseAgent = /** * Builds the base builder for [KtInvocationContextToJavaView], wiring the Kotlin services as Java - * views (unwrapped to the original Java service where possible to avoid extra hops). + * views (unwrapped to the original Java service where possible to avoid extra hops). A wrapped + * service bridges its calls on `dispatcher`. */ @OptIn(FrameworkInternalApi::class) -private fun ktInvocationContextJavaBuilder(ic: KtInvocationContext): JavaInvocationContext.Builder { +private fun ktInvocationContextJavaBuilder( + ic: KtInvocationContext, + dispatcher: CoroutineDispatcher, +): JavaInvocationContext.Builder { val builder = JavaInvocationContext.builder() .invocationId(ic.invocationId) @@ -132,9 +137,9 @@ private fun ktInvocationContextJavaBuilder(ic: KtInvocationContext): JavaInvocat .agent(javaAgentView(ic.agent, ic.frameworkData.callbackContextData)) // Carry the run's RunConfig so a component does not fall back to a default streaming mode. ic.runConfig?.let { builder.runConfig(RunConfigCodec.toJava(it)) } - ic.sessionService?.let { builder.sessionService(ktSessionServiceAsJava(it)) } - ic.artifactService?.let { builder.artifactService(ktArtifactServiceAsJava(it)) } - ic.memoryService?.let { builder.memoryService(ktMemoryServiceAsJava(it)) } + ic.sessionService?.let { builder.sessionService(ktSessionServiceAsJava(it, dispatcher)) } + ic.artifactService?.let { builder.artifactService(ktArtifactServiceAsJava(it, dispatcher)) } + ic.memoryService?.let { builder.memoryService(ktMemoryServiceAsJava(it, dispatcher)) } // ReadonlyContext.userContent() delegates here, so tools and plugins read the turn's content. ic.userContent?.let { builder.userContent(ContentCodec.toJava(it)) } return builder @@ -144,10 +149,12 @@ private fun ktInvocationContextJavaBuilder(ic: KtInvocationContext): JavaInvocat * A Java [JavaInvocationContext] backed live by a Kotlin [KtInvocationContext]: it subclasses the * Java type so casts keep working, and its accessors read and write through to the Kotlin context. * Tools and plugin run-level callbacks share it, so a Java component sees the same context wherever - * it runs. + * it runs. Bridged service calls run on `dispatcher`. */ -internal class KtInvocationContextToJavaView(private val ic: KtInvocationContext) : - JavaInvocationContext(ktInvocationContextJavaBuilder(ic)) { +internal class KtInvocationContextToJavaView( + private val ic: KtInvocationContext, + dispatcher: CoroutineDispatcher, +) : JavaInvocationContext(ktInvocationContextJavaBuilder(ic, dispatcher)) { override fun invocationId(): String = ic.invocationId @@ -175,11 +182,15 @@ internal class KtInvocationContextToJavaView(private val ic: KtInvocationContext /** * Converts a Kotlin [KtToolContext] to a Java [JavaToolContext]. The tool's side effects stay live - * because the invocation context and actions delegate to the Kotlin context by reference. + * because the invocation context and actions delegate to the Kotlin context by reference. Bridged + * service calls run on `dispatcher`. */ -internal fun ktToolContextToJava(context: KtToolContext): JavaToolContext { +internal fun ktToolContextToJava( + context: KtToolContext, + dispatcher: CoroutineDispatcher, +): JavaToolContext { val builder = - JavaToolContext.builder(KtInvocationContextToJavaView(context.invocationContext)) + JavaToolContext.builder(KtInvocationContextToJavaView(context.invocationContext, dispatcher)) .actions(KtEventActionsToJavaView(context.actions)) .functionCallId(context.functionCallId) .eventId(context.eventId) diff --git a/tokt/src/main/kotlin/com/google/adk/tokt/services/JavaArtifactServiceToKt.kt b/tokt/src/main/kotlin/com/google/adk/tokt/services/JavaArtifactServiceToKt.kt index 988f3b053..f83308ad9 100644 --- a/tokt/src/main/kotlin/com/google/adk/tokt/services/JavaArtifactServiceToKt.kt +++ b/tokt/src/main/kotlin/com/google/adk/tokt/services/JavaArtifactServiceToKt.kt @@ -20,9 +20,9 @@ import com.google.adk.artifacts.BaseArtifactService as JavaBaseArtifactService import com.google.adk.kt.artifacts.ArtifactService as KtArtifactService import com.google.adk.kt.sessions.SessionKey import com.google.adk.kt.types.Part as KtPart -import com.google.adk.tokt.InteropDispatcher import com.google.adk.tokt.codecs.PartCodec import com.google.adk.tokt.codecs.sessionId +import kotlinx.coroutines.CoroutineDispatcher import kotlinx.coroutines.rx3.await import kotlinx.coroutines.rx3.awaitSingleOrNull import kotlinx.coroutines.withContext @@ -31,17 +31,19 @@ import kotlinx.coroutines.withContext * A Kotlin [KtArtifactService] backed by an ADK Java [JavaBaseArtifactService] - the reverse of * [KtArtifactServiceToJava] - so the Kotlin runner can drive a Java app's own artifact service when * a Java `Runner` runs on the Kotlin engine. Parts are converted with [PartCodec]. All Java-service - * calls run on InteropDispatcher since a user's artifact service may block on subscribe. + * calls run on `dispatcher` since a user's artifact service may block on subscribe. */ -internal class JavaArtifactServiceToKt(internal val service: JavaBaseArtifactService) : - KtArtifactService { +internal class JavaArtifactServiceToKt( + internal val service: JavaBaseArtifactService, + private val dispatcher: CoroutineDispatcher, +) : KtArtifactService { override suspend fun saveArtifact( sessionKey: SessionKey, filename: String, artifact: KtPart, ): Int = - withContext(InteropDispatcher) { + withContext(dispatcher) { service .saveArtifact( sessionKey.appName, @@ -58,7 +60,7 @@ internal class JavaArtifactServiceToKt(internal val service: JavaBaseArtifactSer filename: String, artifact: KtPart, ): KtPart = - withContext(InteropDispatcher) { + withContext(dispatcher) { PartCodec.fromJavaOrThrow( service .saveAndReloadArtifact( @@ -77,7 +79,7 @@ internal class JavaArtifactServiceToKt(internal val service: JavaBaseArtifactSer filename: String, version: Int?, ): KtPart? = - withContext(InteropDispatcher) { + withContext(dispatcher) { service .loadArtifact( sessionKey.appName, @@ -93,7 +95,7 @@ internal class JavaArtifactServiceToKt(internal val service: JavaBaseArtifactSer } override suspend fun listArtifactKeys(sessionKey: SessionKey): List = - withContext(InteropDispatcher) { + withContext(dispatcher) { service .listArtifactKeys(sessionKey.appName, sessionKey.userId, sessionId(sessionKey)) .await() @@ -101,7 +103,7 @@ internal class JavaArtifactServiceToKt(internal val service: JavaBaseArtifactSer } override suspend fun deleteArtifact(sessionKey: SessionKey, filename: String) { - withContext(InteropDispatcher) { + withContext(dispatcher) { service .deleteArtifact(sessionKey.appName, sessionKey.userId, sessionId(sessionKey), filename) .await() @@ -109,7 +111,7 @@ internal class JavaArtifactServiceToKt(internal val service: JavaBaseArtifactSer } override suspend fun listVersions(sessionKey: SessionKey, filename: String): List = - withContext(InteropDispatcher) { + withContext(dispatcher) { service .listVersions(sessionKey.appName, sessionKey.userId, sessionId(sessionKey), filename) .await() diff --git a/tokt/src/main/kotlin/com/google/adk/tokt/services/JavaMemoryServiceToKt.kt b/tokt/src/main/kotlin/com/google/adk/tokt/services/JavaMemoryServiceToKt.kt index 1fabe218e..5a4f07dea 100644 --- a/tokt/src/main/kotlin/com/google/adk/tokt/services/JavaMemoryServiceToKt.kt +++ b/tokt/src/main/kotlin/com/google/adk/tokt/services/JavaMemoryServiceToKt.kt @@ -20,23 +20,25 @@ import com.google.adk.kt.memory.MemoryService as KtMemoryService import com.google.adk.kt.memory.SearchMemoryResponse as KtSearchMemoryResponse import com.google.adk.kt.sessions.Session as KtSession import com.google.adk.memory.BaseMemoryService as JavaBaseMemoryService -import com.google.adk.tokt.InteropDispatcher import com.google.adk.tokt.codecs.MemoryEntryCodec import com.google.adk.tokt.codecs.ktSessionToJava +import kotlinx.coroutines.CoroutineDispatcher import kotlinx.coroutines.rx3.await import kotlinx.coroutines.withContext /** * A Kotlin [KtMemoryService] backed by an ADK Java [JavaBaseMemoryService] - the reverse of * [KtMemoryServiceToJava] - so the Kotlin runner can drive a Java app's own memory service when a - * Java `Runner` runs on the Kotlin engine. + * Java `Runner` runs on the Kotlin engine. All Java-service calls run on `dispatcher`. */ -internal class JavaMemoryServiceToKt(internal val service: JavaBaseMemoryService) : - KtMemoryService { +internal class JavaMemoryServiceToKt( + internal val service: JavaBaseMemoryService, + private val dispatcher: CoroutineDispatcher, +) : KtMemoryService { override suspend fun addSessionToMemory(session: KtSession) { - // On InteropDispatcher: a user's Java memory service may block on subscribe. - withContext(InteropDispatcher) { service.addSessionToMemory(ktSessionToJava(session)).await() } + // On dispatcher: a user's Java memory service may block on subscribe. + withContext(dispatcher) { service.addSessionToMemory(ktSessionToJava(session)).await() } } override suspend fun searchMemory( @@ -44,7 +46,7 @@ internal class JavaMemoryServiceToKt(internal val service: JavaBaseMemoryService userId: String, query: String, ): KtSearchMemoryResponse = - withContext(InteropDispatcher) { + withContext(dispatcher) { MemoryEntryCodec.fromJava(service.searchMemory(appName, userId, query).await()) } } diff --git a/tokt/src/main/kotlin/com/google/adk/tokt/services/JavaSessionServiceToKt.kt b/tokt/src/main/kotlin/com/google/adk/tokt/services/JavaSessionServiceToKt.kt index 1f7aa51b4..70665fac8 100644 --- a/tokt/src/main/kotlin/com/google/adk/tokt/services/JavaSessionServiceToKt.kt +++ b/tokt/src/main/kotlin/com/google/adk/tokt/services/JavaSessionServiceToKt.kt @@ -25,7 +25,6 @@ import com.google.adk.kt.sessions.SessionKey import com.google.adk.kt.sessions.SessionService as KtSessionService import com.google.adk.sessions.BaseSessionService as JavaBaseSessionService import com.google.adk.sessions.GetSessionConfig as JavaGetSessionConfig -import com.google.adk.tokt.InteropDispatcher import com.google.adk.tokt.codecs.EventCodec import com.google.adk.tokt.codecs.SessionCodec import com.google.adk.tokt.codecs.ktSessionToJava @@ -33,6 +32,7 @@ import com.google.adk.tokt.codecs.sessionId import java.util.Optional import kotlin.time.toJavaInstant import kotlin.time.toKotlinInstant +import kotlinx.coroutines.CoroutineDispatcher import kotlinx.coroutines.rx3.await import kotlinx.coroutines.rx3.awaitSingleOrNull import kotlinx.coroutines.withContext @@ -46,20 +46,22 @@ import kotlinx.coroutines.withContext * `appendEvent` persists through the Java service (keyed by appName/userId/id) and then keeps the * in-memory Kotlin [session] the runner holds in sync via the default implementation. */ -internal class JavaSessionServiceToKt(internal val service: JavaBaseSessionService) : - KtSessionService { +internal class JavaSessionServiceToKt( + internal val service: JavaBaseSessionService, + private val dispatcher: CoroutineDispatcher, +) : KtSessionService { - // All Java-service calls are awaited on InteropDispatcher: a user's Java service (in-memory, - // Vertex, custom) may be blocking-on-subscribe, so it must not run on the coroutine driving the - // agent loop. + // All Java-service calls are awaited on dispatcher: a user's Java service (in-memory, Vertex, + // custom) may be blocking-on-subscribe, so it must not run on the coroutine driving the agent + // loop. override suspend fun createSession(key: SessionKey, state: Map?): KtSession = - withContext(InteropDispatcher) { + withContext(dispatcher) { SessionCodec.fromJava(service.createSession(key.appName, key.userId, state, key.id).await()) } override suspend fun getSession(key: SessionKey, config: KtGetSessionConfig?): KtSession? = - withContext(InteropDispatcher) { + withContext(dispatcher) { service .getSession(key.appName, key.userId, sessionId(key), Optional.ofNullable(config?.toJava())) .awaitSingleOrNull() @@ -67,7 +69,7 @@ internal class JavaSessionServiceToKt(internal val service: JavaBaseSessionServi } override suspend fun listSessions(appName: String, userId: String): KtListSessionsResponse = - withContext(InteropDispatcher) { + withContext(dispatcher) { KtListSessionsResponse( sessions = service.listSessions(appName, userId).await().sessions().map { SessionCodec.fromJava(it) } @@ -75,11 +77,11 @@ internal class JavaSessionServiceToKt(internal val service: JavaBaseSessionServi } override suspend fun closeSession(session: KtSession) { - withContext(InteropDispatcher) { service.closeSession(ktSessionToJava(session)).await() } + withContext(dispatcher) { service.closeSession(ktSessionToJava(session)).await() } } override suspend fun deleteSession(key: SessionKey) { - withContext(InteropDispatcher) { + withContext(dispatcher) { service.deleteSession(key.appName, key.userId, sessionId(key)).await() } } @@ -88,7 +90,7 @@ internal class JavaSessionServiceToKt(internal val service: JavaBaseSessionServi // A well-behaved Java service always returns a response; tolerate a null Single/response (e.g. // an unstubbed test double) as "no events" rather than crashing the Kotlin run. val response = - withContext(InteropDispatcher) { + withContext(dispatcher) { service.listEvents(key.appName, key.userId, sessionId(key))?.await() } return KtListEventsResponse( @@ -101,9 +103,7 @@ internal class JavaSessionServiceToKt(internal val service: JavaBaseSessionServi // Persist through the Java service first: the converted Java session carries the prior events, // so the service appends and persists this event (keyed by appName/userId/id). val javaSession = ktSessionToJava(session) - withContext(InteropDispatcher) { - service.appendEvent(javaSession, EventCodec.toJava(event)).await() - } + withContext(dispatcher) { service.appendEvent(javaSession, EventCodec.toJava(event)).await() } // Keep the in-memory Kotlin session the runner holds in sync (state delta + event list); this // also advances lastUpdateTime to the event timestamp. val appended = super.appendEvent(session, event) diff --git a/tokt/src/main/kotlin/com/google/adk/tokt/services/KtArtifactServiceToJava.kt b/tokt/src/main/kotlin/com/google/adk/tokt/services/KtArtifactServiceToJava.kt index 99ba337af..6906ddb0d 100644 --- a/tokt/src/main/kotlin/com/google/adk/tokt/services/KtArtifactServiceToJava.kt +++ b/tokt/src/main/kotlin/com/google/adk/tokt/services/KtArtifactServiceToJava.kt @@ -20,13 +20,13 @@ import com.google.adk.artifacts.BaseArtifactService as JavaBaseArtifactService import com.google.adk.artifacts.ListArtifactsResponse as JavaListArtifactsResponse import com.google.adk.kt.artifacts.ArtifactService as KtArtifactService import com.google.adk.kt.sessions.SessionKey -import com.google.adk.tokt.InteropDispatcher import com.google.adk.tokt.codecs.PartCodec import com.google.common.collect.ImmutableList import com.google.genai.types.Part as GenaiPart import io.reactivex.rxjava3.core.Completable import io.reactivex.rxjava3.core.Maybe import io.reactivex.rxjava3.core.Single +import kotlinx.coroutines.CoroutineDispatcher import kotlinx.coroutines.rx3.rxCompletable import kotlinx.coroutines.rx3.rxMaybe import kotlinx.coroutines.rx3.rxSingle @@ -34,10 +34,12 @@ import kotlinx.coroutines.rx3.rxSingle /** * A Java [JavaBaseArtifactService] backed by a Kotlin [KtArtifactService] - the reverse of the Java * service wrappers - so a Java agent running under the Kotlin runner sees a Java artifact service - * backed by the Kotlin service. Artifacts are converted with [PartCodec]. + * backed by the Kotlin service, bridged on `dispatcher`. Artifacts are converted with [PartCodec]. */ -internal class KtArtifactServiceToJava(internal val service: KtArtifactService) : - JavaBaseArtifactService { +internal class KtArtifactServiceToJava( + internal val service: KtArtifactService, + private val dispatcher: CoroutineDispatcher, +) : JavaBaseArtifactService { override fun saveArtifact( appName: String, @@ -46,7 +48,7 @@ internal class KtArtifactServiceToJava(internal val service: KtArtifactService) filename: String, artifact: GenaiPart, ): Single = - rxSingle(InteropDispatcher) { + rxSingle(dispatcher) { service.saveArtifact( SessionKey(appName, userId, sessionId), filename, @@ -63,7 +65,7 @@ internal class KtArtifactServiceToJava(internal val service: KtArtifactService) filename: String, artifact: GenaiPart, ): Single = - rxSingle(InteropDispatcher) { + rxSingle(dispatcher) { PartCodec.toJavaOrThrow( service.saveAndReloadArtifact( SessionKey(appName, userId, sessionId), @@ -80,7 +82,7 @@ internal class KtArtifactServiceToJava(internal val service: KtArtifactService) filename: String, version: Int?, ): Maybe = - rxMaybe(InteropDispatcher) { + rxMaybe(dispatcher) { // toJavaOrThrow, not toJava: an artifact that exists but cannot be converted must be // rejected, not reported as absent. service.loadArtifact(SessionKey(appName, userId, sessionId), filename, version)?.let { @@ -96,7 +98,7 @@ internal class KtArtifactServiceToJava(internal val service: KtArtifactService) userId: String, sessionId: String, ): Single = - rxSingle(InteropDispatcher) { + rxSingle(dispatcher) { JavaListArtifactsResponse.builder() .filenames( ImmutableList.copyOf(service.listArtifactKeys(SessionKey(appName, userId, sessionId))) @@ -110,7 +112,7 @@ internal class KtArtifactServiceToJava(internal val service: KtArtifactService) sessionId: String, filename: String, ): Completable = - rxCompletable(InteropDispatcher) { + rxCompletable(dispatcher) { service.deleteArtifact(SessionKey(appName, userId, sessionId), filename) } @@ -121,7 +123,7 @@ internal class KtArtifactServiceToJava(internal val service: KtArtifactService) sessionId: String, filename: String, ): Single> = - rxSingle(InteropDispatcher) { + rxSingle(dispatcher) { ImmutableList.copyOf(service.listVersions(SessionKey(appName, userId, sessionId), filename)) } } diff --git a/tokt/src/main/kotlin/com/google/adk/tokt/services/KtMemoryServiceToJava.kt b/tokt/src/main/kotlin/com/google/adk/tokt/services/KtMemoryServiceToJava.kt index 7a0c139bc..8736f0183 100644 --- a/tokt/src/main/kotlin/com/google/adk/tokt/services/KtMemoryServiceToJava.kt +++ b/tokt/src/main/kotlin/com/google/adk/tokt/services/KtMemoryServiceToJava.kt @@ -20,31 +20,31 @@ import com.google.adk.kt.memory.MemoryService as KtMemoryService import com.google.adk.memory.BaseMemoryService as JavaBaseMemoryService import com.google.adk.memory.SearchMemoryResponse as JavaSearchMemoryResponse import com.google.adk.sessions.Session as JavaSession -import com.google.adk.tokt.InteropDispatcher import com.google.adk.tokt.codecs.MemoryEntryCodec import com.google.adk.tokt.codecs.SessionCodec import io.reactivex.rxjava3.core.Completable import io.reactivex.rxjava3.core.Single +import kotlinx.coroutines.CoroutineDispatcher import kotlinx.coroutines.rx3.rxCompletable import kotlinx.coroutines.rx3.rxSingle /** * A Java [JavaBaseMemoryService] backed by a Kotlin [KtMemoryService] - the reverse of the Java * service wrappers - so a Java agent running under the Kotlin runner sees a Java memory service - * backed by the Kotlin service. + * backed by the Kotlin service, bridged on `dispatcher`. */ -internal class KtMemoryServiceToJava(internal val service: KtMemoryService) : - JavaBaseMemoryService { +internal class KtMemoryServiceToJava( + internal val service: KtMemoryService, + private val dispatcher: CoroutineDispatcher, +) : JavaBaseMemoryService { override fun addSessionToMemory(session: JavaSession): Completable = - rxCompletable(InteropDispatcher) { service.addSessionToMemory(SessionCodec.fromJava(session)) } + rxCompletable(dispatcher) { service.addSessionToMemory(SessionCodec.fromJava(session)) } override fun searchMemory( appName: String, userId: String, query: String, ): Single = - rxSingle(InteropDispatcher) { - MemoryEntryCodec.toJava(service.searchMemory(appName, userId, query)) - } + rxSingle(dispatcher) { MemoryEntryCodec.toJava(service.searchMemory(appName, userId, query)) } } diff --git a/tokt/src/main/kotlin/com/google/adk/tokt/services/KtSessionServiceToJava.kt b/tokt/src/main/kotlin/com/google/adk/tokt/services/KtSessionServiceToJava.kt index 4b2bc7f6e..9aff981cc 100644 --- a/tokt/src/main/kotlin/com/google/adk/tokt/services/KtSessionServiceToJava.kt +++ b/tokt/src/main/kotlin/com/google/adk/tokt/services/KtSessionServiceToJava.kt @@ -25,7 +25,6 @@ import com.google.adk.sessions.GetSessionConfig as JavaGetSessionConfig import com.google.adk.sessions.ListEventsResponse as JavaListEventsResponse import com.google.adk.sessions.ListSessionsResponse as JavaListSessionsResponse import com.google.adk.sessions.Session as JavaSession -import com.google.adk.tokt.InteropDispatcher import com.google.adk.tokt.codecs.EventCodec import com.google.adk.tokt.codecs.KtBackedEventsView import com.google.adk.tokt.codecs.SessionCodec @@ -38,6 +37,7 @@ import java.util.concurrent.ConcurrentMap import kotlin.jvm.optionals.getOrNull import kotlin.time.toJavaInstant import kotlin.time.toKotlinInstant +import kotlinx.coroutines.CoroutineDispatcher import kotlinx.coroutines.rx3.rxCompletable import kotlinx.coroutines.rx3.rxMaybe import kotlinx.coroutines.rx3.rxSingle @@ -45,10 +45,12 @@ import kotlinx.coroutines.rx3.rxSingle /** * A Java [JavaBaseSessionService] backed by a Kotlin [KtSessionService] (reverse of the Java * service wrappers): a Java agent running on the Kotlin runner sees a Java session service whose - * operations run on the Kotlin service. + * operations run on the Kotlin service, bridged on `dispatcher`. */ -internal class KtSessionServiceToJava(internal val service: KtSessionService) : - JavaBaseSessionService { +internal class KtSessionServiceToJava( + internal val service: KtSessionService, + private val dispatcher: CoroutineDispatcher, +) : JavaBaseSessionService { @Deprecated("Deprecated in BaseSessionService") override fun createSession( @@ -57,7 +59,7 @@ internal class KtSessionServiceToJava(internal val service: KtSessionService) : state: ConcurrentMap?, sessionId: String?, ): Single = - rxSingle(InteropDispatcher) { + rxSingle(dispatcher) { ktSessionToJava(service.createSession(SessionKey(appName, userId, sessionId), state)) } @@ -67,14 +69,14 @@ internal class KtSessionServiceToJava(internal val service: KtSessionService) : sessionId: String, config: Optional, ): Maybe = - rxMaybe(InteropDispatcher) { + rxMaybe(dispatcher) { service .getSession(SessionKey(appName, userId, sessionId), config.getOrNull()?.toKotlin()) ?.let { ktSessionToJava(it) } } override fun listSessions(appName: String, userId: String): Single = - rxSingle(InteropDispatcher) { + rxSingle(dispatcher) { val response = service.listSessions(appName, userId) JavaListSessionsResponse.builder() .sessions(response.sessions.map { ktSessionToJava(it) }) @@ -82,19 +84,17 @@ internal class KtSessionServiceToJava(internal val service: KtSessionService) : } override fun closeSession(session: JavaSession): Completable = - rxCompletable(InteropDispatcher) { service.closeSession(SessionCodec.fromJava(session)) } + rxCompletable(dispatcher) { service.closeSession(SessionCodec.fromJava(session)) } override fun deleteSession(appName: String, userId: String, sessionId: String): Completable = - rxCompletable(InteropDispatcher) { - service.deleteSession(SessionKey(appName, userId, sessionId)) - } + rxCompletable(dispatcher) { service.deleteSession(SessionKey(appName, userId, sessionId)) } override fun listEvents( appName: String, userId: String, sessionId: String, ): Single = - rxSingle(InteropDispatcher) { + rxSingle(dispatcher) { val response = service.listEvents(SessionKey(appName, userId, sessionId)) val builder = JavaListEventsResponse.builder().events(response.events.map { EventCodec.toJava(it) }) @@ -108,7 +108,7 @@ internal class KtSessionServiceToJava(internal val service: KtSessionService) : * `Runner` keeps observing the appended state in place. */ override fun appendEvent(session: JavaSession, event: JavaEvent): Single = - rxSingle(InteropDispatcher) { + rxSingle(dispatcher) { val key = SessionKey(session.appName(), session.userId(), session.id()) service.appendEvent(SessionCodec.fromJava(session), EventCodec.fromJava(event)) service.getSession(key)?.let { stored -> diff --git a/tokt/src/main/kotlin/com/google/adk/tokt/services/ServiceAdapters.kt b/tokt/src/main/kotlin/com/google/adk/tokt/services/ServiceAdapters.kt index 384ba6d90..8532bf44a 100644 --- a/tokt/src/main/kotlin/com/google/adk/tokt/services/ServiceAdapters.kt +++ b/tokt/src/main/kotlin/com/google/adk/tokt/services/ServiceAdapters.kt @@ -22,27 +22,47 @@ import com.google.adk.kt.memory.MemoryService as KtMemoryService import com.google.adk.kt.sessions.SessionService as KtSessionService import com.google.adk.memory.BaseMemoryService as JavaBaseMemoryService import com.google.adk.sessions.BaseSessionService as JavaBaseSessionService +import kotlinx.coroutines.CoroutineDispatcher /** * Service adapter factories that unwrap a round-tripped service rather than stacking adapters: * exposing a Kotlin service that is itself a wrapped Java service returns the original Java service * (and vice versa), so a Java -> Kt -> Java round-trip collapses to one reference with no extra - * hop. + * hop. A freshly wrapped service crosses on `dispatcher`; an unwrapped one is returned directly, + * with no adapter and so no dispatcher hop. */ -internal fun ktSessionServiceAsJava(service: KtSessionService): JavaBaseSessionService = - (service as? JavaSessionServiceToKt)?.service ?: KtSessionServiceToJava(service) +internal fun ktSessionServiceAsJava( + service: KtSessionService, + dispatcher: CoroutineDispatcher, +): JavaBaseSessionService = + (service as? JavaSessionServiceToKt)?.service ?: KtSessionServiceToJava(service, dispatcher) -internal fun javaSessionServiceAsKt(service: JavaBaseSessionService): KtSessionService = - (service as? KtSessionServiceToJava)?.service ?: JavaSessionServiceToKt(service) +internal fun javaSessionServiceAsKt( + service: JavaBaseSessionService, + dispatcher: CoroutineDispatcher, +): KtSessionService = + (service as? KtSessionServiceToJava)?.service ?: JavaSessionServiceToKt(service, dispatcher) -internal fun ktArtifactServiceAsJava(service: KtArtifactService): JavaBaseArtifactService = - (service as? JavaArtifactServiceToKt)?.service ?: KtArtifactServiceToJava(service) +internal fun ktArtifactServiceAsJava( + service: KtArtifactService, + dispatcher: CoroutineDispatcher, +): JavaBaseArtifactService = + (service as? JavaArtifactServiceToKt)?.service ?: KtArtifactServiceToJava(service, dispatcher) -internal fun javaArtifactServiceAsKt(service: JavaBaseArtifactService): KtArtifactService = - (service as? KtArtifactServiceToJava)?.service ?: JavaArtifactServiceToKt(service) +internal fun javaArtifactServiceAsKt( + service: JavaBaseArtifactService, + dispatcher: CoroutineDispatcher, +): KtArtifactService = + (service as? KtArtifactServiceToJava)?.service ?: JavaArtifactServiceToKt(service, dispatcher) -internal fun ktMemoryServiceAsJava(service: KtMemoryService): JavaBaseMemoryService = - (service as? JavaMemoryServiceToKt)?.service ?: KtMemoryServiceToJava(service) +internal fun ktMemoryServiceAsJava( + service: KtMemoryService, + dispatcher: CoroutineDispatcher, +): JavaBaseMemoryService = + (service as? JavaMemoryServiceToKt)?.service ?: KtMemoryServiceToJava(service, dispatcher) -internal fun javaMemoryServiceAsKt(service: JavaBaseMemoryService): KtMemoryService = - (service as? KtMemoryServiceToJava)?.service ?: JavaMemoryServiceToKt(service) +internal fun javaMemoryServiceAsKt( + service: JavaBaseMemoryService, + dispatcher: CoroutineDispatcher, +): KtMemoryService = + (service as? KtMemoryServiceToJava)?.service ?: JavaMemoryServiceToKt(service, dispatcher) 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 ff617cdc3..b1861703e 100644 --- a/tokt/src/test/kotlin/com/google/adk/tokt/KtRunnerInteropTest.kt +++ b/tokt/src/test/kotlin/com/google/adk/tokt/KtRunnerInteropTest.kt @@ -40,7 +40,10 @@ import com.google.adk.kt.models.LlmResponse as KtLlmResponse import com.google.adk.kt.runners.InMemoryRunner as KtInMemoryRunner import com.google.adk.kt.runners.Runner as KtRunner import com.google.adk.kt.sessions.GetSessionConfig as KtGetSessionConfig +import com.google.adk.kt.sessions.InMemorySessionService as KtInMemorySessionService +import com.google.adk.kt.sessions.Session as KtSession import com.google.adk.kt.sessions.SessionKey as KtSessionKey +import com.google.adk.kt.sessions.SessionService as KtSessionService 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 @@ -137,6 +140,7 @@ import java.util.Optional import java.util.concurrent.ConcurrentHashMap import java.util.concurrent.ConcurrentMap import java.util.concurrent.CopyOnWriteArrayList +import java.util.concurrent.Executors import java.util.concurrent.atomic.AtomicInteger import java.util.concurrent.atomic.AtomicReference import kotlin.jvm.optionals.getOrNull @@ -151,6 +155,7 @@ import kotlin.test.assertSame import kotlin.test.assertTrue import kotlin.test.fail import kotlin.time.Duration.Companion.seconds +import kotlinx.coroutines.asCoroutineDispatcher import kotlinx.coroutines.flow.Flow import kotlinx.coroutines.flow.flow import kotlinx.coroutines.flow.toList @@ -710,6 +715,125 @@ class KtRunnerInteropTest { assertTrue(JavaAdkToKt.asKtTools(emptyList()).isEmpty()) } + @Test + fun javaAdkToKt_customDispatcher_runsAdaptedComponentsOnIt() = runBlocking { + // A custom dispatcher passed to the forward entry points must actually carry the adapted Java + // component's (possibly blocking) calls, rather than the default Dispatchers.IO. + val dispatcher = + Executors.newSingleThreadExecutor { r -> Thread(r, "tokt-custom-dispatcher") } + .asCoroutineDispatcher() + try { + val modelThread = AtomicReference() + val toolThread = AtomicReference() + val model = + object : JavaBaseLlm("java-model") { + private var step = 0 + + override fun generateContent( + llmRequest: JavaLlmRequest, + stream: Boolean, + ): Flowable { + modelThread.set(Thread.currentThread().name) + val content = + if (step++ == 0) modelFunctionCall("java_echo", mapOf("text" to "hi")) + else modelText("done") + return Flowable.just(JavaLlmResponse.builder().content(content).build()) + } + + override fun connect(llmRequest: JavaLlmRequest): JavaBaseLlmConnection = + throw UnsupportedOperationException() + } + val tool = + object : JavaBaseTool("java_echo", "echoes") { + override fun declaration(): Optional = + Optional.of(GenaiFunctionDeclaration.builder().name("java_echo").build()) + + @JvmSuppressWildcards + override fun runAsync( + args: Map, + toolContext: JavaToolContext, + ): Single> { + toolThread.set(Thread.currentThread().name) + return Single.just(mapOf("echoed" to (args["text"] ?: ""))) + } + } + val agent = + KtLlmAgent( + name = "a", + model = JavaAdkToKt.asKtModel(model, dispatcher), + tools = listOf(JavaAdkToKt.asKtTool(tool, dispatcher)), + ) + val runner = KtInMemoryRunner(agent, appName = "app") + + runner.turn() + + // Coroutines may append " @coroutine#N" to the thread name, so match the pool thread's + // prefix. + assertTrue( + modelThread.get().orEmpty().startsWith("tokt-custom-dispatcher"), + "the adapted Java model should run on the supplied dispatcher, not Dispatchers.IO; ran on " + + modelThread.get(), + ) + assertTrue( + toolThread.get().orEmpty().startsWith("tokt-custom-dispatcher"), + "the adapted Java tool should run on the supplied dispatcher, not Dispatchers.IO; ran on " + + toolThread.get(), + ) + } finally { + dispatcher.close() + } + } + + @Test + fun asJavaRunner_customDispatcher_runsReverseServiceAdaptersOnIt() { + // The reverse direction (asJavaRunner) bridges Java RxJava calls back onto the Kotlin engine; a + // reverse service adapter must run on the supplied dispatcher, not the default Dispatchers.IO. + val dispatcher = + Executors.newSingleThreadExecutor { r -> Thread(r, "tokt-reverse-dispatcher") } + .asCoroutineDispatcher() + try { + val sessionThread = AtomicReference() + // A native Kotlin session service that records the thread its getSession runs on. + val recording = + object : KtSessionService by KtInMemorySessionService() { + override suspend fun getSession( + key: KtSessionKey, + config: KtGetSessionConfig?, + ): KtSession? { + sessionThread.set(Thread.currentThread().name) + return null + } + } + val ktRunner = + KtInMemoryRunner( + app = + KtApp( + appName = "app", + rootAgent = + KtLlmAgent( + name = "a", + model = JavaAdkToKt.asKtModel(SequentialJavaModel(emptyList())), + ), + ), + sessionService = recording, + ) + val javaRunner = KotlinAdkToJava.asJavaRunner(ktRunner, dispatcher) + + // Route through the Java-facing reverse session-service adapter (KtSessionServiceToJava + // .getSession = rxMaybe(dispatcher) { service.getSession(...) }). + val unused = + javaRunner.sessionService().getSession("app", "u", "s", Optional.empty()).blockingGet() + + assertTrue( + sessionThread.get().orEmpty().startsWith("tokt-reverse-dispatcher"), + "the reverse session-service adapter should run on the supplied dispatcher; ran on " + + sessionThread.get(), + ) + } finally { + dispatcher.close() + } + } + @Test fun ktRunner_stateDeltaWrittenToAnAdaptedJavaSessionService_mapsOnlyTheRemovalSentinel() = runBlocking {