From f5df016c4c961459999b3b1b7aa250bc66597a69 Mon Sep 17 00:00:00 2001 From: Huy Nguyen <290170845+huynguyen20011@users.noreply.github.com> Date: Thu, 8 Oct 2026 19:10:45 -0700 Subject: [PATCH 1/6] feat(tooling): Improve tool usage lifecycle and timeout handling - Ensures that the tool approval awaitable is properly completed regardless of the exit path. - Implements robust timeout handling within `SessionManager` to clear state on timeout. - Exposes a dedicated callback in `ChatClient` for managing tool completion promises. BREAKING CHANGE: The `sendStreamingMessageWithCallback` function now accepts `onToolFinishedAwaitable` to explicitly manage the completion of the tool's pending promise. --- .../main/kotlin/io/askimo/ui/chat/ChatView.kt | 1 - .../io/askimo/ui/session/SessionManager.kt | 18 +++ .../core/providers/ChatClientExtensions.kt | 59 ++++++++++ .../ChatClientExtensionsToolApprovalTest.kt | 104 ++++++++++++++++++ 4 files changed, 181 insertions(+), 1 deletion(-) diff --git a/desktop-shared/src/main/kotlin/io/askimo/ui/chat/ChatView.kt b/desktop-shared/src/main/kotlin/io/askimo/ui/chat/ChatView.kt index ef3fdf515..38a4e48c6 100644 --- a/desktop-shared/src/main/kotlin/io/askimo/ui/chat/ChatView.kt +++ b/desktop-shared/src/main/kotlin/io/askimo/ui/chat/ChatView.kt @@ -1338,7 +1338,6 @@ fun chatView( Text( text = stringResource("chat.tool.approval.prompt", pendingToolApproval.toolName), style = AppTextStyles.caption, - color = MaterialTheme.colorScheme.onPrimaryContainer, modifier = Modifier.weight(1f), ) Row(horizontalArrangement = Arrangement.spacedBy(Spacing.small)) { diff --git a/desktop-shared/src/main/kotlin/io/askimo/ui/session/SessionManager.kt b/desktop-shared/src/main/kotlin/io/askimo/ui/session/SessionManager.kt index 144d89f26..e995630d7 100644 --- a/desktop-shared/src/main/kotlin/io/askimo/ui/session/SessionManager.kt +++ b/desktop-shared/src/main/kotlin/io/askimo/ui/session/SessionManager.kt @@ -51,6 +51,7 @@ import kotlinx.coroutines.sync.Mutex import kotlinx.coroutines.sync.withLock import kotlinx.coroutines.withContext import java.util.UUID +import java.util.concurrent.CompletableFuture import java.util.concurrent.ConcurrentHashMap import kotlin.time.Clock @@ -431,6 +432,23 @@ class SessionManager( ) } }, + onToolApprovalTimedOut = { + streamingScope.launch { thread.clearApproval() } + }, + // Awaitable companion to onToolFinished: used only on approval + // timeout/denial so markToolDone() applies before the exception + // propagates — otherwise the tool could be persisted as RUNNING. + onToolFinishedAwaitable = { toolName, arguments, result, hasFailed -> + CompletableFuture().apply { + streamingScope.launch { + try { + thread.markToolDone(toolName, arguments, result, hasFailed) + } finally { + complete(Unit) + } + } + } + }, chatMemory = sessionContext.memory, ) } finally { diff --git a/shared-jvm/src/main/kotlin/io/askimo/core/providers/ChatClientExtensions.kt b/shared-jvm/src/main/kotlin/io/askimo/core/providers/ChatClientExtensions.kt index af4ed1ce7..7b100d1ef 100644 --- a/shared-jvm/src/main/kotlin/io/askimo/core/providers/ChatClientExtensions.kt +++ b/shared-jvm/src/main/kotlin/io/askimo/core/providers/ChatClientExtensions.kt @@ -37,6 +37,7 @@ import io.askimo.core.intent.ToolApprovalPolicy import io.askimo.core.intent.ToolConfig import io.askimo.core.intent.ToolRegistry import io.askimo.core.intent.defaultApprovalPolicy +import io.askimo.core.logging.currentFileLogger import io.askimo.core.logging.logger import io.askimo.core.memory.SessionConversationSummary import io.askimo.core.memory.UserMemorySummary @@ -49,9 +50,12 @@ import kotlinx.serialization.json.JsonElement import kotlinx.serialization.json.JsonObject import kotlinx.serialization.json.JsonPrimitive import kotlinx.serialization.json.jsonArray +import java.util.concurrent.CompletableFuture import java.util.concurrent.CountDownLatch import java.util.concurrent.TimeUnit +private val log = currentFileLogger() + /** * Checks whether any message in this exception's cause chain contains [needle] * (case-insensitive). Unlike a single `e.message` / `e.cause?.message` check, this @@ -86,6 +90,44 @@ internal sealed class StreamingErrorResult { data class Terminal(val message: String) : StreamingErrorResult() } +/** + * Notifies tool completion on the approval-timeout/denial paths, awaiting it when possible + * before the caller propagates its exception. + * + * Prefers [onToolFinishedAwaitable] when non-null, awaiting its [CompletableFuture] (bounded by + * [timeoutMs]) so e.g. `SessionManager.markToolDone` applies before a later `catch` block reads + * stale (still-RUNNING) state. Falls back to firing [onToolFinished] without waiting — preserving + * its original behavior for callers that don't supply the awaitable variant. + * + * @param onToolFinished Legacy, non-awaitable callback — invoked without waiting if [onToolFinishedAwaitable] is null. + * @param onToolFinishedAwaitable Optional awaitable variant — when present, takes precedence and is awaited. + * @param toolName Name of the tool + * @param arguments Tool arguments + * @param result Tool result + * @param hasFailed Whether the tool failed + * @param timeoutMs Timeout for waiting on the completion future + */ +private fun notifyToolFinishedAndAwaitIfPossible( + onToolFinished: ((String, String?, String?, Boolean) -> Unit)?, + onToolFinishedAwaitable: ((String, String?, String?, Boolean) -> CompletableFuture?)?, + toolName: String, + arguments: String?, + result: String?, + hasFailed: Boolean, + timeoutMs: Long, +) { + if (onToolFinishedAwaitable != null) { + val completion = onToolFinishedAwaitable.invoke(toolName, arguments, result, hasFailed) + try { + completion?.get(timeoutMs, TimeUnit.MILLISECONDS) + } catch (e: Exception) { + log.warn("Tool completion callback timed out for $toolName", e) + } + } else { + onToolFinished?.invoke(toolName, arguments, result, hasFailed) + } +} + /** * Pure classification function: maps a streaming [Throwable] to a [StreamingErrorResult]. * @@ -153,6 +195,10 @@ internal fun classifyStreamingError( * @param userContents the contents sent by user * @param onToken Optional callback function that is invoked for each token received from the model * @param onFollowUpSuggestion Optional callback for follow-up suggestions based on AI response + * @param onToolFinishedAwaitable Awaitable variant of [onToolFinished], used instead of it on the + * approval-timeout/denial paths: its [CompletableFuture] is awaited (bounded by [toolApprovalTimeoutMs]) + * before throwing, so callers like `SessionManager` can finish state updates (e.g. `markToolDone`) + * before the tool is read as still RUNNING elsewhere. * @return The complete response from the language model as a string */ fun ChatClient.sendStreamingMessageWithCallback( @@ -183,6 +229,13 @@ fun ChatClient.sendStreamingMessageWithCallback( * - Invoke `deny` to cancel the tool call (surfaces as a [ToolExecutionException]). */ onToolApprovalRequired: ((toolName: String, arguments: String?, approve: () -> Unit, deny: () -> Unit) -> Unit)? = null, + /** + * Called when the approval wait in [onToolApprovalRequired] times out (neither `approve` + * nor `deny` was invoked). Lets the caller clear pending-approval UI state that its own + * `approve`/`deny` closures never got to run — e.g. `SessionManager.StreamingThread.clearApproval()` + * — otherwise an approval banner stays visible forever. + */ + onToolApprovalTimedOut: (() -> Unit)? = null, /** * How long to wait for [onToolApprovalRequired] to invoke `approve`/`deny` before treating * the request as timed out. Defaults to [io.askimo.core.config.ModelsConfig.toolApprovalTimeoutMs]; @@ -196,6 +249,7 @@ fun ChatClient.sendStreamingMessageWithCallback( * is left without a matching `tool_result`, which Anthropic's API rejects on the next turn. */ chatMemory: ChatMemory? = null, + onToolFinishedAwaitable: ((toolName: String, arguments: String?, result: String?, hasFailed: Boolean) -> CompletableFuture?)? = null, ): String { val log = logger() @@ -333,6 +387,10 @@ fun ChatClient.sendStreamingMessageWithCallback( chatMemory?.add(ToolExecutionResultMessage.from(req, msg)) } pendingToolRequests.clear() + notifyToolFinishedAndAwaitIfPossible(onToolFinished, onToolFinishedAwaitable, toolName, arguments, msg, true, toolApprovalTimeoutMs) + // Neither approve() nor deny() fired, so the caller's pending-approval + // UI state was never cleared by either closure above. + onToolApprovalTimedOut?.invoke() throw ToolExecutionException(toolName = toolName, errorDetails = msg) } if (!approved) { @@ -341,6 +399,7 @@ fun ChatClient.sendStreamingMessageWithCallback( chatMemory?.add(ToolExecutionResultMessage.from(req, msg)) } pendingToolRequests.clear() + notifyToolFinishedAndAwaitIfPossible(onToolFinished, onToolFinishedAwaitable, toolName, arguments, msg, true, toolApprovalTimeoutMs) throw ToolExecutionException(toolName = toolName, errorDetails = msg) } } diff --git a/shared-jvm/src/test/kotlin/io/askimo/core/providers/ChatClientExtensionsToolApprovalTest.kt b/shared-jvm/src/test/kotlin/io/askimo/core/providers/ChatClientExtensionsToolApprovalTest.kt index f613270bc..000a69f6e 100644 --- a/shared-jvm/src/test/kotlin/io/askimo/core/providers/ChatClientExtensionsToolApprovalTest.kt +++ b/shared-jvm/src/test/kotlin/io/askimo/core/providers/ChatClientExtensionsToolApprovalTest.kt @@ -30,7 +30,9 @@ import org.junit.jupiter.api.BeforeEach import org.junit.jupiter.api.Test import org.junit.jupiter.api.io.TempDir import java.nio.file.Path +import java.util.concurrent.CompletableFuture import java.util.function.Consumer +import kotlin.test.assertEquals import kotlin.test.assertFailsWith import kotlin.test.assertTrue @@ -227,4 +229,106 @@ class ChatClientExtensionsToolApprovalTest { assertTrue(requestA.id() in resultIds, "Expected a tool_result for requestA (${requestA.id()}), got: $resultIds") assertTrue(requestB.id() in resultIds, "Expected a tool_result for requestB (${requestB.id()}), got: $resultIds") } + + /** + * Regression test for the async/sync race documented on `notifyToolFinishedAndAwaitIfPossible`: + * the timeout branch must synchronously await the caller's `onToolFinishedAwaitable` completion + * future *before* invoking `onToolApprovalTimedOut`/throwing — otherwise a caller like + * `SessionManager` that persists UI state inside `onToolApprovalTimedOut` would observe the + * tool still marked RUNNING. Asserting the exact invocation order (not just "both eventually + * ran") is what actually catches a regression back to the fire-and-forget launch. + */ + @Test + fun `approval timeout invokes onToolFinishedAwaitable before onToolApprovalTimedOut, both exactly once`() { + val chatMemory = FakeChatMemory() + val request = toolRequest("call_solo", "toolTimeout") + val invocationOrder = mutableListOf() + var toolFinishedCalls = 0 + var timedOutCalls = 0 + var toolFinishedHasFailed: Boolean? = null + + val chatClient = object : ChatClient { + override fun sendMessageStreaming(userContents: List): TokenStream = FakeTokenStream { + chatMemory.add(AiMessage.from(listOf(request))) + beforeToolExecutionHandler?.accept(beforeToolExecution(request)) + } + + override fun sendMessage(prompt: String): String = error("not used") + } + + assertFailsWith { + chatClient.sendStreamingMessageWithCallback( + userContents = listOf(), + resolvedTools = listOf(requireApprovalConfig("toolTimeout")), + // Never invokes approve/deny — forces the timeout branch. + onToolApprovalRequired = { _, _, _, _ -> }, + onToolFinishedAwaitable = { _, _, _, hasFailed -> + toolFinishedCalls++ + toolFinishedHasFailed = hasFailed + invocationOrder += "onToolFinishedAwaitable" + CompletableFuture.completedFuture(Unit) + }, + onToolApprovalTimedOut = { + timedOutCalls++ + invocationOrder += "onToolApprovalTimedOut" + }, + toolApprovalTimeoutMs = 50L, + chatMemory = chatMemory, + ) + } + + assertEquals(1, toolFinishedCalls, "Expected onToolFinishedAwaitable to be invoked exactly once, got $toolFinishedCalls") + assertEquals(1, timedOutCalls, "Expected onToolApprovalTimedOut to be invoked exactly once, got $timedOutCalls") + assertEquals(true, toolFinishedHasFailed, "Expected onToolFinishedAwaitable's hasFailed to be true on timeout") + assertEquals( + listOf("onToolFinishedAwaitable", "onToolApprovalTimedOut"), + invocationOrder, + "Expected onToolFinishedAwaitable's completion future to be awaited before onToolApprovalTimedOut fires", + ) + } + + /** + * Companion to the timeout test above: denial must still notify `onToolFinishedAwaitable` (so UI + * state like a RUNNING tool entry gets flipped to DONE/failed) but must NOT invoke + * `onToolApprovalTimedOut` — that callback exists solely to clear pending-approval UI state + * left behind by a timeout, which never applies on an explicit deny. + */ + @Test + fun `approval denial invokes onToolFinishedAwaitable but never onToolApprovalTimedOut`() { + val chatMemory = FakeChatMemory() + val request = toolRequest("call_deny", "toolDeny") + var toolFinishedCalls = 0 + var timedOutCalls = 0 + var toolFinishedHasFailed: Boolean? = null + + val chatClient = object : ChatClient { + override fun sendMessageStreaming(userContents: List): TokenStream = FakeTokenStream { + chatMemory.add(AiMessage.from(listOf(request))) + beforeToolExecutionHandler?.accept(beforeToolExecution(request)) + } + + override fun sendMessage(prompt: String): String = error("not used") + } + + assertFailsWith { + chatClient.sendStreamingMessageWithCallback( + userContents = listOf(), + resolvedTools = listOf(requireApprovalConfig("toolDeny")), + onToolApprovalRequired = { _, _, _, deny -> deny() }, + onToolFinishedAwaitable = { _, _, _, hasFailed -> + toolFinishedCalls++ + toolFinishedHasFailed = hasFailed + CompletableFuture.completedFuture(Unit) + }, + onToolApprovalTimedOut = { + timedOutCalls++ + }, + chatMemory = chatMemory, + ) + } + + assertEquals(1, toolFinishedCalls, "Expected onToolFinishedAwaitable to be invoked exactly once on denial, got $toolFinishedCalls") + assertEquals(true, toolFinishedHasFailed, "Expected onToolFinishedAwaitable's hasFailed to be true on denial") + assertEquals(0, timedOutCalls, "onToolApprovalTimedOut must never fire on explicit denial, got $timedOutCalls calls") + } } From 2c9955fdf5837def8ea4e48c2595cd48cea44274 Mon Sep 17 00:00:00 2001 From: Huy Nguyen <290170845+huynguyen20011@users.noreply.github.com> Date: Thu, 8 Oct 2026 19:24:40 -0700 Subject: [PATCH 2/6] Update per feedbacks --- .../core/providers/ChatClientExtensions.kt | 7 +++- .../ChatClientExtensionsToolApprovalTest.kt | 41 ++++++++++++++----- 2 files changed, 35 insertions(+), 13 deletions(-) diff --git a/shared-jvm/src/main/kotlin/io/askimo/core/providers/ChatClientExtensions.kt b/shared-jvm/src/main/kotlin/io/askimo/core/providers/ChatClientExtensions.kt index 7b100d1ef..7562568fe 100644 --- a/shared-jvm/src/main/kotlin/io/askimo/core/providers/ChatClientExtensions.kt +++ b/shared-jvm/src/main/kotlin/io/askimo/core/providers/ChatClientExtensions.kt @@ -387,10 +387,13 @@ fun ChatClient.sendStreamingMessageWithCallback( chatMemory?.add(ToolExecutionResultMessage.from(req, msg)) } pendingToolRequests.clear() - notifyToolFinishedAndAwaitIfPossible(onToolFinished, onToolFinishedAwaitable, toolName, arguments, msg, true, toolApprovalTimeoutMs) // Neither approve() nor deny() fired, so the caller's pending-approval - // UI state was never cleared by either closure above. + // UI state was never cleared by either closure above. Fire this + // immediately — not after awaiting onToolFinishedAwaitable below, + // which can itself block for up to toolApprovalTimeoutMs and would + // otherwise leave the approval banner visible for twice as long. onToolApprovalTimedOut?.invoke() + notifyToolFinishedAndAwaitIfPossible(onToolFinished, onToolFinishedAwaitable, toolName, arguments, msg, true, toolApprovalTimeoutMs) throw ToolExecutionException(toolName = toolName, errorDetails = msg) } if (!approved) { diff --git a/shared-jvm/src/test/kotlin/io/askimo/core/providers/ChatClientExtensionsToolApprovalTest.kt b/shared-jvm/src/test/kotlin/io/askimo/core/providers/ChatClientExtensionsToolApprovalTest.kt index 000a69f6e..984ecb043 100644 --- a/shared-jvm/src/test/kotlin/io/askimo/core/providers/ChatClientExtensionsToolApprovalTest.kt +++ b/shared-jvm/src/test/kotlin/io/askimo/core/providers/ChatClientExtensionsToolApprovalTest.kt @@ -31,6 +31,7 @@ import org.junit.jupiter.api.Test import org.junit.jupiter.api.io.TempDir import java.nio.file.Path import java.util.concurrent.CompletableFuture +import java.util.concurrent.atomic.AtomicBoolean import java.util.function.Consumer import kotlin.test.assertEquals import kotlin.test.assertFailsWith @@ -231,21 +232,23 @@ class ChatClientExtensionsToolApprovalTest { } /** - * Regression test for the async/sync race documented on `notifyToolFinishedAndAwaitIfPossible`: - * the timeout branch must synchronously await the caller's `onToolFinishedAwaitable` completion - * future *before* invoking `onToolApprovalTimedOut`/throwing — otherwise a caller like - * `SessionManager` that persists UI state inside `onToolApprovalTimedOut` would observe the - * tool still marked RUNNING. Asserting the exact invocation order (not just "both eventually - * ran") is what actually catches a regression back to the fire-and-forget launch. + * `onToolApprovalTimedOut` must fire immediately on timeout — before awaiting + * `onToolFinishedAwaitable`'s completion future, which can itself block for up to + * `toolApprovalTimeoutMs`. Otherwise a caller clearing UI state in `onToolApprovalTimedOut` + * would see it stay visible for up to twice the configured timeout. + * + * Returns a future completed after a delay (not [CompletableFuture.completedFuture]) so an + * implementation that ignores it instead of calling `future.get()` returns too early to pass. */ @Test - fun `approval timeout invokes onToolFinishedAwaitable before onToolApprovalTimedOut, both exactly once`() { + fun `approval timeout invokes onToolApprovalTimedOut before awaiting onToolFinishedAwaitable, both exactly once`() { val chatMemory = FakeChatMemory() val request = toolRequest("call_solo", "toolTimeout") val invocationOrder = mutableListOf() var toolFinishedCalls = 0 var timedOutCalls = 0 var toolFinishedHasFailed: Boolean? = null + val futureCompleted = AtomicBoolean(false) val chatClient = object : ChatClient { override fun sendMessageStreaming(userContents: List): TokenStream = FakeTokenStream { @@ -266,13 +269,24 @@ class ChatClientExtensionsToolApprovalTest { toolFinishedCalls++ toolFinishedHasFailed = hasFailed invocationOrder += "onToolFinishedAwaitable" - CompletableFuture.completedFuture(Unit) + val future = CompletableFuture() + // Completed only after a delay, from a background thread — forces a caller + // that actually awaits this future to block until then; one that ignores it + // (the bug this test guards against) returns well before futureCompleted flips. + Thread { + Thread.sleep(150) + futureCompleted.set(true) + future.complete(Unit) + }.start() + future }, onToolApprovalTimedOut = { timedOutCalls++ invocationOrder += "onToolApprovalTimedOut" }, - toolApprovalTimeoutMs = 50L, + // Large enough that the 150ms delayed completion above always wins the race + // against this timeout inside notifyToolFinishedAndAwaitIfPossible's get() call. + toolApprovalTimeoutMs = 300L, chatMemory = chatMemory, ) } @@ -281,9 +295,14 @@ class ChatClientExtensionsToolApprovalTest { assertEquals(1, timedOutCalls, "Expected onToolApprovalTimedOut to be invoked exactly once, got $timedOutCalls") assertEquals(true, toolFinishedHasFailed, "Expected onToolFinishedAwaitable's hasFailed to be true on timeout") assertEquals( - listOf("onToolFinishedAwaitable", "onToolApprovalTimedOut"), + listOf("onToolApprovalTimedOut", "onToolFinishedAwaitable"), invocationOrder, - "Expected onToolFinishedAwaitable's completion future to be awaited before onToolApprovalTimedOut fires", + "Expected onToolApprovalTimedOut to fire immediately on timeout, before awaiting onToolFinishedAwaitable's completion future", + ) + assertTrue( + futureCompleted.get(), + "Expected the call to block until onToolFinishedAwaitable's future completed — " + + "it returned before the delayed completion ran, meaning the future was never actually awaited", ) } From e80182b2284ddc526aab48754b73affa0233a181 Mon Sep 17 00:00:00 2001 From: Huy Nguyen <290170845+huynguyen20011@users.noreply.github.com> Date: Thu, 8 Oct 2026 19:45:46 -0700 Subject: [PATCH 3/6] Update per feedback --- .../io/askimo/ui/session/SessionManager.kt | 69 ++++++++++--- .../ChatClientExtensionsToolApprovalTest.kt | 96 +++++++++++-------- 2 files changed, 116 insertions(+), 49 deletions(-) diff --git a/desktop-shared/src/main/kotlin/io/askimo/ui/session/SessionManager.kt b/desktop-shared/src/main/kotlin/io/askimo/ui/session/SessionManager.kt index e995630d7..fee121ea1 100644 --- a/desktop-shared/src/main/kotlin/io/askimo/ui/session/SessionManager.kt +++ b/desktop-shared/src/main/kotlin/io/askimo/ui/session/SessionManager.kt @@ -45,6 +45,7 @@ import kotlinx.coroutines.flow.MutableStateFlow import kotlinx.coroutines.flow.StateFlow import kotlinx.coroutines.flow.asStateFlow import kotlinx.coroutines.flow.filterIsInstance +import kotlinx.coroutines.joinAll import kotlinx.coroutines.launch import kotlinx.coroutines.runBlocking import kotlinx.coroutines.sync.Mutex @@ -53,6 +54,7 @@ import kotlinx.coroutines.withContext import java.util.UUID import java.util.concurrent.CompletableFuture import java.util.concurrent.ConcurrentHashMap +import java.util.concurrent.ConcurrentLinkedQueue import kotlin.time.Clock /** @@ -206,7 +208,10 @@ class SessionManager( suspend fun markToolRunning(toolName: String, arguments: String?) { mutex.withLock { val alreadyRunning = _timeline.value.any { - it is TurnTimelineEntry.Tool && it.toolCall.toolName == toolName && it.toolCall.status == ToolCallStatus.RUNNING + it is TurnTimelineEntry.Tool && + it.toolCall.toolName == toolName && + it.toolCall.arguments == arguments && + it.toolCall.status == ToolCallStatus.RUNNING } if (!alreadyRunning) { _timeline.value += TurnTimelineEntry.Tool( @@ -222,8 +227,15 @@ class SessionManager( suspend fun markToolDone(toolName: String, arguments: String?, result: String?, hasFailed: Boolean) { mutex.withLock { val list = _timeline.value + // Matched on (toolName, arguments) rather than toolName alone — otherwise two + // parallel invocations of the SAME tool (same name, different arguments) would + // collide: the second markToolRunning would be deduped as "already running" and + // this lookup would flip the wrong call's entry to DONE. val idx = list.indexOfLast { - it is TurnTimelineEntry.Tool && it.toolCall.toolName == toolName && it.toolCall.status == ToolCallStatus.RUNNING + it is TurnTimelineEntry.Tool && + it.toolCall.toolName == toolName && + it.toolCall.arguments == arguments && + it.toolCall.status == ToolCallStatus.RUNNING } // Preserve the RUNNING entry's original start time (instead of "now") so a // future "took Ns" label stays possible even after the live timer stops. @@ -326,16 +338,18 @@ class SessionManager( _hasFailed = MutableStateFlow(false), ) - // Register this thread - activeThreads[sessionId] = thread - - log.debug("Streaming thread $threadId for session $sessionId started. Active streams: ${activeThreads.size}") - // Prepare context and save user message to DB val promptWithContext = chatSessionService.prepareContextAndGetPromptForChat(sessionId, userMessage, willSaveUserMessage) log.debug("Saved prompt for session $sessionId, starting streaming") thread.job = streamingScope.launch { + // Tracks every fire-and-forget timeline-mutating coroutine (onToken, + // onThinkingToken, onToolStarted, onToolFinished) launched while + // sendStreamingMessageWithCallback is in-flight below. Joined — both on success + // and on any exception — before reading thread.timeline/getCurrentContent(), so + // persistence never observes a timeline missing the last chunk(s)/tool update(s) + // still mid-flight the instant the call returns or throws. + val pendingUpdateJobs = ConcurrentLinkedQueue() try { if (mode is CreationMode.Chat) { var capturedInputTokens: Int? = null @@ -370,13 +384,25 @@ class SessionManager( val fullResponse = try { val sessionContext = chatSessionService.getOrCreateContextForSession(sessionId) + // Tracks each tool call's in-flight markToolRunning job, keyed by + // (toolName, arguments) rather than toolName alone — otherwise two + // parallel invocations of the SAME tool (same name, different + // arguments) would collide: the second's markToolRunning would be + // deduped as "already running" and markToolDone could join/flip the + // wrong call's entry. markToolDone joins this first — onToolStarted/ + // onToolFinished each launch their update on streamingScope + // independently, and without this, a delayed RUNNING-append coroutine + // can land *after* a faster DONE-append one, leaving the timeline's + // last entry for that tool RUNNING (collapsedEffectiveTools treats it + // as effective) even though the tool already finished. + val pendingToolStartJobs = ConcurrentHashMap, Job>() sessionContext.chatClient .sendStreamingMessageWithCallback( projectId = projectId, userContents = promptWithContext, enabledServerIds = enabledServerIds, onToken = { token -> - streamingScope.launch { + pendingUpdateJobs += streamingScope.launch { thread.appendChunk(token) } }, @@ -391,17 +417,20 @@ class SessionManager( log.debug("Token usage for session $sessionId: input=$input, output=$output, total=$total, duration=${durationMs}ms") }, onToolStarted = { toolName, arguments -> - streamingScope.launch { + val job = streamingScope.launch { thread.markToolRunning(toolName, arguments) } + pendingToolStartJobs[toolName to arguments] = job + pendingUpdateJobs += job }, onToolFinished = { toolName, arguments, result, hasFailed -> - streamingScope.launch { + pendingUpdateJobs += streamingScope.launch { + pendingToolStartJobs.remove(toolName to arguments)?.join() thread.markToolDone(toolName, arguments, result, hasFailed) } }, onThinkingToken = { token -> - streamingScope.launch { + pendingUpdateJobs += streamingScope.launch { thread.appendThinkingChunk(token) } }, @@ -442,6 +471,7 @@ class SessionManager( CompletableFuture().apply { streamingScope.launch { try { + pendingToolStartJobs.remove(toolName to arguments)?.join() thread.markToolDone(toolName, arguments, result, hasFailed) } finally { complete(Unit) @@ -457,6 +487,11 @@ class SessionManager( if (needsMessageCorrelation) ProxyChatContext.clear() } + // Ensure every timeline mutation triggered during the call above has + // actually landed before persisting/reading the timeline below — see + // pendingUpdateJobs' declaration comment. + pendingUpdateJobs.joinAll() + val savedMessage = chatSessionService.saveAiResponse( sessionId = sessionId, response = fullResponse, @@ -517,6 +552,11 @@ class SessionManager( log.error("Error while sending message to chat session $sessionId", e) thread.markFailed() + // Same join as the success path above — the exception can propagate while + // onToken/onToolStarted/onToolFinished/onThinkingToken updates are still + // in-flight on streamingScope. + pendingUpdateJobs.joinAll() + val partialResponse = thread.getCurrentContent() val failedResponse = if (e is ConfigurationErrorException) { e.displayMessage @@ -564,6 +604,13 @@ class SessionManager( } } + // Registered only now that thread.job holds the real coroutine handle (not the + // placeholder Job() from construction above) — otherwise a stopStream() call landing + // in the gap between registration and this assignment would cancel the inert + // placeholder instead of the actual stream. + activeThreads[sessionId] = thread + log.debug("Streaming thread $threadId for session $sessionId started. Active streams: ${activeThreads.size}") + return threadId } diff --git a/shared-jvm/src/test/kotlin/io/askimo/core/providers/ChatClientExtensionsToolApprovalTest.kt b/shared-jvm/src/test/kotlin/io/askimo/core/providers/ChatClientExtensionsToolApprovalTest.kt index 984ecb043..e2da27bcf 100644 --- a/shared-jvm/src/test/kotlin/io/askimo/core/providers/ChatClientExtensionsToolApprovalTest.kt +++ b/shared-jvm/src/test/kotlin/io/askimo/core/providers/ChatClientExtensionsToolApprovalTest.kt @@ -31,7 +31,9 @@ import org.junit.jupiter.api.Test import org.junit.jupiter.api.io.TempDir import java.nio.file.Path import java.util.concurrent.CompletableFuture -import java.util.concurrent.atomic.AtomicBoolean +import java.util.concurrent.CountDownLatch +import java.util.concurrent.Executors +import java.util.concurrent.TimeUnit import java.util.function.Consumer import kotlin.test.assertEquals import kotlin.test.assertFailsWith @@ -237,8 +239,12 @@ class ChatClientExtensionsToolApprovalTest { * `toolApprovalTimeoutMs`. Otherwise a caller clearing UI state in `onToolApprovalTimedOut` * would see it stay visible for up to twice the configured timeout. * - * Returns a future completed after a delay (not [CompletableFuture.completedFuture]) so an - * implementation that ignores it instead of calling `future.get()` returns too early to pass. + * Runs the call under test on a separate executor and coordinates via a latch/a + * test-controlled [CompletableFuture] — not elapsed-time margins — so this can't flake + * under a loaded CI worker: `future` only completes when *this test* completes it, so an + * implementation that actually awaits it is guaranteed to still be blocked (task not done) + * right after `onToolFinishedAwaitable` fires, while one that ignores the returned future + * completes almost immediately regardless of scheduling delays. */ @Test fun `approval timeout invokes onToolApprovalTimedOut before awaiting onToolFinishedAwaitable, both exactly once`() { @@ -248,7 +254,9 @@ class ChatClientExtensionsToolApprovalTest { var toolFinishedCalls = 0 var timedOutCalls = 0 var toolFinishedHasFailed: Boolean? = null - val futureCompleted = AtomicBoolean(false) + + val awaitableInvoked = CountDownLatch(1) + val future = CompletableFuture() val chatClient = object : ChatClient { override fun sendMessageStreaming(userContents: List): TokenStream = FakeTokenStream { @@ -259,36 +267,53 @@ class ChatClientExtensionsToolApprovalTest { override fun sendMessage(prompt: String): String = error("not used") } - assertFailsWith { - chatClient.sendStreamingMessageWithCallback( - userContents = listOf(), - resolvedTools = listOf(requireApprovalConfig("toolTimeout")), - // Never invokes approve/deny — forces the timeout branch. - onToolApprovalRequired = { _, _, _, _ -> }, - onToolFinishedAwaitable = { _, _, _, hasFailed -> - toolFinishedCalls++ - toolFinishedHasFailed = hasFailed - invocationOrder += "onToolFinishedAwaitable" - val future = CompletableFuture() - // Completed only after a delay, from a background thread — forces a caller - // that actually awaits this future to block until then; one that ignores it - // (the bug this test guards against) returns well before futureCompleted flips. - Thread { - Thread.sleep(150) - futureCompleted.set(true) - future.complete(Unit) - }.start() - future - }, - onToolApprovalTimedOut = { - timedOutCalls++ - invocationOrder += "onToolApprovalTimedOut" - }, - // Large enough that the 150ms delayed completion above always wins the race - // against this timeout inside notifyToolFinishedAndAwaitIfPossible's get() call. - toolApprovalTimeoutMs = 300L, - chatMemory = chatMemory, + val executor = Executors.newSingleThreadExecutor() + try { + val callTask = executor.submit { + assertFailsWith { + chatClient.sendStreamingMessageWithCallback( + userContents = listOf(), + resolvedTools = listOf(requireApprovalConfig("toolTimeout")), + // Never invokes approve/deny — forces the timeout branch. + onToolApprovalRequired = { _, _, _, _ -> }, + onToolFinishedAwaitable = { _, _, _, hasFailed -> + toolFinishedCalls++ + toolFinishedHasFailed = hasFailed + invocationOrder += "onToolFinishedAwaitable" + awaitableInvoked.countDown() + future + }, + onToolApprovalTimedOut = { + timedOutCalls++ + invocationOrder += "onToolApprovalTimedOut" + }, + // Large relative to the latch/executor coordination below: the completion + // await inside the implementation must not self-timeout before this test + // explicitly completes `future` — see the isDone check further down. + toolApprovalTimeoutMs = 1000L, + chatMemory = chatMemory, + ) + } + } + + assertTrue( + awaitableInvoked.await(5, TimeUnit.SECONDS), + "Timed out waiting for onToolFinishedAwaitable to be invoked", ) + + // `future` is still incomplete at this point (we only complete it below), so a caller + // that actually awaits it must still be blocked — not yet done — no matter how slow + // this check runs under CI load. + assertTrue( + !callTask.isDone, + "Expected the call to still be blocked awaiting onToolFinishedAwaitable's future, " + + "but it already completed — meaning the future was never actually awaited", + ) + + future.complete(Unit) + callTask.get(5, TimeUnit.SECONDS) + } finally { + executor.shutdownNow() } assertEquals(1, toolFinishedCalls, "Expected onToolFinishedAwaitable to be invoked exactly once, got $toolFinishedCalls") @@ -299,11 +324,6 @@ class ChatClientExtensionsToolApprovalTest { invocationOrder, "Expected onToolApprovalTimedOut to fire immediately on timeout, before awaiting onToolFinishedAwaitable's completion future", ) - assertTrue( - futureCompleted.get(), - "Expected the call to block until onToolFinishedAwaitable's future completed — " + - "it returned before the delayed completion ran, meaning the future was never actually awaited", - ) } /** From 665ac666c30f1ab81dc9cf289e90421c9f9e6a78 Mon Sep 17 00:00:00 2001 From: Huy Nguyen <290170845+huynguyen20011@users.noreply.github.com> Date: Thu, 8 Oct 2026 20:02:11 -0700 Subject: [PATCH 4/6] Update --- .../io/askimo/ui/session/SessionManager.kt | 99 ++++++++++--------- 1 file changed, 50 insertions(+), 49 deletions(-) diff --git a/desktop-shared/src/main/kotlin/io/askimo/ui/session/SessionManager.kt b/desktop-shared/src/main/kotlin/io/askimo/ui/session/SessionManager.kt index fee121ea1..0a7590f0a 100644 --- a/desktop-shared/src/main/kotlin/io/askimo/ui/session/SessionManager.kt +++ b/desktop-shared/src/main/kotlin/io/askimo/ui/session/SessionManager.kt @@ -8,6 +8,7 @@ import androidx.compose.runtime.getValue import androidx.compose.runtime.mutableStateOf import androidx.compose.runtime.setValue import androidx.compose.runtime.snapshots.Snapshot +import dev.langchain4j.data.message.Content import io.askimo.core.chat.domain.ChatMessage import io.askimo.core.chat.domain.ChatSession import io.askimo.core.chat.domain.Project @@ -38,6 +39,7 @@ import io.askimo.core.vision.ImageProcessor import io.askimo.ui.chat.ChatViewModel import io.askimo.ui.chat.CreationMode import kotlinx.coroutines.CoroutineScope +import kotlinx.coroutines.CoroutineStart import kotlinx.coroutines.Dispatchers import kotlinx.coroutines.Job import kotlinx.coroutines.SupervisorJob @@ -308,47 +310,31 @@ class SessionManager( } } - // Check if this session already has an active (not yet complete) stream - val existingThread = activeThreads[sessionId] - if (existingThread != null && !existingThread.isComplete.value) { - log.warn("Session $sessionId already has an active stream") - return null - } - // If the thread is complete but not yet cleaned up by the ViewModel, remove it now - // so a new stream can start (e.g. user retries before the completion handler fires) - if (existingThread != null) { - activeThreads.remove(sessionId) - log.debug("Removed stale completed thread for session $sessionId before starting new stream") - } - - // Check global stream limit + // Soft cap — approximate, not atomic with the reservation below. if (activeThreads.size >= MAX_CONCURRENT_STREAMS) { log.warn("Max concurrent streams ($MAX_CONCURRENT_STREAMS) reached") return null } val threadId = "${sessionId}_${System.currentTimeMillis()}" - val thread = StreamingThread( threadId = threadId, sessionId = sessionId, - job = Job(), + job = Job(), // placeholder, replaced below before reservation _timeline = MutableStateFlow(emptyList()), _isComplete = MutableStateFlow(false), _hasFailed = MutableStateFlow(false), ) - // Prepare context and save user message to DB - val promptWithContext = chatSessionService.prepareContextAndGetPromptForChat(sessionId, userMessage, willSaveUserMessage) - log.debug("Saved prompt for session $sessionId, starting streaming") + // Assigned after reservation succeeds — the LAZY coroutine below only reads this once started. + lateinit var promptWithContext: List - thread.job = streamingScope.launch { - // Tracks every fire-and-forget timeline-mutating coroutine (onToken, - // onThinkingToken, onToolStarted, onToolFinished) launched while - // sendStreamingMessageWithCallback is in-flight below. Joined — both on success - // and on any exception — before reading thread.timeline/getCurrentContent(), so - // persistence never observes a timeline missing the last chunk(s)/tool update(s) - // still mid-flight the instant the call returns or throws. + // LAZY: gives a real, stable, cancellable Job before the coroutine body runs — without + // passing an external Job into launch()'s context, which breaks structured concurrency. + thread.job = streamingScope.launch(start = CoroutineStart.LAZY) { + // Tracks fire-and-forget timeline-mutating jobs (onToken, onThinkingToken, + // onToolStarted, onToolFinished); joined before any thread.timeline read below so + // persistence never misses an update still in flight. val pendingUpdateJobs = ConcurrentLinkedQueue() try { if (mode is CreationMode.Chat) { @@ -384,17 +370,10 @@ class SessionManager( val fullResponse = try { val sessionContext = chatSessionService.getOrCreateContextForSession(sessionId) - // Tracks each tool call's in-flight markToolRunning job, keyed by - // (toolName, arguments) rather than toolName alone — otherwise two - // parallel invocations of the SAME tool (same name, different - // arguments) would collide: the second's markToolRunning would be - // deduped as "already running" and markToolDone could join/flip the - // wrong call's entry. markToolDone joins this first — onToolStarted/ - // onToolFinished each launch their update on streamingScope - // independently, and without this, a delayed RUNNING-append coroutine - // can land *after* a faster DONE-append one, leaving the timeline's - // last entry for that tool RUNNING (collapsedEffectiveTools treats it - // as effective) even though the tool already finished. + // Keyed by (toolName, arguments), not toolName alone, so two parallel + // calls to the same tool don't collide. markToolDone joins the matching + // start job first so a delayed RUNNING-append can't land after a faster + // DONE-append and leave the tool stuck RUNNING. val pendingToolStartJobs = ConcurrentHashMap, Job>() sessionContext.chatClient .sendStreamingMessageWithCallback( @@ -487,9 +466,7 @@ class SessionManager( if (needsMessageCorrelation) ProxyChatContext.clear() } - // Ensure every timeline mutation triggered during the call above has - // actually landed before persisting/reading the timeline below — see - // pendingUpdateJobs' declaration comment. + // Wait for in-flight timeline updates before reading it below. pendingUpdateJobs.joinAll() val savedMessage = chatSessionService.saveAiResponse( @@ -552,9 +529,7 @@ class SessionManager( log.error("Error while sending message to chat session $sessionId", e) thread.markFailed() - // Same join as the success path above — the exception can propagate while - // onToken/onToolStarted/onToolFinished/onThinkingToken updates are still - // in-flight on streamingScope. + // Same reason as the success path — updates may still be in flight. pendingUpdateJobs.joinAll() val partialResponse = thread.getCurrentContent() @@ -604,12 +579,38 @@ class SessionManager( } } - // Registered only now that thread.job holds the real coroutine handle (not the - // placeholder Job() from construction above) — otherwise a stopStream() call landing - // in the gap between registration and this assignment would cancel the inert - // placeholder instead of the actual stream. - activeThreads[sessionId] = thread - log.debug("Streaming thread $threadId for session $sessionId started. Active streams: ${activeThreads.size}") + // Atomically reserve this session's slot: reject if a stream is already active, + // otherwise overwrite a stale/completed entry with `thread`. Keeps the check-then-insert + // from being a TOCTOU race where two concurrent sendMessage() calls both succeed. + var reserved = true + activeThreads.compute(sessionId) { _, existing -> + if (existing != null && !existing.isComplete.value) { + reserved = false + existing + } else { + thread + } + } + if (!reserved) { + log.warn("Session $sessionId already has an active stream") + thread.job.cancel() + return null + } + log.debug("Streaming thread $threadId for session $sessionId reserved. Active streams: ${activeThreads.size}") + + // Release the reservation if setup fails, so a retry isn't blocked by a dead thread. + try { + promptWithContext = chatSessionService.prepareContextAndGetPromptForChat(sessionId, userMessage, willSaveUserMessage) + } catch (e: Exception) { + activeThreads.remove(sessionId, thread) + thread.job.cancel() + throw e + } + log.debug("Saved prompt for session $sessionId, starting streaming") + + // Starts the coroutine only now — cancelling it earlier (e.g. via a racing stopStream()) + // would have made this a no-op instead. + thread.job.start() return threadId } From 1e3cc34c4d7121a6eec901759a7c5d2540c6ce7a Mon Sep 17 00:00:00 2001 From: Huy Nguyen <290170845+huynguyen20011@users.noreply.github.com> Date: Thu, 8 Oct 2026 20:26:13 -0700 Subject: [PATCH 5/6] Update --- .../io/askimo/ui/session/SessionManager.kt | 51 +++++----- .../core/providers/ChatClientExtensions.kt | 44 ++++++--- .../ChatClientExtensionsToolApprovalTest.kt | 95 ++++++++++++++++++- .../io/askimo/core/chat/dto/ToolCallInfo.kt | 8 +- 4 files changed, 156 insertions(+), 42 deletions(-) diff --git a/desktop-shared/src/main/kotlin/io/askimo/ui/session/SessionManager.kt b/desktop-shared/src/main/kotlin/io/askimo/ui/session/SessionManager.kt index 0a7590f0a..5ba03ee17 100644 --- a/desktop-shared/src/main/kotlin/io/askimo/ui/session/SessionManager.kt +++ b/desktop-shared/src/main/kotlin/io/askimo/ui/session/SessionManager.kt @@ -206,37 +206,34 @@ class SessionManager( } } - /** Appends a RUNNING tool entry at its true chronological position. Deduplicates: no-op if already tracked as RUNNING. */ - suspend fun markToolRunning(toolName: String, arguments: String?) { + /** Appends a RUNNING tool entry at its true chronological position. Deduplicates: no-op if already tracked as RUNNING under [id]. */ + suspend fun markToolRunning(id: String, toolName: String, arguments: String?) { mutex.withLock { val alreadyRunning = _timeline.value.any { it is TurnTimelineEntry.Tool && - it.toolCall.toolName == toolName && - it.toolCall.arguments == arguments && + it.toolCall.requestId == id && it.toolCall.status == ToolCallStatus.RUNNING } if (!alreadyRunning) { _timeline.value += TurnTimelineEntry.Tool( // Full, untruncated arguments — only the persisted copy is capped // (see truncatedForStorage()). - ToolCallInfo(toolName = toolName, status = ToolCallStatus.RUNNING, arguments = arguments), + ToolCallInfo(toolName = toolName, status = ToolCallStatus.RUNNING, arguments = arguments, requestId = id), ) } } } /** Flips the matching RUNNING tool entry to DONE in-place, preserving its original chronological position. */ - suspend fun markToolDone(toolName: String, arguments: String?, result: String?, hasFailed: Boolean) { + suspend fun markToolDone(id: String, toolName: String, arguments: String?, result: String?, hasFailed: Boolean) { mutex.withLock { val list = _timeline.value - // Matched on (toolName, arguments) rather than toolName alone — otherwise two - // parallel invocations of the SAME tool (same name, different arguments) would - // collide: the second markToolRunning would be deduped as "already running" and - // this lookup would flip the wrong call's entry to DONE. + // Matched on the request id — the model's unique id for this specific + // invocation — rather than (toolName, arguments), which collides when the + // model calls the same tool twice with identical arguments in one parallel batch. val idx = list.indexOfLast { it is TurnTimelineEntry.Tool && - it.toolCall.toolName == toolName && - it.toolCall.arguments == arguments && + it.toolCall.requestId == id && it.toolCall.status == ToolCallStatus.RUNNING } // Preserve the RUNNING entry's original start time (instead of "now") so a @@ -252,6 +249,7 @@ class SessionManager( result = result, hasFailed = hasFailed, startedAtMillis = startedAtMillis, + requestId = id, ), ) _timeline.value = if (idx >= 0) { @@ -370,11 +368,12 @@ class SessionManager( val fullResponse = try { val sessionContext = chatSessionService.getOrCreateContextForSession(sessionId) - // Keyed by (toolName, arguments), not toolName alone, so two parallel - // calls to the same tool don't collide. markToolDone joins the matching - // start job first so a delayed RUNNING-append can't land after a faster - // DONE-append and leave the tool stuck RUNNING. - val pendingToolStartJobs = ConcurrentHashMap, Job>() + // Keyed by the tool request id (not toolName/arguments) so two parallel + // calls to the same tool with identical arguments don't collide. + // markToolDone joins the matching start job first so a delayed + // RUNNING-append can't land after a faster DONE-append and leave the + // tool stuck RUNNING. + val pendingToolStartJobs = ConcurrentHashMap() sessionContext.chatClient .sendStreamingMessageWithCallback( projectId = projectId, @@ -395,17 +394,17 @@ class SessionManager( capturedDurationMs = durationMs log.debug("Token usage for session $sessionId: input=$input, output=$output, total=$total, duration=${durationMs}ms") }, - onToolStarted = { toolName, arguments -> + onToolStarted = { id, toolName, arguments -> val job = streamingScope.launch { - thread.markToolRunning(toolName, arguments) + thread.markToolRunning(id, toolName, arguments) } - pendingToolStartJobs[toolName to arguments] = job + pendingToolStartJobs[id] = job pendingUpdateJobs += job }, - onToolFinished = { toolName, arguments, result, hasFailed -> + onToolFinished = { id, toolName, arguments, result, hasFailed -> pendingUpdateJobs += streamingScope.launch { - pendingToolStartJobs.remove(toolName to arguments)?.join() - thread.markToolDone(toolName, arguments, result, hasFailed) + pendingToolStartJobs.remove(id)?.join() + thread.markToolDone(id, toolName, arguments, result, hasFailed) } }, onThinkingToken = { token -> @@ -446,12 +445,12 @@ class SessionManager( // Awaitable companion to onToolFinished: used only on approval // timeout/denial so markToolDone() applies before the exception // propagates — otherwise the tool could be persisted as RUNNING. - onToolFinishedAwaitable = { toolName, arguments, result, hasFailed -> + onToolFinishedAwaitable = { id, toolName, arguments, result, hasFailed -> CompletableFuture().apply { streamingScope.launch { try { - pendingToolStartJobs.remove(toolName to arguments)?.join() - thread.markToolDone(toolName, arguments, result, hasFailed) + pendingToolStartJobs.remove(id)?.join() + thread.markToolDone(id, toolName, arguments, result, hasFailed) } finally { complete(Unit) } diff --git a/shared-jvm/src/main/kotlin/io/askimo/core/providers/ChatClientExtensions.kt b/shared-jvm/src/main/kotlin/io/askimo/core/providers/ChatClientExtensions.kt index 7562568fe..e8bc800a8 100644 --- a/shared-jvm/src/main/kotlin/io/askimo/core/providers/ChatClientExtensions.kt +++ b/shared-jvm/src/main/kotlin/io/askimo/core/providers/ChatClientExtensions.kt @@ -101,6 +101,7 @@ internal sealed class StreamingErrorResult { * * @param onToolFinished Legacy, non-awaitable callback — invoked without waiting if [onToolFinishedAwaitable] is null. * @param onToolFinishedAwaitable Optional awaitable variant — when present, takes precedence and is awaited. + * @param id [ToolExecutionRequest.id] of the invocation this update belongs to * @param toolName Name of the tool * @param arguments Tool arguments * @param result Tool result @@ -108,8 +109,9 @@ internal sealed class StreamingErrorResult { * @param timeoutMs Timeout for waiting on the completion future */ private fun notifyToolFinishedAndAwaitIfPossible( - onToolFinished: ((String, String?, String?, Boolean) -> Unit)?, - onToolFinishedAwaitable: ((String, String?, String?, Boolean) -> CompletableFuture?)?, + onToolFinished: ((String, String, String?, String?, Boolean) -> Unit)?, + onToolFinishedAwaitable: ((String, String, String?, String?, Boolean) -> CompletableFuture?)?, + id: String, toolName: String, arguments: String?, result: String?, @@ -117,14 +119,14 @@ private fun notifyToolFinishedAndAwaitIfPossible( timeoutMs: Long, ) { if (onToolFinishedAwaitable != null) { - val completion = onToolFinishedAwaitable.invoke(toolName, arguments, result, hasFailed) + val completion = onToolFinishedAwaitable.invoke(id, toolName, arguments, result, hasFailed) try { completion?.get(timeoutMs, TimeUnit.MILLISECONDS) } catch (e: Exception) { log.warn("Tool completion callback timed out for $toolName", e) } } else { - onToolFinished?.invoke(toolName, arguments, result, hasFailed) + onToolFinished?.invoke(id, toolName, arguments, result, hasFailed) } } @@ -208,8 +210,13 @@ fun ChatClient.sendStreamingMessageWithCallback( onToken: (String) -> Unit = {}, onFollowUpSuggestion: ((FollowUpSuggestion) -> Unit)? = null, onTokenUsage: ((inputTokens: Int, outputTokens: Int, totalTokens: Int, durationMs: Long) -> Unit)? = null, - onToolStarted: ((toolName: String, arguments: String?) -> Unit)? = null, - onToolFinished: ((toolName: String, arguments: String?, result: String?, hasFailed: Boolean) -> Unit)? = null, + /** + * `id` is [ToolExecutionRequest.id] — the model's unique id for this specific invocation, not + * just `toolName`. Needed because the model can call the same tool twice with identical + * arguments in one parallel batch, which `toolName`+arguments alone can't disambiguate. + */ + onToolStarted: ((id: String, toolName: String, arguments: String?) -> Unit)? = null, + onToolFinished: ((id: String, toolName: String, arguments: String?, result: String?, hasFailed: Boolean) -> Unit)? = null, onThinkingToken: ((String) -> Unit)? = null, /** * Resolved [ToolConfig] list for the current session. @@ -249,7 +256,7 @@ fun ChatClient.sendStreamingMessageWithCallback( * is left without a matching `tool_result`, which Anthropic's API rejects on the next turn. */ chatMemory: ChatMemory? = null, - onToolFinishedAwaitable: ((toolName: String, arguments: String?, result: String?, hasFailed: Boolean) -> CompletableFuture?)? = null, + onToolFinishedAwaitable: ((id: String, toolName: String, arguments: String?, result: String?, hasFailed: Boolean) -> CompletableFuture?)? = null, ): String { val log = logger() @@ -337,9 +344,10 @@ fun ChatClient.sendStreamingMessageWithCallback( } done.countDown() }.beforeToolExecution { before -> + val id = before.request().id() val toolName = before.request().name() val arguments = before.request().arguments() - onToolStarted?.invoke(toolName, arguments) + onToolStarted?.invoke(id, toolName, arguments) // Lazily seed `pendingToolRequests` from the latest tool-calling // AiMessage already persisted to chatMemory — LangChain4j appends @@ -383,7 +391,11 @@ fun ChatClient.sendStreamingMessageWithCallback( ) if (!latch.await(toolApprovalTimeoutMs, TimeUnit.MILLISECONDS)) { val msg = LocalizationManager.getString("chat.tool.approval.timed_out", toolName) - pendingToolRequests.values.forEach { req -> + // Snapshot before clearing — every pending sibling needs its own + // completion notification, not just the current tool. Unioned with + // the current request in case chatMemory is null (unseeded map). + val terminatedRequests = (pendingToolRequests.values + before.request()).distinctBy { it.id() } + terminatedRequests.forEach { req -> chatMemory?.add(ToolExecutionResultMessage.from(req, msg)) } pendingToolRequests.clear() @@ -393,16 +405,22 @@ fun ChatClient.sendStreamingMessageWithCallback( // which can itself block for up to toolApprovalTimeoutMs and would // otherwise leave the approval banner visible for twice as long. onToolApprovalTimedOut?.invoke() - notifyToolFinishedAndAwaitIfPossible(onToolFinished, onToolFinishedAwaitable, toolName, arguments, msg, true, toolApprovalTimeoutMs) + terminatedRequests.forEach { req -> + notifyToolFinishedAndAwaitIfPossible(onToolFinished, onToolFinishedAwaitable, req.id(), req.name(), req.arguments(), msg, true, toolApprovalTimeoutMs) + } throw ToolExecutionException(toolName = toolName, errorDetails = msg) } if (!approved) { val msg = LocalizationManager.getString("chat.tool.approval.denied", toolName) - pendingToolRequests.values.forEach { req -> + // Same reasoning as the timeout branch above. + val terminatedRequests = (pendingToolRequests.values + before.request()).distinctBy { it.id() } + terminatedRequests.forEach { req -> chatMemory?.add(ToolExecutionResultMessage.from(req, msg)) } pendingToolRequests.clear() - notifyToolFinishedAndAwaitIfPossible(onToolFinished, onToolFinishedAwaitable, toolName, arguments, msg, true, toolApprovalTimeoutMs) + terminatedRequests.forEach { req -> + notifyToolFinishedAndAwaitIfPossible(onToolFinished, onToolFinishedAwaitable, req.id(), req.name(), req.arguments(), msg, true, toolApprovalTimeoutMs) + } throw ToolExecutionException(toolName = toolName, errorDetails = msg) } } @@ -413,7 +431,7 @@ fun ChatClient.sendStreamingMessageWithCallback( val arguments = tool.request().arguments() val result = tool.result() val hasFailed = tool.hasFailed() - onToolFinished?.invoke(toolName, arguments, result, hasFailed) + onToolFinished?.invoke(tool.request().id(), toolName, arguments, result, hasFailed) } .onError { e -> errorOccurred = true diff --git a/shared-jvm/src/test/kotlin/io/askimo/core/providers/ChatClientExtensionsToolApprovalTest.kt b/shared-jvm/src/test/kotlin/io/askimo/core/providers/ChatClientExtensionsToolApprovalTest.kt index e2da27bcf..b0b943963 100644 --- a/shared-jvm/src/test/kotlin/io/askimo/core/providers/ChatClientExtensionsToolApprovalTest.kt +++ b/shared-jvm/src/test/kotlin/io/askimo/core/providers/ChatClientExtensionsToolApprovalTest.kt @@ -233,6 +233,49 @@ class ChatClientExtensionsToolApprovalTest { assertTrue(requestB.id() in resultIds, "Expected a tool_result for requestB (${requestB.id()}), got: $resultIds") } + /** + * Companion to the test above: a sibling whose `beforeToolExecution` already fired + * `onToolStarted` (creating a RUNNING timeline entry) must also get an `onToolFinished` + * notification on timeout/denial — not just a synthetic `chatMemory` result — otherwise its + * RUNNING entry is never flipped to DONE/failed in the persisted timeline. + */ + @Test + fun `approval timeout notifies onToolFinished for every pending sibling request`() { + val chatMemory = FakeChatMemory() + val requestA = toolRequest("call_A", "toolA") + val requestB = toolRequest("call_B", "toolB") + val finishedIds = mutableListOf() + + val chatClient = object : ChatClient { + override fun sendMessageStreaming(userContents: List): TokenStream = FakeTokenStream { + chatMemory.add(AiMessage.from(listOf(requestA, requestB))) + // Simulates both siblings already having started (e.g. the model requested + // both in parallel and the framework invoked beforeToolExecution for each + // before either finished) — only requestB then blocks on approval and times out. + beforeToolExecutionHandler?.accept(beforeToolExecution(requestA)) + } + + override fun sendMessage(prompt: String): String = error("not used") + } + + assertFailsWith { + chatClient.sendStreamingMessageWithCallback( + userContents = listOf(), + resolvedTools = listOf(requireApprovalConfig("toolA"), requireApprovalConfig("toolB")), + onToolApprovalRequired = { _, _, _, _ -> }, + onToolFinished = { id, _, _, _, _ -> finishedIds += id }, + toolApprovalTimeoutMs = 50L, + chatMemory = chatMemory, + ) + } + + assertEquals( + setOf(requestA.id(), requestB.id()), + finishedIds.toSet(), + "Expected onToolFinished for every pending sibling request, got: $finishedIds", + ) + } + /** * `onToolApprovalTimedOut` must fire immediately on timeout — before awaiting * `onToolFinishedAwaitable`'s completion future, which can itself block for up to @@ -276,7 +319,7 @@ class ChatClientExtensionsToolApprovalTest { resolvedTools = listOf(requireApprovalConfig("toolTimeout")), // Never invokes approve/deny — forces the timeout branch. onToolApprovalRequired = { _, _, _, _ -> }, - onToolFinishedAwaitable = { _, _, _, hasFailed -> + onToolFinishedAwaitable = { _, _, _, _, hasFailed -> toolFinishedCalls++ toolFinishedHasFailed = hasFailed invocationOrder += "onToolFinishedAwaitable" @@ -326,6 +369,54 @@ class ChatClientExtensionsToolApprovalTest { ) } + /** + * `id` ([ToolExecutionRequest.id]) must correlate `onToolStarted`/`onToolFinished` pairs, not + * `(toolName, arguments)` — the model can issue the exact same tool with identical arguments + * twice in one parallel batch, which `(toolName, arguments)` can't disambiguate but each + * call's unique id can. + */ + @Test + fun `onToolStarted and onToolFinished correlate by id, not toolName plus arguments`() { + val chatMemory = FakeChatMemory() + // Same tool name AND same arguments — only the id differs between the two calls. + val requestA = toolRequest("call_1", "duplicateTool") + val requestB = toolRequest("call_2", "duplicateTool") + + val startedIds = mutableListOf() + val finishedIds = mutableListOf() + + val chatClient = object : ChatClient { + override fun sendMessageStreaming(userContents: List): TokenStream = FakeTokenStream { + chatMemory.add(AiMessage.from(listOf(requestA, requestB))) + beforeToolExecutionHandler?.accept(beforeToolExecution(requestA)) + beforeToolExecutionHandler?.accept(beforeToolExecution(requestB)) + toolExecutedHandler?.accept( + ToolExecution.builder().request(requestB).result("result-B") + .invocationContext(InvocationContext.builder().build()).build(), + ) + toolExecutedHandler?.accept( + ToolExecution.builder().request(requestA).result("result-A") + .invocationContext(InvocationContext.builder().build()).build(), + ) + completeResponseHandler?.accept( + ChatResponse.builder().aiMessage(AiMessage.from("done")).build(), + ) + } + + override fun sendMessage(prompt: String): String = error("not used") + } + + chatClient.sendStreamingMessageWithCallback( + userContents = listOf(), + onToolStarted = { id, _, _ -> startedIds += id }, + onToolFinished = { id, _, _, _, _ -> finishedIds += id }, + chatMemory = chatMemory, + ) + + assertEquals(listOf("call_1", "call_2"), startedIds, "Expected both distinct call ids to be reported as started") + assertEquals(listOf("call_2", "call_1"), finishedIds, "Expected both distinct call ids to be reported as finished, in completion order") + } + /** * Companion to the timeout test above: denial must still notify `onToolFinishedAwaitable` (so UI * state like a RUNNING tool entry gets flipped to DONE/failed) but must NOT invoke @@ -354,7 +445,7 @@ class ChatClientExtensionsToolApprovalTest { userContents = listOf(), resolvedTools = listOf(requireApprovalConfig("toolDeny")), onToolApprovalRequired = { _, _, _, deny -> deny() }, - onToolFinishedAwaitable = { _, _, _, hasFailed -> + onToolFinishedAwaitable = { _, _, _, _, hasFailed -> toolFinishedCalls++ toolFinishedHasFailed = hasFailed CompletableFuture.completedFuture(Unit) diff --git a/shared/src/commonMain/kotlin/io/askimo/core/chat/dto/ToolCallInfo.kt b/shared/src/commonMain/kotlin/io/askimo/core/chat/dto/ToolCallInfo.kt index 7c5e3aa0b..325d32cc3 100644 --- a/shared/src/commonMain/kotlin/io/askimo/core/chat/dto/ToolCallInfo.kt +++ b/shared/src/commonMain/kotlin/io/askimo/core/chat/dto/ToolCallInfo.kt @@ -34,7 +34,7 @@ enum class ToolCallStatus { * [MAX_FIELD_LENGTH] (see [truncated]) is applied only once, right before persisting a turn * (via [io.askimo.core.chat.dto.truncatedForStorage]) — some tools can otherwise produce * megabytes of raw text, bloating the persisted `content_json` column and any sync payload. - * Unlike [io.askimo.core.agent.ExternalAgent]'s `onToolCall` `detail` (a short display label, + * Unlike `ExternalAgent`'s `onToolCall` `detail` (a short display label, * never truncated), these fields retain enough raw content to stay useful for debugging/audit * once truncated — just not unbounded. * @@ -47,6 +47,11 @@ enum class ToolCallStatus { * elapsed timer while [status] is [ToolCallStatus.RUNNING] — without it, a long-running tool * gives no feedback that the AI is still working. Marked [Transient] — session-only, never * persisted, since it's meaningless for historical tool calls loaded back from `content_json`. + * @param requestId `ToolExecutionRequest.id`, the model's unique id for this specific + * invocation. Used only to match a RUNNING entry to its DONE update when the same tool is + * called more than once with identical arguments in one parallel batch — `(toolName, + * arguments)` alone can't distinguish those. Marked [Transient]: meaningless once reloaded + * from persisted `content_json`, and not every tool-calling model provider supplies one. */ @Serializable data class ToolCallInfo( @@ -56,6 +61,7 @@ data class ToolCallInfo( val result: String? = null, val hasFailed: Boolean = false, @Transient val startedAtMillis: Long = Clock.System.now().toEpochMilliseconds(), + @Transient val requestId: String? = null, ) { companion object { /** Max length of [arguments]/[result] before truncation — see [truncated]. */ From e7f215442306f97646e1cd85659c8b5b772db1fe Mon Sep 17 00:00:00 2001 From: Huy Nguyen <290170845+huynguyen20011@users.noreply.github.com> Date: Thu, 8 Oct 2026 20:41:55 -0700 Subject: [PATCH 6/6] Update per feedbacks --- .../core/providers/ChatClientExtensions.kt | 104 ++++++++++++------ 1 file changed, 73 insertions(+), 31 deletions(-) diff --git a/shared-jvm/src/main/kotlin/io/askimo/core/providers/ChatClientExtensions.kt b/shared-jvm/src/main/kotlin/io/askimo/core/providers/ChatClientExtensions.kt index e8bc800a8..c4eabcbd7 100644 --- a/shared-jvm/src/main/kotlin/io/askimo/core/providers/ChatClientExtensions.kt +++ b/shared-jvm/src/main/kotlin/io/askimo/core/providers/ChatClientExtensions.kt @@ -118,15 +118,14 @@ private fun notifyToolFinishedAndAwaitIfPossible( hasFailed: Boolean, timeoutMs: Long, ) { - if (onToolFinishedAwaitable != null) { - val completion = onToolFinishedAwaitable.invoke(id, toolName, arguments, result, hasFailed) - try { - completion?.get(timeoutMs, TimeUnit.MILLISECONDS) - } catch (e: Exception) { - log.warn("Tool completion callback timed out for $toolName", e) + try { + if (onToolFinishedAwaitable != null) { + onToolFinishedAwaitable.invoke(id, toolName, arguments, result, hasFailed)?.get(timeoutMs, TimeUnit.MILLISECONDS) + } else { + onToolFinished?.invoke(id, toolName, arguments, result, hasFailed) } - } else { - onToolFinished?.invoke(id, toolName, arguments, result, hasFailed) + } catch (e: Exception) { + log.warn("Tool completion callback failed or timed out for $toolName", e) } } @@ -308,11 +307,19 @@ fun ChatClient.sendStreamingMessageWithCallback( sendMessageStreaming(userContents) .onPartialResponse { chunk -> sb.append(chunk) - onToken(chunk) + try { + onToken(chunk) + } catch (e: Exception) { + log.warn("onToken callback threw", e) + } }.onPartialThinking { thinking -> val text = thinking.text() if (!text.isNullOrEmpty()) { - onThinkingToken?.invoke(text) + try { + onThinkingToken?.invoke(text) + } catch (e: Exception) { + log.warn("onThinkingToken callback threw", e) + } } } .onCompleteResponse { response -> @@ -322,12 +329,16 @@ fun ChatClient.sendStreamingMessageWithCallback( // Fire per-message token usage callback before counting down if (onTokenUsage != null && tokenUsage != null) { val duration = System.currentTimeMillis() - streamStartTime - onTokenUsage( - tokenUsage.inputTokenCount() ?: 0, - tokenUsage.outputTokenCount() ?: 0, - tokenUsage.totalTokenCount() ?: 0, - duration, - ) + try { + onTokenUsage( + tokenUsage.inputTokenCount() ?: 0, + tokenUsage.outputTokenCount() ?: 0, + tokenUsage.totalTokenCount() ?: 0, + duration, + ) + } catch (e: Exception) { + log.warn("onTokenUsage callback threw", e) + } } if (GeneratedImageHelper.hasGeneratedImages(aiMessage)) { @@ -338,7 +349,11 @@ fun ChatClient.sendStreamingMessageWithCallback( val mimeType = image.mimeType() ?: "image/png" val markdownImage = "\n![Generated Image](data:$mimeType;base64,$base64Data)\n" sb.append(markdownImage) - onToken(markdownImage) + try { + onToken(markdownImage) + } catch (e: Exception) { + log.warn("onToken callback threw", e) + } } } } @@ -347,7 +362,11 @@ fun ChatClient.sendStreamingMessageWithCallback( val id = before.request().id() val toolName = before.request().name() val arguments = before.request().arguments() - onToolStarted?.invoke(id, toolName, arguments) + try { + onToolStarted?.invoke(id, toolName, arguments) + } catch (e: Exception) { + log.warn("onToolStarted callback threw for $toolName", e) + } // Lazily seed `pendingToolRequests` from the latest tool-calling // AiMessage already persisted to chatMemory — LangChain4j appends @@ -380,15 +399,22 @@ fun ChatClient.sendStreamingMessageWithCallback( if (effectivePolicy == ToolApprovalPolicy.REQUIRE_APPROVAL) { val latch = CountDownLatch(1) var approved = false - onToolApprovalRequired.invoke( - toolName, - arguments, - { - approved = true - latch.countDown() - }, - { latch.countDown() }, - ) + try { + onToolApprovalRequired.invoke( + toolName, + arguments, + { + approved = true + latch.countDown() + }, + { latch.countDown() }, + ) + } catch (e: Exception) { + // Treat a throwing callback as an implicit deny (fail-safe) + // rather than hanging for the full approval timeout. + log.warn("onToolApprovalRequired callback threw for $toolName — treating as denied", e) + latch.countDown() + } if (!latch.await(toolApprovalTimeoutMs, TimeUnit.MILLISECONDS)) { val msg = LocalizationManager.getString("chat.tool.approval.timed_out", toolName) // Snapshot before clearing — every pending sibling needs its own @@ -404,7 +430,11 @@ fun ChatClient.sendStreamingMessageWithCallback( // immediately — not after awaiting onToolFinishedAwaitable below, // which can itself block for up to toolApprovalTimeoutMs and would // otherwise leave the approval banner visible for twice as long. - onToolApprovalTimedOut?.invoke() + try { + onToolApprovalTimedOut?.invoke() + } catch (e: Exception) { + log.warn("onToolApprovalTimedOut callback threw", e) + } terminatedRequests.forEach { req -> notifyToolFinishedAndAwaitIfPossible(onToolFinished, onToolFinishedAwaitable, req.id(), req.name(), req.arguments(), msg, true, toolApprovalTimeoutMs) } @@ -431,7 +461,11 @@ fun ChatClient.sendStreamingMessageWithCallback( val arguments = tool.request().arguments() val result = tool.result() val hasFailed = tool.hasFailed() - onToolFinished?.invoke(tool.request().id(), toolName, arguments, result, hasFailed) + try { + onToolFinished?.invoke(tool.request().id(), toolName, arguments, result, hasFailed) + } catch (e: Exception) { + log.warn("onToolFinished callback threw for $toolName", e) + } } .onError { e -> errorOccurred = true @@ -482,7 +516,11 @@ fun ChatClient.sendStreamingMessageWithCallback( } terminalErrorMessage = message sb.append(message) - onToken(message) + try { + onToken(message) + } catch (e2: Exception) { + log.warn("onToken callback threw", e2) + } done.countDown() } }.start() @@ -509,7 +547,11 @@ fun ChatClient.sendStreamingMessageWithCallback( if (followUpSuggestion != null) { log.debug("Detected follow-up opportunity (Stage 2): ${followUpSuggestion.question}") - onFollowUpSuggestion(followUpSuggestion) + try { + onFollowUpSuggestion(followUpSuggestion) + } catch (e: Exception) { + log.warn("onFollowUpSuggestion callback threw", e) + } } }