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..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 @@ -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 @@ -45,13 +47,16 @@ 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 import kotlinx.coroutines.sync.withLock 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 /** @@ -201,28 +206,35 @@ 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.status == ToolCallStatus.RUNNING + it is TurnTimelineEntry.Tool && + 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 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.status == ToolCallStatus.RUNNING + it is TurnTimelineEntry.Tool && + it.toolCall.requestId == id && + 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. @@ -237,6 +249,7 @@ class SessionManager( result = result, hasFailed = hasFailed, startedAtMillis = startedAtMillis, + requestId = id, ), ) _timeline.value = if (idx >= 0) { @@ -295,46 +308,32 @@ 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), ) - // Register this thread - activeThreads[sessionId] = thread - - log.debug("Streaming thread $threadId for session $sessionId started. Active streams: ${activeThreads.size}") + // Assigned after reservation succeeds — the LAZY coroutine below only reads this once started. + lateinit var promptWithContext: List - // 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 { + // 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) { var capturedInputTokens: Int? = null @@ -369,13 +368,19 @@ class SessionManager( val fullResponse = try { val sessionContext = chatSessionService.getOrCreateContextForSession(sessionId) + // 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, userContents = promptWithContext, enabledServerIds = enabledServerIds, onToken = { token -> - streamingScope.launch { + pendingUpdateJobs += streamingScope.launch { thread.appendChunk(token) } }, @@ -389,18 +394,21 @@ class SessionManager( capturedDurationMs = durationMs log.debug("Token usage for session $sessionId: input=$input, output=$output, total=$total, duration=${durationMs}ms") }, - onToolStarted = { toolName, arguments -> - streamingScope.launch { - thread.markToolRunning(toolName, arguments) + onToolStarted = { id, toolName, arguments -> + val job = streamingScope.launch { + thread.markToolRunning(id, toolName, arguments) } + pendingToolStartJobs[id] = job + pendingUpdateJobs += job }, - onToolFinished = { toolName, arguments, result, hasFailed -> - streamingScope.launch { - thread.markToolDone(toolName, arguments, result, hasFailed) + onToolFinished = { id, toolName, arguments, result, hasFailed -> + pendingUpdateJobs += streamingScope.launch { + pendingToolStartJobs.remove(id)?.join() + thread.markToolDone(id, toolName, arguments, result, hasFailed) } }, onThinkingToken = { token -> - streamingScope.launch { + pendingUpdateJobs += streamingScope.launch { thread.appendThinkingChunk(token) } }, @@ -431,6 +439,24 @@ 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 = { id, toolName, arguments, result, hasFailed -> + CompletableFuture().apply { + streamingScope.launch { + try { + pendingToolStartJobs.remove(id)?.join() + thread.markToolDone(id, toolName, arguments, result, hasFailed) + } finally { + complete(Unit) + } + } + } + }, chatMemory = sessionContext.memory, ) } finally { @@ -439,6 +465,9 @@ class SessionManager( if (needsMessageCorrelation) ProxyChatContext.clear() } + // Wait for in-flight timeline updates before reading it below. + pendingUpdateJobs.joinAll() + val savedMessage = chatSessionService.saveAiResponse( sessionId = sessionId, response = fullResponse, @@ -499,6 +528,9 @@ class SessionManager( log.error("Error while sending message to chat session $sessionId", e) thread.markFailed() + // Same reason as the success path — updates may still be in flight. + pendingUpdateJobs.joinAll() + val partialResponse = thread.getCurrentContent() val failedResponse = if (e is ConfigurationErrorException) { e.displayMessage @@ -546,6 +578,39 @@ class SessionManager( } } + // 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 } 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..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 @@ -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,45 @@ 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 id [ToolExecutionRequest.id] of the invocation this update belongs to + * @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?, String?, Boolean) -> Unit)?, + onToolFinishedAwaitable: ((String, String, String?, String?, Boolean) -> CompletableFuture?)?, + id: String, + toolName: String, + arguments: String?, + result: String?, + hasFailed: Boolean, + timeoutMs: Long, +) { + try { + if (onToolFinishedAwaitable != null) { + onToolFinishedAwaitable.invoke(id, toolName, arguments, result, hasFailed)?.get(timeoutMs, TimeUnit.MILLISECONDS) + } else { + onToolFinished?.invoke(id, toolName, arguments, result, hasFailed) + } + } catch (e: Exception) { + log.warn("Tool completion callback failed or timed out for $toolName", e) + } +} + /** * Pure classification function: maps a streaming [Throwable] to a [StreamingErrorResult]. * @@ -153,6 +196,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( @@ -162,8 +209,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. @@ -183,6 +235,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 +255,7 @@ fun ChatClient.sendStreamingMessageWithCallback( * is left without a matching `tool_result`, which Anthropic's API rejects on the next turn. */ chatMemory: ChatMemory? = null, + onToolFinishedAwaitable: ((id: String, toolName: String, arguments: String?, result: String?, hasFailed: Boolean) -> CompletableFuture?)? = null, ): String { val log = logger() @@ -247,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 -> @@ -261,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)) { @@ -277,15 +349,24 @@ 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) + } } } } done.countDown() }.beforeToolExecution { before -> + val id = before.request().id() val toolName = before.request().name() val arguments = before.request().arguments() - onToolStarted?.invoke(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 @@ -318,29 +399,58 @@ 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) - 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() + // Neither approve() nor deny() fired, so the caller's pending-approval + // 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. + 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) + } 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() + terminatedRequests.forEach { req -> + notifyToolFinishedAndAwaitIfPossible(onToolFinished, onToolFinishedAwaitable, req.id(), req.name(), req.arguments(), msg, true, toolApprovalTimeoutMs) + } throw ToolExecutionException(toolName = toolName, errorDetails = msg) } } @@ -351,7 +461,11 @@ fun ChatClient.sendStreamingMessageWithCallback( val arguments = tool.request().arguments() val result = tool.result() val hasFailed = tool.hasFailed() - onToolFinished?.invoke(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 @@ -402,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() @@ -429,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) + } } } 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..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 @@ -30,7 +30,12 @@ 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.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 import kotlin.test.assertTrue @@ -227,4 +232,233 @@ 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") } + + /** + * 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 + * `toolApprovalTimeoutMs`. Otherwise a caller clearing UI state in `onToolApprovalTimedOut` + * would see it stay visible for up to twice the configured timeout. + * + * 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`() { + val chatMemory = FakeChatMemory() + val request = toolRequest("call_solo", "toolTimeout") + val invocationOrder = mutableListOf() + var toolFinishedCalls = 0 + var timedOutCalls = 0 + var toolFinishedHasFailed: Boolean? = null + + val awaitableInvoked = CountDownLatch(1) + val future = CompletableFuture() + + 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") + } + + 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") + 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("onToolApprovalTimedOut", "onToolFinishedAwaitable"), + invocationOrder, + "Expected onToolApprovalTimedOut to fire immediately on timeout, before awaiting onToolFinishedAwaitable's completion future", + ) + } + + /** + * `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 + * `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") + } } 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]. */