feat: let the ADK Java <-> Kotlin interop run on a caller-supplied dispatcher

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

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<JavaBaseTool>): List<KtBaseTool> = javaTools.map { asKtTool(it) }
+  @JvmOverloads
+  fun asKtTools(
+    javaTools: List<JavaBaseTool>,
+    dispatcher: CoroutineDispatcher = InteropDispatcher,
+  ): List<KtBaseTool> = 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<JavaBaseToolset>): List<KtToolset> = javaToolsets.map {
-    asKtToolset(it)
-  }
+  @JvmOverloads
+  fun asKtToolsets(
+    javaToolsets: List<JavaBaseToolset>,
+    dispatcher: CoroutineDispatcher = InteropDispatcher,
+  ): List<KtToolset> = 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<JavaPlugin>): List<KtPlugin> = javaPlugins.map {
-    asKtPlugin(it)
-  }
+  @JvmOverloads
+  fun asKtPlugins(
+    javaPlugins: List<JavaPlugin>,
+    dispatcher: CoroutineDispatcher = InteropDispatcher,
+  ): List<KtPlugin> = 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)
 }

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)
 }

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,

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<LlmResponse>` 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<LlmResponse> =
     flow {
         emitAll(
@@ -54,5 +57,5 @@ internal class JavaModelToKt(private val javaLlm: JavaBaseLlm) : Model {
           }
         )
       }
-      .flowOn(InteropDispatcher)
+      .flowOn(dispatcher)
 }