2 files changed
+578
-38
lib/features/on_device/data/on_device_chat_service.dart
+280
-31
@@ -18,6 +18,9 @@ class OnDeviceChatService implements ChatService {
18
19
final OnDeviceInferenceService _gemmaService;
20
final Set<_InferenceRun> _activeRuns = <_InferenceRun>{};
21
+ _RetainedConversation? _retainedConversation;
22
+ int _nextRunSequence = 0;
23
+ int _latestRunSequence = 0;
24
bool _isDisposed = false;
25
26
@override
@@ -32,7 +35,9 @@ class OnDeviceChatService implements ChatService {
35
}) {
36
late final _InferenceRun run;
37
final controller = StreamController<ChatResponse>();
35
- run = _InferenceRun(controller);
38
+ final sequence = ++_nextRunSequence;
39
+ _latestRunSequence = sequence;
40
+ run = _InferenceRun(controller, sequence);
41
controller.onCancel = () => _cancelRun(run, notifyListener: false);
42
43
if (_isDisposed) {
@@ -61,13 +66,32 @@ class OnDeviceChatService implements ChatService {
66
ChatParameters params,
67
) async {
68
try {
64
- final systemInstruction = params.systemPrompt?.isNotEmpty == true
65
- ? params.systemPrompt
66
- : null;
67
- final session = await _gemmaService.createChat(
68
- systemInstruction: systemInstruction,
69
- tools: const [],
70
- );
69
+ final input = _prepareInput(modelId, messages, params);
70
+ if (input == null) {
71
+ await _finishWithError(run, 'On-device chat requires a user message.');
72
+ return;
73
+ }
74
+ run.input = input;
75
+
76
+ final retained = _retainedConversation;
77
+
78
+ late final OnDeviceInferenceSession session;
79
+ late final bool reusedSession;
80
+ if (retained != null && retained.canContinueWith(input, modelId)) {
81
+ _retainedConversation = null;
82
+ session = retained.session;
83
+ reusedSession = true;
84
+ } else {
85
+ if (retained != null && !input.isAuxiliary) {
86
+ _retainedConversation = null;
87
+ await retained.session.close();
88
+ }
89
+ session = await _gemmaService.createChat(
90
+ systemInstruction: input.baseSystemInstruction,
91
+ tools: const [],
92
+ );
93
+ reusedSession = false;
94
+ }
95
run.session = session;
96
97
if (run.isCancelled) {
@@ -75,13 +99,9 @@ class OnDeviceChatService implements ChatService {
99
return;
100
}
101
78
- for (final message in _convertMessages(messages)) {
79
- if (run.isCancelled) {
80
- await _closeSession(run, stopGeneration: false);
81
- return;
82
- }
83
- await session.addQueryChunk(message);
84
- }
102
+ await session.addQueryChunk(
103
+ input.toGemmaMessage(includeHistory: !reusedSession),
104
+ );
105
106
if (run.isCancelled) {
107
await _closeSession(run, stopGeneration: false);
@@ -129,6 +149,7 @@ class OnDeviceChatService implements ChatService {
149
final content =
150
run.textDecoder?.decodeChunk(response.token) ?? response.token;
151
if (content.isNotEmpty) {
152
+ run.generatedContent.write(content);
153
run.controller.add(
154
ChatResponse(type: ChatResponseType.message, content: content),
155
);
@@ -173,6 +194,7 @@ class OnDeviceChatService implements ChatService {
194
final finalToken = run.textDecoder?.flush() ?? '';
195
final finalReasoning = run.reasoningDecoder?.flush() ?? '';
196
if (finalToken.isNotEmpty) {
197
+ run.generatedContent.write(finalToken);
198
run.controller.add(
199
ChatResponse(type: ChatResponseType.message, content: finalToken),
200
);
@@ -190,7 +212,7 @@ class OnDeviceChatService implements ChatService {
212
}
213
}
214
193
- await _closeSession(run, stopGeneration: false);
215
+ await _retainOrCloseSession(run);
216
await _closeController(run, notifyListener: !run.isCancelled);
217
_activeRuns.remove(run);
218
_completeRunFuture(run);
@@ -292,6 +314,34 @@ class OnDeviceChatService implements ChatService {
314
}
315
}
316
317
+ Future<void> _retainOrCloseSession(_InferenceRun run) async {
318
+ final session = run.session;
319
+ run.session = null;
320
+ final input = run.input;
321
+ if (session == null) return;
322
+
323
+ if (run.isCancelled ||
324
+ input == null ||
325
+ input.isAuxiliary ||
326
+ run.sequence != _latestRunSequence ||
327
+ run.generatedContent.isEmpty) {
328
+ await session.close();
329
+ return;
330
+ }
331
+
332
+ final previous = _retainedConversation;
333
+ _retainedConversation = _RetainedConversation(
334
+ session: session,
335
+ conversationId: input.conversationId,
336
+ modelId: input.modelId,
337
+ baseSystemInstruction: input.baseSystemInstruction,
338
+ timeline: input.completedTimeline(run.generatedContent.toString()),
339
+ );
340
+ if (previous != null && !identical(previous.session, session)) {
341
+ await previous.session.close();
342
+ }
343
+ }
344
+
345
Future<void> _closeController(
346
_InferenceRun run, {
347
required bool notifyListener,
@@ -303,27 +353,79 @@ class OnDeviceChatService implements ChatService {
353
await run.controller.close();
354
}
355
306
- List<gemma.Message> _convertMessages(List<Message> messages) {
307
- var list = messages;
308
- if (list.isNotEmpty &&
309
- list.last.role == MessageRole.assistant &&
310
- list.last.content.isEmpty) {
311
- list = list.sublist(0, list.length - 1);
312
- }
313
- return list
356
+ _PreparedInput? _prepareInput(
357
+ String modelId,
358
+ List<Message> messages,
359
+ ChatParameters params,
360
+ ) {
361
+ final relevant = messages
362
.where(
363
(message) =>
364
message.role == MessageRole.user ||
365
message.role == MessageRole.assistant ||
366
message.role == MessageRole.system,
367
)
320
- .map(
321
- (message) => gemma.Message.text(
322
- text: message.content,
323
- isUser: message.role == MessageRole.user,
324
- ),
368
+ .toList(growable: true);
369
+
370
+ Message? pendingAssistant;
371
+ if (relevant.isNotEmpty &&
372
+ relevant.last.role == MessageRole.assistant &&
373
+ relevant.last.content.isEmpty) {
374
+ pendingAssistant = relevant.removeLast();
375
+ }
376
+
377
+ final systemMessages = relevant
378
+ .where(
379
+ (message) =>
380
+ message.role == MessageRole.system &&
381
+ message.content.trim().isNotEmpty,
382
+ )
383
+ .map((message) => message.content.trim())
384
+ .toList(growable: false);
385
+ final timeline = relevant
386
+ .where(
387
+ (message) =>
388
+ message.role == MessageRole.user ||
389
+ message.role == MessageRole.assistant,
390
)
326
- .toList();
391
+ .map(_MessageSnapshot.fromMessage)
392
+ .toList(growable: false);
393
+ if (timeline.isEmpty) return null;
394
+
395
+ final currentMessage = timeline.last;
396
+ 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);
402
+ final messageSystemInstruction = systemMessages.isEmpty
403
+ ? null
404
+ : systemMessages.join('\n\n');
405
+ final baseSystemInstruction = isAuxiliary
406
+ ? _nonEmpty(params.systemPrompt) ?? messageSystemInstruction
407
+ : messageSystemInstruction ?? _nonEmpty(params.systemPrompt);
408
+ if (currentMessage.role != MessageRole.user) {
409
+ Log.warning(
410
+ 'Rebuilding an on-device session for assistant continuation; '
411
+ 'history will be supplied as a transcript.',
412
+ );
413
+ }
414
+
415
+ return _PreparedInput(
416
+ conversationId: currentMessage.conversationId,
417
+ modelId: modelId,
418
+ baseSystemInstruction: baseSystemInstruction,
419
+ history: history,
420
+ currentMessage: currentMessage,
421
+ pendingAssistantId: pendingAssistant?.id,
422
+ isAuxiliary: isAuxiliary,
423
+ );
424
+ }
425
+
426
+ static String? _nonEmpty(String? value) {
427
+ final trimmed = value?.trim();
428
+ return trimmed == null || trimmed.isEmpty ? null : trimmed;
429
}
430
431
@override
@@ -337,18 +439,26 @@ class OnDeviceChatService implements ChatService {
439
if (_isDisposed) return;
440
_isDisposed = true;
441
cancelStream();
442
+ final retained = _retainedConversation;
443
+ _retainedConversation = null;
444
+ if (retained != null) {
445
+ unawaited(retained.session.close());
446
+ }
447
}
448
}
449
450
class _InferenceRun {
344
- _InferenceRun(this.controller);
451
+ _InferenceRun(this.controller, this.sequence);
452
453
final StreamController<ChatResponse> controller;
454
+ final int sequence;
455
StreamSubscription<gemma.ModelResponse>? subscription;
456
OnDeviceInferenceSession? session;
457
Completer<void>? streamCompleter;
458
BpeDecoder? textDecoder;
459
BpeDecoder? reasoningDecoder;
460
+ _PreparedInput? input;
461
+ final StringBuffer generatedContent = StringBuffer();
462
bool generationStarted = false;
463
bool isCancelled = false;
464
bool _isFinishing = false;
@@ -359,3 +469,142 @@ class _InferenceRun {
469
return true;
470
}
471
}
472
+
473
+class _PreparedInput {
474
+ const _PreparedInput({
475
+ required this.conversationId,
476
+ required this.modelId,
477
+ required this.baseSystemInstruction,
478
+ required this.history,
479
+ required this.currentMessage,
480
+ required this.pendingAssistantId,
481
+ required this.isAuxiliary,
482
+ });
483
+
484
+ final String conversationId;
485
+ final String modelId;
486
+ final String? baseSystemInstruction;
487
+ final List<_MessageSnapshot> history;
488
+ final _MessageSnapshot currentMessage;
489
+ final String? pendingAssistantId;
490
+ final bool isAuxiliary;
491
+
492
+ gemma.Message toGemmaMessage({required bool includeHistory}) {
493
+ if (!includeHistory || history.isEmpty) {
494
+ return currentMessage.toGemmaMessage();
495
+ }
496
+ final transcript = history
497
+ .map(
498
+ (message) =>
499
+ '${message.role == MessageRole.user ? 'User' : 'Assistant'}: '
500
+ '${message.content}',
501
+ )
502
+ .join('\n\n');
503
+ final currentLabel = currentMessage.role == MessageRole.user
504
+ ? 'Current user message'
505
+ : 'Assistant response to continue';
506
+ return gemma.Message.text(
507
+ text:
508
+ 'Earlier conversation transcript (context only):\n\n$transcript\n\n'
509
+ '$currentLabel:\n${currentMessage.content}',
510
+ isUser: currentMessage.role == MessageRole.user,
511
+ );
512
+ }
513
+
514
+ List<_MessageSnapshot> completedTimeline(String response) {
515
+ if (currentMessage.role == MessageRole.assistant) {
516
+ return [
517
+ ...history,
518
+ currentMessage.copyWith(content: '${currentMessage.content}$response'),
519
+ ];
520
+ }
521
+ return [
522
+ ...history,
523
+ currentMessage,
524
+ _MessageSnapshot(
525
+ id: pendingAssistantId ?? 'generated-${currentMessage.id}',
526
+ conversationId: conversationId,
527
+ role: MessageRole.assistant,
528
+ content: response,
529
+ ),
530
+ ];
531
+ }
532
+}
533
+
534
+class _RetainedConversation {
535
+ const _RetainedConversation({
536
+ required this.session,
537
+ required this.conversationId,
538
+ required this.modelId,
539
+ required this.baseSystemInstruction,
540
+ required this.timeline,
541
+ });
542
+
543
+ final OnDeviceInferenceSession session;
544
+ final String conversationId;
545
+ final String modelId;
546
+ final String? baseSystemInstruction;
547
+ final List<_MessageSnapshot> timeline;
548
+
549
+ bool canContinueWith(_PreparedInput input, String requestedModelId) {
550
+ return input.currentMessage.role == MessageRole.user &&
551
+ conversationId == input.conversationId &&
552
+ modelId == requestedModelId &&
553
+ baseSystemInstruction == input.baseSystemInstruction &&
554
+ _listEquals(timeline, input.history);
555
+ }
556
+
557
+ static bool _listEquals(
558
+ List<_MessageSnapshot> first,
559
+ List<_MessageSnapshot> second,
560
+ ) {
561
+ if (first.length != second.length) return false;
562
+ for (var index = 0; index < first.length; index++) {
563
+ if (first[index] != second[index]) return false;
564
+ }
565
+ return true;
566
+ }
567
+}
568
+
569
+class _MessageSnapshot {
570
+ const _MessageSnapshot({
571
+ required this.id,
572
+ required this.conversationId,
573
+ required this.role,
574
+ required this.content,
575
+ });
576
+
577
+ factory _MessageSnapshot.fromMessage(Message message) => _MessageSnapshot(
578
+ id: message.id,
579
+ conversationId: message.conversationId,
580
+ role: message.role,
581
+ content: message.content,
582
+ );
583
+
584
+ final String id;
585
+ final String conversationId;
586
+ final MessageRole role;
587
+ final String content;
588
+
589
+ gemma.Message toGemmaMessage() =>
590
+ gemma.Message.text(text: content, isUser: role == MessageRole.user);
591
+
592
+ _MessageSnapshot copyWith({String? content}) => _MessageSnapshot(
593
+ id: id,
594
+ conversationId: conversationId,
595
+ role: role,
596
+ content: content ?? this.content,
597
+ );
598
+
599
+ @override
600
+ bool operator ==(Object other) {
601
+ return other is _MessageSnapshot &&
602
+ id == other.id &&
603
+ conversationId == other.conversationId &&
604
+ role == other.role &&
605
+ content == other.content;
606
+ }
607
+
608
+ @override
609
+ int get hashCode => Object.hash(id, conversationId, role, content);
610
+}
test/features/on_device/on_device_chat_service_lifecycle_test.dart
+298
-7
@@ -53,7 +53,11 @@ void main() {
53
.listen(secondResponses.add)
54
.asFuture<void>();
55
56
- await _waitFor(() => inferenceService.sessions.length == 2);
56
+ await _waitFor(
57
+ () =>
58
+ inferenceService.sessions.length == 2 &&
59
+ inferenceService.sessions.every((session) => session.didGenerate),
60
+ );
61
final first = inferenceService.sessions[0];
62
final second = inferenceService.sessions[1];
63
@@ -65,7 +69,7 @@ void main() {
69
70
expect(inferenceService.createCount, 2);
71
expect(first.closeCount, 1);
68
- expect(second.closeCount, 1);
72
+ expect(second.closeCount, 0);
73
expect(first.stopCount, 0);
74
expect(second.stopCount, 0);
75
expect(
@@ -79,6 +83,279 @@ void main() {
83
},
84
);
85
86
+ test('reuses a chat session and only sends the new user turn', () async {
87
+ final firstDone = chatService
88
+ .sendMessage(
89
+ server: _server(),
90
+ modelId: 'qwen3-0.6b',
91
+ messages: [
92
+ _message('first', id: 'user-1'),
93
+ _message('', id: 'assistant-1', role: MessageRole.assistant),
94
+ ],
95
+ params: ChatParameters.defaults(),
96
+ )
97
+ .listen((_) {})
98
+ .asFuture<void>();
99
+
100
+ await _waitFor(
101
+ () =>
102
+ inferenceService.sessions.length == 1 &&
103
+ inferenceService.sessions.single.generationCount == 1,
104
+ );
105
+ final session = inferenceService.sessions.single;
106
+ session.responses.add(const gemma.TextResponse('first answer'));
107
+ await session.responses.close();
108
+ await firstDone;
109
+
110
+ final secondDone = chatService
111
+ .sendMessage(
112
+ server: _server(),
113
+ modelId: 'qwen3-0.6b',
114
+ messages: [
115
+ _message('first', id: 'user-1'),
116
+ _message(
117
+ 'first answer',
118
+ id: 'assistant-1',
119
+ role: MessageRole.assistant,
120
+ ),
121
+ _message('second', id: 'user-2'),
122
+ _message('', id: 'assistant-2', role: MessageRole.assistant),
123
+ ],
124
+ params: ChatParameters.defaults(),
125
+ )
126
+ .listen((_) {})
127
+ .asFuture<void>();
128
+
129
+ await _waitFor(() => session.generationCount == 2);
130
+ session.responses.add(const gemma.TextResponse('second answer'));
131
+ await session.responses.close();
132
+ await secondDone;
133
+
134
+ expect(inferenceService.createCount, 1);
135
+ expect(session.messages.map((message) => message.text).toList(), [
136
+ 'first',
137
+ 'second',
138
+ ]);
139
+ expect(session.messages.every((message) => message.isUser), isTrue);
140
+ expect(session.closeCount, 0);
141
+ });
142
+
143
+ test(
144
+ 'passes system content once and never stages it as a chat turn',
145
+ () async {
146
+ final done = chatService
147
+ .sendMessage(
148
+ server: _server(),
149
+ modelId: 'qwen3-0.6b',
150
+ messages: [
151
+ _message(
152
+ 'You are concise.',
153
+ id: 'system',
154
+ role: MessageRole.system,
155
+ ),
156
+ _message('hello', id: 'user'),
157
+ _message('', id: 'assistant', role: MessageRole.assistant),
158
+ ],
159
+ params: ChatParameters.defaults().copyWith(
160
+ systemPrompt: 'You are concise.',
161
+ ),
162
+ )
163
+ .listen((_) {})
164
+ .asFuture<void>();
165
+
166
+ await _waitFor(
167
+ () =>
168
+ inferenceService.sessions.isNotEmpty &&
169
+ inferenceService.sessions.single.didGenerate,
170
+ );
171
+ final session = inferenceService.sessions.single;
172
+ session.responses.add(const gemma.TextResponse('hi'));
173
+ await session.responses.close();
174
+ await done;
175
+
176
+ expect(inferenceService.systemInstructions, ['You are concise.']);
177
+ expect(session.messages.map((message) => message.text).toList(), [
178
+ 'hello',
179
+ ]);
180
+ },
181
+ );
182
+
183
+ test('restores existing history as a role-labelled transcript', () async {
184
+ final done = chatService
185
+ .sendMessage(
186
+ server: _server(),
187
+ modelId: 'qwen3-0.6b',
188
+ messages: [
189
+ _message('first', id: 'user-1'),
190
+ _message(
191
+ 'first answer',
192
+ id: 'assistant-1',
193
+ role: MessageRole.assistant,
194
+ ),
195
+ _message('second', id: 'user-2'),
196
+ _message('', id: 'assistant-2', role: MessageRole.assistant),
197
+ ],
198
+ params: ChatParameters.defaults(),
199
+ )
200
+ .listen((_) {})
201
+ .asFuture<void>();
202
+
203
+ await _waitFor(
204
+ () =>
205
+ inferenceService.sessions.isNotEmpty &&
206
+ inferenceService.sessions.single.didGenerate,
207
+ );
208
+ final session = inferenceService.sessions.single;
209
+ session.responses.add(const gemma.TextResponse('second answer'));
210
+ await session.responses.close();
211
+ await done;
212
+
213
+ expect(session.messages, hasLength(1));
214
+ expect(
215
+ session.messages.single.text,
216
+ contains('User: first\n\nAssistant: first answer'),
217
+ );
218
+ expect(
219
+ session.messages.single.text,
220
+ contains('Current user message:\nsecond'),
221
+ );
222
+ expect(inferenceService.systemInstructions.single, isNull);
223
+ });
224
+
225
+ test('rebuilds the session when the upstream timeline diverges', () async {
226
+ final firstDone = chatService
227
+ .sendMessage(
228
+ server: _server(),
229
+ modelId: 'qwen3-0.6b',
230
+ messages: [
231
+ _message('first', id: 'user-1'),
232
+ _message('', id: 'assistant-1', role: MessageRole.assistant),
233
+ ],
234
+ params: ChatParameters.defaults(),
235
+ )
236
+ .listen((_) {})
237
+ .asFuture<void>();
238
+ await _waitFor(
239
+ () =>
240
+ inferenceService.sessions.length == 1 &&
241
+ inferenceService.sessions.single.didGenerate,
242
+ );
243
+ final firstSession = inferenceService.sessions.single;
244
+ firstSession.responses.add(const gemma.TextResponse('first answer'));
245
+ await firstSession.responses.close();
246
+ await firstDone;
247
+
248
+ final branchDone = chatService
249
+ .sendMessage(
250
+ server: _server(),
251
+ modelId: 'qwen3-0.6b',
252
+ messages: [
253
+ _message('edited first', id: 'user-1'),
254
+ _message('second', id: 'user-2'),
255
+ _message('', id: 'assistant-2', role: MessageRole.assistant),
256
+ ],
257
+ params: ChatParameters.defaults(),
258
+ )
259
+ .listen((_) {})
260
+ .asFuture<void>();
261
+ await _waitFor(
262
+ () =>
263
+ inferenceService.sessions.length == 2 &&
264
+ inferenceService.sessions.last.didGenerate,
265
+ );
266
+ final branchSession = inferenceService.sessions.last;
267
+ branchSession.responses.add(const gemma.TextResponse('branch answer'));
268
+ await branchSession.responses.close();
269
+ await branchDone;
270
+
271
+ expect(firstSession.closeCount, 1);
272
+ expect(branchSession.messages, hasLength(1));
273
+ expect(branchSession.messages.single.text, contains('User: edited first'));
274
+ });
275
+
276
+ test('auxiliary generation does not replace the retained chat', () async {
277
+ final firstDone = chatService
278
+ .sendMessage(
279
+ server: _server(),
280
+ modelId: 'qwen3-0.6b',
281
+ messages: [
282
+ _message('first', id: 'user-1'),
283
+ _message('', id: 'assistant-1', role: MessageRole.assistant),
284
+ ],
285
+ params: ChatParameters.defaults(),
286
+ )
287
+ .listen((_) {})
288
+ .asFuture<void>();
289
+ await _waitFor(
290
+ () =>
291
+ inferenceService.sessions.length == 1 &&
292
+ inferenceService.sessions.single.didGenerate,
293
+ );
294
+ final chatSession = inferenceService.sessions.single;
295
+ chatSession.responses.add(const gemma.TextResponse('first answer'));
296
+ await chatSession.responses.close();
297
+ await firstDone;
298
+
299
+ final titleDone = chatService
300
+ .sendMessage(
301
+ server: _server(),
302
+ modelId: 'qwen3-0.6b',
303
+ messages: [
304
+ _message('first', id: 'user-1'),
305
+ _message(
306
+ 'first answer',
307
+ id: 'assistant-1',
308
+ role: MessageRole.assistant,
309
+ ),
310
+ _message('Create a title', id: 'title-generation-prompt'),
311
+ ],
312
+ params: ChatParameters.defaults().copyWith(
313
+ systemPrompt: 'Generate only a title.',
314
+ ),
315
+ )
316
+ .listen((_) {})
317
+ .asFuture<void>();
318
+ await _waitFor(
319
+ () =>
320
+ inferenceService.sessions.length == 2 &&
321
+ inferenceService.sessions.last.didGenerate,
322
+ );
323
+ final titleSession = inferenceService.sessions.last;
324
+ titleSession.responses.add(const gemma.TextResponse('A title'));
325
+ await titleSession.responses.close();
326
+ await titleDone;
327
+
328
+ final secondDone = chatService
329
+ .sendMessage(
330
+ server: _server(),
331
+ modelId: 'qwen3-0.6b',
332
+ messages: [
333
+ _message('first', id: 'user-1'),
334
+ _message(
335
+ 'first answer',
336
+ id: 'assistant-1',
337
+ role: MessageRole.assistant,
338
+ ),
339
+ _message('second', id: 'user-2'),
340
+ _message('', id: 'assistant-2', role: MessageRole.assistant),
341
+ ],
342
+ params: ChatParameters.defaults(),
343
+ )
344
+ .listen((_) {})
345
+ .asFuture<void>();
346
+ await _waitFor(() => chatSession.generationCount == 2);
347
+ chatSession.responses.add(const gemma.TextResponse('second answer'));
348
+ await chatSession.responses.close();
349
+ await secondDone;
350
+
351
+ expect(inferenceService.createCount, 2);
352
+ expect(titleSession.closeCount, 1);
353
+ expect(chatSession.messages.map((message) => message.text).toList(), [
354
+ 'first',
355
+ 'second',
356
+ ]);
357
+ });
358
+
359
test('cancelling one subscription only closes its session', () async {
360
final firstSubscription = chatService
361
.sendMessage(
@@ -134,7 +411,11 @@ void main() {
411
.listen(responses.add)
412
.asFuture<void>();
413
137
- await _waitFor(() => inferenceService.sessions.isNotEmpty);
414
+ await _waitFor(
415
+ () =>
416
+ inferenceService.sessions.isNotEmpty &&
417
+ inferenceService.sessions.single.didGenerate,
418
+ );
419
inferenceService.sessions.single.responses.addError(
420
StateError('expected inference failure'),
421
StackTrace.current,
@@ -184,6 +465,7 @@ void main() {
465
466
class _FakeInferenceService implements OnDeviceInferenceService {
467
final List<_FakeInferenceSession> sessions = [];
468
+ final List<String?> systemInstructions = [];
469
int createCount = 0;
470
471
@override
@@ -195,6 +477,7 @@ class _FakeInferenceService implements OnDeviceInferenceService {
477
List<gemma.Tool> tools = const [],
478
}) async {
479
createCount++;
480
+ systemInstructions.add(systemInstruction);
481
final session = _FakeInferenceSession();
482
sessions.add(session);
483
return session;
@@ -202,11 +485,13 @@ class _FakeInferenceService implements OnDeviceInferenceService {
485
}
486
487
class _FakeInferenceSession implements OnDeviceInferenceSession {
205
- final StreamController<gemma.ModelResponse> responses = StreamController();
488
+ final List<StreamController<gemma.ModelResponse>> _responseStreams = [];
489
final List<gemma.Message> messages = [];
490
int stopCount = 0;
491
int closeCount = 0;
492
bool didGenerate = false;
493
+ int get generationCount => _responseStreams.length;
494
+ StreamController<gemma.ModelResponse> get responses => _responseStreams.last;
495
496
@override
497
Future<void> addQueryChunk(gemma.Message message) async {
@@ -216,6 +501,8 @@ class _FakeInferenceSession implements OnDeviceInferenceSession {
501
@override
502
Stream<gemma.ModelResponse> generateChatResponseAsync() {
503
didGenerate = true;
504
+ final responses = StreamController<gemma.ModelResponse>();
505
+ _responseStreams.add(responses);
506
return responses.stream;
507
}
508
@@ -238,10 +525,14 @@ Future<void> _waitFor(bool Function() condition) async {
525
fail('Condition was not met');
526
}
527
241
-Message _message(String content) => Message(
242
- id: content,
528
+Message _message(
529
+ String content, {
530
+ String? id,
531
+ MessageRole role = MessageRole.user,
532
+}) => Message(
533
+ id: id ?? content,
534
conversationId: 'conversation',
244
- role: MessageRole.user,
535
+ role: role,
536
content: content,
537
createdAt: DateTime.utc(2026, 8, 9),
538
);