fix(on-device): bound retained-session memory and tidy auxiliary-id coupling
- add a 10-minute idle TTL for retained chat sessions to bound background memory cost when a user abandons a conversation mid-stream; the timer is cancelled on reuse and on dispose - extract the auxiliary message-id set to a documented top-level constant so the coupling with title_generation_service and smart_reply_service is visible at the call site - add a BPE round-trip test verifying the retained assistant snapshot matches the streamed content verbatim (the BPE byte set cannot form multi-byte UTF-8 sequences, so the flush tail is always ASCII) - add tests for TTL elapse and reuse re-arming
abdulmominsakib committed
Aug 31, 2026 at 12:26 UTC
f4d0b64499ef7749098ba15308ce1f9a54584740
2 files changed
+288
-6
lib/features/on_device/data/on_device_chat_service.dart
+52
-6
@@ -13,12 +13,39 @@ import '../../chat/data/tools/tool_definition.dart';
13
import '../../servers/data/models/server.dart';
14
import 'on_device_gemma_service.dart';
15
16
+/// Synthetic message ids whose presence signals an auxiliary generation
17
+/// (title, smart-reply, AI-suggested user message) rather than a real chat
18
+/// turn. Auxiliary runs are isolated from any retained chat session so their
19
+/// system prompts and context do not bleed into the user's primary
20
+/// conversation. Keep this list in sync with the synthetic ids emitted by
21
+/// `title_generation_service.dart` and `smart_reply_service.dart`.
22
+const Set<String> _auxiliaryMessageIds = {
23
+ 'title-generation-prompt',
24
+ 'smart-reply-prompt',
25
+ 'ai-user-response-prompt',
26
+};
27
+
28
+/// How long an on-device chat session is retained after a successful turn
29
+/// without further activity. Retained sessions keep their KV cache in memory;
30
+/// this bounds the cost of background retention when a user abandons a
31
+/// conversation mid-stream.
32
+const Duration _defaultRetainedSessionTtl = Duration(minutes: 10);
33
+
34
class OnDeviceChatService implements ChatService {
17
- OnDeviceChatService(this._gemmaService);
35
+ OnDeviceChatService(
36
+ this._gemmaService, {
37
+ Duration retainedSessionTtl = _defaultRetainedSessionTtl,
38
+ // The retained-session TTL is a separate constructor parameter (rather
39
+ // than an initializing formal) so it can have a default value; the
40
+ // lint suggestion doesn't apply when a default is involved.
41
+ // ignore: prefer_initializing_formals
42
+ }) : _retainedSessionTtl = retainedSessionTtl;
43
44
final OnDeviceInferenceService _gemmaService;
45
+ final Duration _retainedSessionTtl;
46
final Set<_InferenceRun> _activeRuns = <_InferenceRun>{};
47
_RetainedConversation? _retainedConversation;
48
+ Timer? _retainedExpirationTimer;
49
int _nextRunSequence = 0;
50
int _latestRunSequence = 0;
51
bool _isDisposed = false;
@@ -79,11 +106,13 @@ class OnDeviceChatService implements ChatService {
106
late final bool reusedSession;
107
if (retained != null && retained.canContinueWith(input, modelId)) {
108
_retainedConversation = null;
109
+ _cancelRetainedExpiration();
110
session = retained.session;
111
reusedSession = true;
112
} else {
113
if (retained != null && !input.isAuxiliary) {
114
_retainedConversation = null;
115
+ _cancelRetainedExpiration();
116
await retained.session.close();
117
}
118
session = await _gemmaService.createChat(
@@ -330,6 +359,7 @@ class OnDeviceChatService implements ChatService {
359
}
360
361
final previous = _retainedConversation;
362
+ _cancelRetainedExpiration();
363
_retainedConversation = _RetainedConversation(
364
session: session,
365
conversationId: input.conversationId,
@@ -337,11 +367,30 @@ class OnDeviceChatService implements ChatService {
367
baseSystemInstruction: input.baseSystemInstruction,
368
timeline: input.completedTimeline(run.generatedContent.toString()),
369
);
370
+ _armRetainedExpiration();
371
if (previous != null && !identical(previous.session, session)) {
372
await previous.session.close();
373
}
374
}
375
376
+ void _armRetainedExpiration() {
377
+ _retainedExpirationTimer?.cancel();
378
+ _retainedExpirationTimer = Timer(_retainedSessionTtl, () {
379
+ if (_isDisposed) return;
380
+ final retained = _retainedConversation;
381
+ _retainedConversation = null;
382
+ _retainedExpirationTimer = null;
383
+ if (retained != null) {
384
+ unawaited(retained.session.close());
385
+ }
386
+ });
387
+ }
388
+
389
+ void _cancelRetainedExpiration() {
390
+ _retainedExpirationTimer?.cancel();
391
+ _retainedExpirationTimer = null;
392
+ }
393
+
394
Future<void> _closeController(
395
_InferenceRun run, {
396
required bool notifyListener,
@@ -394,11 +443,7 @@ class OnDeviceChatService implements ChatService {
443
444
final currentMessage = timeline.last;
445
final history = timeline.sublist(0, timeline.length - 1);
397
- final isAuxiliary = const {
398
- 'title-generation-prompt',
399
- 'smart-reply-prompt',
400
- 'ai-user-response-prompt',
401
- }.contains(currentMessage.id);
446
+ final isAuxiliary = _auxiliaryMessageIds.contains(currentMessage.id);
447
final messageSystemInstruction = systemMessages.isEmpty
448
? null
449
: systemMessages.join('\n\n');
@@ -439,6 +484,7 @@ class OnDeviceChatService implements ChatService {
484
if (_isDisposed) return;
485
_isDisposed = true;
486
cancelStream();
487
+ _cancelRetainedExpiration();
488
final retained = _retainedConversation;
489
_retainedConversation = null;
490
if (retained != null) {
test/features/on_device/on_device_chat_service_lifecycle_test.dart
+236
@@ -356,6 +356,242 @@ void main() {
356
]);
357
});
358
359
+ test(
360
+ 'BPE-decoded tokens including the flush tail are reused verbatim',
361
+ () async {
362
+ // Drive a BPE model with tokens that exercise the byte range covered by
363
+ // _bpeToByte (0–31, 127, 128–159). The BPE byte set cannot form any
364
+ // multi-byte UTF-8 sequence, so the flush tail is always ASCII and
365
+ // concatenates cleanly to the streamed chunks.
366
+ const bpeTokens = <String>['Ġhello', 'ġĢģ', 'Ġworld'];
367
+
368
+ final firstResponses = <ChatResponse>[];
369
+ final firstDone = chatService
370
+ .sendMessage(
371
+ server: _server(),
372
+ modelId: 'qwen3-0.6b',
373
+ messages: [
374
+ _message('hi', id: 'user-1'),
375
+ _message('', id: 'assistant-1', role: MessageRole.assistant),
376
+ ],
377
+ params: ChatParameters.defaults(),
378
+ )
379
+ .listen(firstResponses.add)
380
+ .asFuture<void>();
381
+ await _waitFor(
382
+ () =>
383
+ inferenceService.sessions.length == 1 &&
384
+ inferenceService.sessions.single.didGenerate,
385
+ );
386
+ final session = inferenceService.sessions.single;
387
+ for (final token in bpeTokens) {
388
+ session.responses.add(gemma.TextResponse(token));
389
+ }
390
+ await session.responses.close();
391
+ await firstDone;
392
+
393
+ final streamedContent = firstResponses
394
+ .where((response) => response.type == ChatResponseType.message)
395
+ .map((response) => response.content ?? '')
396
+ .join();
397
+
398
+ // The retained assistant snapshot must equal the streamed content
399
+ // exactly. We verify this indirectly: feeding the streamed content
400
+ // back as the prior assistant message must result in a reuse (no
401
+ // rebuild), because the retained timeline already matches.
402
+ final secondDone = chatService
403
+ .sendMessage(
404
+ server: _server(),
405
+ modelId: 'qwen3-0.6b',
406
+ messages: [
407
+ _message('hi', id: 'user-1'),
408
+ _message(
409
+ streamedContent,
410
+ id: 'assistant-1',
411
+ role: MessageRole.assistant,
412
+ ),
413
+ _message('again', id: 'user-2'),
414
+ _message('', id: 'assistant-2', role: MessageRole.assistant),
415
+ ],
416
+ params: ChatParameters.defaults(),
417
+ )
418
+ .listen((_) {})
419
+ .asFuture<void>();
420
+ await _waitFor(() => session.generationCount == 2);
421
+ session.responses.add(const gemma.TextResponse('reply'));
422
+ await session.responses.close();
423
+ await secondDone;
424
+
425
+ expect(inferenceService.createCount, 1);
426
+ expect(session.closeCount, 0);
427
+ expect(session.messages.map((message) => message.text).toList(), [
428
+ 'hi',
429
+ 'again',
430
+ ]);
431
+ },
432
+ );
433
+
434
+ test(
435
+ 'retained sessions are closed after the TTL elapses',
436
+ () async {
437
+ final ttlService = OnDeviceChatService(
438
+ inferenceService,
439
+ retainedSessionTtl: const Duration(milliseconds: 50),
440
+ );
441
+ try {
442
+ final firstDone = ttlService
443
+ .sendMessage(
444
+ server: _server(),
445
+ modelId: 'qwen3-0.6b',
446
+ messages: [
447
+ _message('hi', id: 'user-1'),
448
+ _message(
449
+ '',
450
+ id: 'assistant-1',
451
+ role: MessageRole.assistant,
452
+ ),
453
+ ],
454
+ params: ChatParameters.defaults(),
455
+ )
456
+ .listen((_) {})
457
+ .asFuture<void>();
458
+ await _waitFor(
459
+ () =>
460
+ inferenceService.sessions.length == 1 &&
461
+ inferenceService.sessions.single.didGenerate,
462
+ );
463
+ final session = inferenceService.sessions.single;
464
+ session.responses.add(const gemma.TextResponse('hello'));
465
+ await session.responses.close();
466
+ await firstDone;
467
+
468
+ // Session is retained immediately after completion.
469
+ expect(session.closeCount, 0);
470
+
471
+ // After the TTL elapses, the retained session is closed.
472
+ await _waitFor(() => session.closeCount == 1);
473
+
474
+ // A subsequent send with a different conversation id forces a
475
+ // rebuild — the retained session has been closed by the TTL.
476
+ final secondDone = ttlService
477
+ .sendMessage(
478
+ server: _server(),
479
+ modelId: 'qwen3-0.6b',
480
+ messages: [
481
+ Message(
482
+ id: 'user-1',
483
+ conversationId: 'other-conversation',
484
+ role: MessageRole.user,
485
+ content: 'hi',
486
+ createdAt: DateTime.utc(2026, 8, 9),
487
+ ),
488
+ Message(
489
+ id: 'assistant-1',
490
+ conversationId: 'other-conversation',
491
+ role: MessageRole.assistant,
492
+ content: '',
493
+ createdAt: DateTime.utc(2026, 8, 9),
494
+ ),
495
+ ],
496
+ params: ChatParameters.defaults(),
497
+ )
498
+ .listen((_) {})
499
+ .asFuture<void>();
500
+ await _waitFor(
501
+ () =>
502
+ inferenceService.sessions.length == 2 &&
503
+ inferenceService.sessions.last.didGenerate,
504
+ );
505
+ final secondSession = inferenceService.sessions.last;
506
+ secondSession.responses.add(const gemma.TextResponse('reply'));
507
+ await secondSession.responses.close();
508
+ await secondDone;
509
+
510
+ expect(inferenceService.createCount, 2);
511
+ } finally {
512
+ ttlService.dispose();
513
+ }
514
+ },
515
+ );
516
+
517
+ test(
518
+ 'reusing a retained session re-arms its TTL from the reuse point',
519
+ () async {
520
+ final ttlService = OnDeviceChatService(
521
+ inferenceService,
522
+ retainedSessionTtl: const Duration(milliseconds: 100),
523
+ );
524
+ try {
525
+ final firstDone = ttlService
526
+ .sendMessage(
527
+ server: _server(),
528
+ modelId: 'qwen3-0.6b',
529
+ messages: [
530
+ _message('hi', id: 'user-1'),
531
+ _message(
532
+ '',
533
+ id: 'assistant-1',
534
+ role: MessageRole.assistant,
535
+ ),
536
+ ],
537
+ params: ChatParameters.defaults(),
538
+ )
539
+ .listen((_) {})
540
+ .asFuture<void>();
541
+ await _waitFor(
542
+ () =>
543
+ inferenceService.sessions.length == 1 &&
544
+ inferenceService.sessions.single.didGenerate,
545
+ );
546
+ final session = inferenceService.sessions.single;
547
+ session.responses.add(const gemma.TextResponse('hello'));
548
+ await session.responses.close();
549
+ await firstDone;
550
+
551
+ // Reuse at t≈0 (TTL is 100ms). The reuse should re-arm the timer
552
+ // for a fresh 100ms from now.
553
+ await Future<void>.delayed(const Duration(milliseconds: 30));
554
+ final secondDone = ttlService
555
+ .sendMessage(
556
+ server: _server(),
557
+ modelId: 'qwen3-0.6b',
558
+ messages: [
559
+ _message('hi', id: 'user-1'),
560
+ _message(
561
+ 'hello',
562
+ id: 'assistant-1',
563
+ role: MessageRole.assistant,
564
+ ),
565
+ _message('again', id: 'user-2'),
566
+ _message(
567
+ '',
568
+ id: 'assistant-2',
569
+ role: MessageRole.assistant,
570
+ ),
571
+ ],
572
+ params: ChatParameters.defaults(),
573
+ )
574
+ .listen((_) {})
575
+ .asFuture<void>();
576
+ await _waitFor(() => session.generationCount == 2);
577
+ session.responses.add(const gemma.TextResponse('reply'));
578
+ await session.responses.close();
579
+ await secondDone;
580
+
581
+ // At t≈30ms+ε, the original TTL (100ms from first completion) has
582
+ // not yet elapsed, so the original timer would still be alive if it
583
+ // had not been cancelled. We assert no close yet.
584
+ expect(session.closeCount, 0);
585
+
586
+ // After 120ms total since reuse, the re-armed TTL has elapsed and
587
+ // the session should be closed by exactly one TTL.
588
+ await _waitFor(() => session.closeCount == 1);
589
+ } finally {
590
+ ttlService.dispose();
591
+ }
592
+ },
593
+ );
594
+
595
test('cancelling one subscription only closes its session', () async {
596
final firstSubscription = chatService
597
.sendMessage(