Skip to content

Commit 791fa7f

Browse files
kirill-markinkirill-markin
committed
Harden chat stop with run ids
1 parent f106e96 commit 791fa7f

31 files changed

Lines changed: 821 additions & 113 deletions

File tree

‎apps/android/data/local/src/main/java/com/flashcardsopensourceapp/data/local/ai/AiChatRemoteService.kt‎

Lines changed: 26 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,7 @@ import com.flashcardsopensourceapp.data.local.model.AiChatServerConfig
2222
import com.flashcardsopensourceapp.data.local.model.AiChatToolCall
2323
import com.flashcardsopensourceapp.data.local.model.AiChatToolCallStatus
2424
import com.flashcardsopensourceapp.data.local.model.AiChatStartRunRequest
25+
import com.flashcardsopensourceapp.data.local.model.AiChatStopRunRequest
2526
import com.flashcardsopensourceapp.data.local.model.AiToolCallRequest
2627
import com.flashcardsopensourceapp.data.local.model.AiChatBootstrapResponse
2728
import com.flashcardsopensourceapp.data.local.model.AiChatLiveEvent
@@ -249,8 +250,7 @@ class AiChatRemoteService(
249250
suspend fun stopRun(
250251
apiBaseUrl: String,
251252
authorizationHeader: String,
252-
sessionId: String,
253-
workspaceId: String?
253+
request: AiChatStopRunRequest
254254
): AiChatStopRunResponse = withContext(dispatchers.io) {
255255
val connection = openConnection(
256256
apiBaseUrl = apiBaseUrl,
@@ -264,11 +264,7 @@ class AiChatRemoteService(
264264
connection.doOutput = true
265265
connection.outputStream.use { outputStream ->
266266
outputStream.write(
267-
putOptionalWorkspaceId(
268-
payload = JSONObject()
269-
.put("sessionId", sessionId),
270-
workspaceId = workspaceId
271-
).toString().toByteArray(StandardCharsets.UTF_8)
267+
encodeStopRunRequest(request = request).toString().toByteArray(StandardCharsets.UTF_8)
272268
)
273269
}
274270
val responseBody = readResponseBody(connection = connection)
@@ -395,6 +391,19 @@ class AiChatRemoteService(
395391
)
396392
}
397393

394+
private fun encodeStopRunRequest(request: AiChatStopRunRequest): JSONObject {
395+
val payload: JSONObject = JSONObject()
396+
.put("sessionId", request.sessionId)
397+
398+
return putOptionalRunId(
399+
payload = putOptionalWorkspaceId(
400+
payload = payload,
401+
workspaceId = request.workspaceId
402+
),
403+
runId = request.runId
404+
)
405+
}
406+
398407
private fun putOptionalWorkspaceId(
399408
payload: JSONObject,
400409
workspaceId: String?
@@ -405,6 +414,16 @@ class AiChatRemoteService(
405414
return payload
406415
}
407416

417+
private fun putOptionalRunId(
418+
payload: JSONObject,
419+
runId: String?
420+
): JSONObject {
421+
runId?.takeIf { value -> value.isNotBlank() }?.let { resolvedRunId ->
422+
payload.put("runId", resolvedRunId)
423+
}
424+
return payload
425+
}
426+
408427
private fun putOptionalUiLocale(
409428
payload: JSONObject,
410429
uiLocale: String?

‎apps/android/data/local/src/main/java/com/flashcardsopensourceapp/data/local/model/AiChatModels.kt‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -306,6 +306,13 @@ data class AiChatNewSessionRequest(
306306
val uiLocale: String?,
307307
)
308308

309+
data class AiChatStopRunRequest(
310+
val sessionId: String,
311+
val workspaceId: String?,
312+
// TODO: Remove optional runId and make it required after most users have updated to the latest version. This is a legacy path.
313+
val runId: String?,
314+
)
315+
309316
data class AiToolCallRequest(
310317
val toolCallId: String,
311318
val name: String,

‎apps/android/data/local/src/main/java/com/flashcardsopensourceapp/data/local/repository/AiRepositories.kt‎

Lines changed: 7 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@ import com.flashcardsopensourceapp.data.local.model.AiChatNewSessionRequest
1616
import com.flashcardsopensourceapp.data.local.model.AiChatResumeDiagnostics
1717
import com.flashcardsopensourceapp.data.local.model.AiChatSessionProvisioningResult
1818
import com.flashcardsopensourceapp.data.local.model.AiChatSessionSnapshot
19+
import com.flashcardsopensourceapp.data.local.model.AiChatStopRunRequest
1920
import com.flashcardsopensourceapp.data.local.model.AiChatStopRunResponse
2021
import com.flashcardsopensourceapp.data.local.model.AiChatStartRunRequest
2122
import com.flashcardsopensourceapp.data.local.model.AiChatStartRunResponse
@@ -338,14 +339,17 @@ class LocalAiChatRepository(
338339
}
339340
}
340341

341-
override suspend fun stopRun(workspaceId: String?, sessionId: String): AiChatStopRunResponse {
342+
override suspend fun stopRun(workspaceId: String?, sessionId: String, runId: String?): AiChatStopRunResponse {
342343
val remoteWorkspaceId = requireRemoteWorkspaceId(workspaceId = workspaceId)
343344
val session = authorizedSession(workspaceId = remoteWorkspaceId)
344345
return aiChatRemoteService.stopRun(
345346
apiBaseUrl = session.apiBaseUrl,
346347
authorizationHeader = session.authorizationHeader,
347-
sessionId = sessionId,
348-
workspaceId = remoteWorkspaceId
348+
request = AiChatStopRunRequest(
349+
sessionId = sessionId,
350+
workspaceId = remoteWorkspaceId,
351+
runId = runId
352+
)
349353
)
350354
}
351355

‎apps/android/data/local/src/main/java/com/flashcardsopensourceapp/data/local/repository/Repositories.kt‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -206,5 +206,5 @@ interface AiChatRepository {
206206
afterCursor: String?,
207207
resumeDiagnostics: AiChatResumeDiagnostics?
208208
): Flow<AiChatLiveEvent>
209-
suspend fun stopRun(workspaceId: String?, sessionId: String): AiChatStopRunResponse
209+
suspend fun stopRun(workspaceId: String?, sessionId: String, runId: String?): AiChatStopRunResponse
210210
}

‎apps/android/data/local/src/test/java/com/flashcardsopensourceapp/data/local/ai/AiChatRemoteWireTest.kt‎

Lines changed: 47 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@ import com.flashcardsopensourceapp.data.local.model.AiChatNewSessionRequest
88
import com.flashcardsopensourceapp.data.local.model.AiChatResumeDiagnostics
99
import com.flashcardsopensourceapp.data.local.model.AiChatRunTerminalOutcome
1010
import com.flashcardsopensourceapp.data.local.model.AiChatStartRunRequest
11+
import com.flashcardsopensourceapp.data.local.model.AiChatStopRunRequest
1112
import com.flashcardsopensourceapp.data.local.model.AiChatWireContentPart
1213
import com.flashcardsopensourceapp.data.local.model.CloudServiceConfigurationMode
1314
import com.sun.net.httpserver.HttpServer
@@ -437,7 +438,7 @@ class AiChatRemoteWireTest {
437438
}
438439

439440
@Test
440-
fun stopRunIncludesWorkspaceIdWhenPresent() = runBlocking {
441+
fun stopRunIncludesWorkspaceIdAndRunIdWhenPresent() = runBlocking {
441442
val requestBodyRef = AtomicReference("")
442443
val server = HttpServer.create(InetSocketAddress("127.0.0.1", 0), 0)
443444
server.createContext("/chat/stop") { exchange ->
@@ -459,13 +460,56 @@ class AiChatRemoteWireTest {
459460
service.stopRun(
460461
apiBaseUrl = "http://127.0.0.1:${server.address.port}",
461462
authorizationHeader = "Bearer token-1",
462-
sessionId = "session-1",
463-
workspaceId = testWorkspaceId
463+
request = AiChatStopRunRequest(
464+
sessionId = "session-1",
465+
workspaceId = testWorkspaceId,
466+
runId = "run-1"
467+
)
468+
)
469+
470+
val requestBody = JSONObject(requestBodyRef.get())
471+
assertEquals("session-1", requestBody.getString("sessionId"))
472+
assertEquals(testWorkspaceId, requestBody.getString("workspaceId"))
473+
assertEquals("run-1", requestBody.getString("runId"))
474+
} finally {
475+
server.stop(0)
476+
}
477+
}
478+
479+
@Test
480+
fun stopRunOmitsBlankRunId() = runBlocking {
481+
val requestBodyRef = AtomicReference("")
482+
val server = HttpServer.create(InetSocketAddress("127.0.0.1", 0), 0)
483+
server.createContext("/chat/stop") { exchange ->
484+
requestBodyRef.set(exchange.requestBody.bufferedReader().use { reader -> reader.readText() })
485+
val body = """
486+
{
487+
"sessionId": "session-1",
488+
"stopped": true,
489+
"stillRunning": false
490+
}
491+
""".trimIndent().toByteArray()
492+
exchange.sendResponseHeaders(200, body.size.toLong())
493+
exchange.responseBody.use { outputStream -> outputStream.write(body) }
494+
}
495+
server.start()
496+
497+
try {
498+
val service = makeRemoteService()
499+
service.stopRun(
500+
apiBaseUrl = "http://127.0.0.1:${server.address.port}",
501+
authorizationHeader = "Bearer token-1",
502+
request = AiChatStopRunRequest(
503+
sessionId = "session-1",
504+
workspaceId = testWorkspaceId,
505+
runId = ""
506+
)
464507
)
465508

466509
val requestBody = JSONObject(requestBodyRef.get())
467510
assertEquals("session-1", requestBody.getString("sessionId"))
468511
assertEquals(testWorkspaceId, requestBody.getString("workspaceId"))
512+
assertFalse(requestBody.has("runId"))
469513
} finally {
470514
server.stop(0)
471515
}

‎apps/android/feature/ai/src/main/java/com/flashcardsopensourceapp/feature/ai/runtime/AiChatLiveStreamCoordinator.kt‎

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -98,6 +98,24 @@ internal class AiChatLiveStreamCoordinator(
9898
context.persistCurrentState()
9999
}
100100

101+
fun reconcileConversationAfterStopNoop() {
102+
context.activeLiveJob?.cancel(
103+
cause = CancellationException("AI live attach cancelled because the stop response did not stop the active run.")
104+
)
105+
context.activeLiveJob = null
106+
context.runtimeStateMutable.update { state ->
107+
state.copy(
108+
activeRun = null,
109+
isLiveAttached = false,
110+
composerPhase = AiComposerPhase.IDLE,
111+
repairStatus = null,
112+
errorMessage = ""
113+
)
114+
}
115+
context.persistCurrentState()
116+
restartConversationBootstrap(true, null)
117+
}
118+
101119
private fun attachLiveStream(
102120
workspaceId: String?,
103121
sessionId: String,

‎apps/android/feature/ai/src/main/java/com/flashcardsopensourceapp/feature/ai/runtime/AiChatSendCoordinator.kt‎

Lines changed: 11 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -174,12 +174,14 @@ internal class AiChatSendCoordinator(
174174
}
175175

176176
fun stopStreaming() {
177-
if (context.runtimeStateMutable.value.composerPhase != AiComposerPhase.RUNNING) {
177+
val currentState: AiChatRuntimeState = context.runtimeStateMutable.value
178+
if (currentState.composerPhase != AiComposerPhase.RUNNING) {
178179
return
179180
}
180181

181-
val sessionId = context.runtimeStateMutable.value.persistedState.chatSessionId
182-
val workspaceId = context.runtimeStateMutable.value.workspaceId
182+
val sessionId: String = currentState.persistedState.chatSessionId
183+
val workspaceId: String? = currentState.workspaceId
184+
val runId: String? = currentState.activeRun?.runId?.ifBlank { null }
183185

184186
context.runtimeStateMutable.update { state ->
185187
state.copy(
@@ -196,8 +198,13 @@ internal class AiChatSendCoordinator(
196198
if (sessionId.isNotBlank()) {
197199
val response = context.aiChatRepository.stopRun(
198200
workspaceId = workspaceId,
199-
sessionId = sessionId
201+
sessionId = sessionId,
202+
runId = runId
200203
)
204+
if (response.stopped.not()) {
205+
liveStreamCoordinator.reconcileConversationAfterStopNoop()
206+
return@launch
207+
}
201208
if (response.stopped && response.stillRunning.not()) {
202209
liveStreamCoordinator.finalizeStoppedConversation()
203210
return@launch

‎apps/android/feature/ai/src/test/java/com/flashcardsopensourceapp/feature/ai/runtime/AiChatRuntimeBootstrapAndLiveStreamTest.kt‎

Lines changed: 72 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3,9 +3,10 @@ package com.flashcardsopensourceapp.feature.ai.runtime
33
import com.flashcardsopensourceapp.data.local.model.AiChatAttachment
44
import com.flashcardsopensourceapp.data.local.model.AiChatContentPart
55
import com.flashcardsopensourceapp.data.local.model.AiChatLiveEvent
6+
import com.flashcardsopensourceapp.data.local.model.AiChatRunTerminalOutcome
7+
import com.flashcardsopensourceapp.data.local.model.AiChatStopRunResponse
68
import com.flashcardsopensourceapp.data.local.model.AiChatToolCall
79
import com.flashcardsopensourceapp.data.local.model.AiChatToolCallStatus
8-
import com.flashcardsopensourceapp.data.local.model.AiChatRunTerminalOutcome
910
import com.flashcardsopensourceapp.data.local.model.CloudAccountState
1011
import com.flashcardsopensourceapp.data.local.model.makeDefaultAiChatPersistedState
1112
import com.flashcardsopensourceapp.feature.ai.AiEntryPrefill
@@ -286,6 +287,76 @@ class AiChatRuntimeBootstrapAndLiveStreamTest {
286287
advanceUntilIdle()
287288
}
288289

290+
@Test
291+
fun stopStreamingSendsActiveRunIdWhenKnown() = runTest {
292+
val repository = FakeAiChatRepository()
293+
val liveEvents = MutableSharedFlow<AiChatLiveEvent>()
294+
repository.bootstrapResponses += makeBootstrapResponse(
295+
sessionId = "session-1",
296+
activeRun = makeActiveRun(runId = "run-1", cursor = "5")
297+
)
298+
repository.liveFlows["run-1"] = liveEvents
299+
val runtime = makeRuntime(scope = this, repository = repository)
300+
301+
runtime.onScreenVisible()
302+
runtime.updateAccessContext(makeAccessContext(workspaceId = defaultTestWorkspaceId))
303+
advanceUntilIdle()
304+
305+
assertEquals(AiComposerPhase.RUNNING, runtime.state.value.composerPhase)
306+
307+
runtime.stopStreaming()
308+
advanceUntilIdle()
309+
310+
assertEquals(listOf(defaultTestWorkspaceId), repository.stopRunWorkspaceIds)
311+
assertEquals(listOf("session-1"), repository.stopRunSessionIds)
312+
assertEquals(listOf("run-1"), repository.stopRunIds)
313+
314+
runtime.onScreenHidden()
315+
advanceUntilIdle()
316+
}
317+
318+
@Test
319+
fun stopStreamingReloadsBootstrapWhenStopRunReturnsNoop() = runTest {
320+
val repository = FakeAiChatRepository()
321+
val liveEvents = MutableSharedFlow<AiChatLiveEvent>()
322+
val replacementLiveEvents = MutableSharedFlow<AiChatLiveEvent>()
323+
repository.bootstrapResponses += makeBootstrapResponse(
324+
sessionId = "session-1",
325+
activeRun = makeActiveRun(runId = "run-1", cursor = "5")
326+
)
327+
repository.bootstrapResponses += makeBootstrapResponse(
328+
sessionId = "session-1",
329+
activeRun = makeActiveRun(runId = "run-2", cursor = "8")
330+
)
331+
repository.liveFlows["run-1"] = liveEvents
332+
repository.liveFlows["run-2"] = replacementLiveEvents
333+
repository.stopRunResponse = AiChatStopRunResponse(
334+
sessionId = "session-1",
335+
stopped = false,
336+
stillRunning = true
337+
)
338+
val runtime = makeRuntime(scope = this, repository = repository)
339+
340+
runtime.onScreenVisible()
341+
runtime.updateAccessContext(makeAccessContext(workspaceId = defaultTestWorkspaceId))
342+
advanceUntilIdle()
343+
344+
assertEquals(AiComposerPhase.RUNNING, runtime.state.value.composerPhase)
345+
assertEquals("run-1", runtime.state.value.activeRun?.runId)
346+
347+
runtime.stopStreaming()
348+
advanceUntilIdle()
349+
350+
assertEquals(listOf("ensured-session-1", "session-1"), repository.loadBootstrapSessionIds)
351+
assertEquals(listOf("run-1", "run-2"), repository.attachRunIds)
352+
assertEquals(AiComposerPhase.RUNNING, runtime.state.value.composerPhase)
353+
assertEquals("run-2", runtime.state.value.activeRun?.runId)
354+
assertTrue(runtime.state.value.isLiveAttached)
355+
356+
runtime.onScreenHidden()
357+
advanceUntilIdle()
358+
}
359+
289360
@Test
290361
fun bootstrapWithActiveRunAndTrailingToolCallTriggersAutoSyncOnTerminalCompletion() = runTest {
291362
val repository = FakeAiChatRepository()

‎apps/android/feature/ai/src/test/java/com/flashcardsopensourceapp/feature/ai/runtime/AiChatRuntimeTestSupport.kt‎

Lines changed: 12 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -225,6 +225,9 @@ internal class FakeAiChatRepository : AiChatRepository {
225225
val ensureSessionRequests: MutableList<String> = mutableListOf()
226226
val transcribeAudioWorkspaceIds: MutableList<String?> = mutableListOf()
227227
val transcribeAudioSessionIds: MutableList<String> = mutableListOf()
228+
val stopRunWorkspaceIds: MutableList<String?> = mutableListOf()
229+
val stopRunSessionIds: MutableList<String> = mutableListOf()
230+
val stopRunIds: MutableList<String?> = mutableListOf()
228231
var nextEnsureSessionId: String = "ensured-session-1"
229232
var transcribeAudioResponse: AiChatTranscriptionResult = AiChatTranscriptionResult(
230233
text = "transcribed speech",
@@ -237,6 +240,7 @@ internal class FakeAiChatRepository : AiChatRepository {
237240
var startRunCalls: Int = 0
238241
var lastStartRunState: AiChatPersistedState? = null
239242
var lastStartRunUiLocale: String? = null
243+
var stopRunResponse: AiChatStopRunResponse? = null
240244
var startRunResponse: AiChatStartRunResponse = AiChatAcceptedConversationEnvelope(
241245
accepted = true,
242246
sessionId = "session-1",
@@ -443,7 +447,14 @@ internal class FakeAiChatRepository : AiChatRepository {
443447
return liveFlows[runId] ?: emptyFlow()
444448
}
445449

446-
override suspend fun stopRun(workspaceId: String?, sessionId: String): AiChatStopRunResponse {
450+
override suspend fun stopRun(workspaceId: String?, sessionId: String, runId: String?): AiChatStopRunResponse {
451+
stopRunWorkspaceIds += workspaceId
452+
stopRunSessionIds += sessionId
453+
stopRunIds += runId
454+
val configuredResponse: AiChatStopRunResponse? = stopRunResponse
455+
if (configuredResponse != null) {
456+
return configuredResponse
457+
}
447458
return AiChatStopRunResponse(
448459
sessionId = sessionId,
449460
stopped = true,

0 commit comments

Comments
 (0)