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 {