4 files changed
+82
-2
lib/features/on_device/data/on_device_chat_service.dart
+5
-2
@@ -136,7 +136,9 @@ class OnDeviceChatService implements ChatService {
136
137
late final OnDeviceInferenceSession session;
138
late final bool reusedSession;
139
- if (retained != null && retained.canContinueWith(input, modelId)) {
139
+ if (retained != null &&
140
+ _gemmaService.supportsChatSessionReuse &&
141
+ retained.canContinueWith(input, modelId)) {
142
_retainedConversation = null;
143
_cancelRetainedExpiration();
144
session = retained.session;
@@ -385,7 +387,8 @@ class OnDeviceChatService implements ChatService {
387
final input = run.input;
388
if (session == null) return;
389
388
- if (run.isCancelled ||
390
+ if (!_gemmaService.supportsChatSessionReuse ||
391
+ run.isCancelled ||
392
input == null ||
393
input.isAuxiliary ||
394
run.sequence != _latestRunSequence ||
lib/features/on_device/data/on_device_gemma_service.dart
+5
@@ -21,6 +21,7 @@ abstract interface class OnDeviceInferenceSession {
21
abstract interface class OnDeviceInferenceService {
22
bool get isLoaded;
23
bool get currentModelSupportsVision;
24
+ bool get supportsChatSessionReuse;
25
26
Future<OnDeviceInferenceSession> createChat({
27
String? systemInstruction,
@@ -69,6 +70,10 @@ class OnDeviceGemmaService implements OnDeviceInferenceService {
70
71
@override
72
bool get isLoaded => _model != null && !_isDisposed;
73
+
74
+ @override
75
+ bool get supportsChatSessionReuse =>
76
+ _model != null && _model!.fileType != ModelFileType.builtIn;
77
bool get isDisposed => _isDisposed;
78
79
/// The filename of the model restored by the native model manager on app
test/features/on_device/on_device_chat_attachments_test.dart
+3
@@ -522,6 +522,9 @@ class _FakeInferenceService implements OnDeviceInferenceService {
522
@override
523
bool get isLoaded => true;
524
525
+ @override
526
+ bool get supportsChatSessionReuse => true;
527
+
528
@override
529
bool currentModelSupportsVision = false;
530
test/features/on_device/on_device_chat_service_lifecycle_test.dart
+69
@@ -140,6 +140,71 @@ void main() {
140
expect(session.closeCount, 0);
141
});
142
143
+ test(
144
+ 'rebuilds runtimes that close their session after each response',
145
+ () async {
146
+ inferenceService.canReuseChatSession = false;
147
+
148
+ final firstDone = chatService
149
+ .sendMessage(
150
+ server: _server(),
151
+ modelId: 'apple-foundation-models',
152
+ messages: [
153
+ _message('first', id: 'user-1'),
154
+ _message('', id: 'assistant-1', role: MessageRole.assistant),
155
+ ],
156
+ params: ChatParameters.defaults(),
157
+ )
158
+ .listen((_) {})
159
+ .asFuture<void>();
160
+
161
+ await _waitFor(
162
+ () =>
163
+ inferenceService.sessions.length == 1 &&
164
+ inferenceService.sessions.single.didGenerate,
165
+ );
166
+ final firstSession = inferenceService.sessions.single;
167
+ firstSession.responses.add(const gemma.TextResponse('first answer'));
168
+ await firstSession.responses.close();
169
+ await firstDone;
170
+
171
+ expect(firstSession.closeCount, 1);
172
+
173
+ final secondDone = chatService
174
+ .sendMessage(
175
+ server: _server(),
176
+ modelId: 'apple-foundation-models',
177
+ messages: [
178
+ _message('first', id: 'user-1'),
179
+ _message(
180
+ 'first answer',
181
+ id: 'assistant-1',
182
+ role: MessageRole.assistant,
183
+ ),
184
+ _message('second', id: 'user-2'),
185
+ _message('', id: 'assistant-2', role: MessageRole.assistant),
186
+ ],
187
+ params: ChatParameters.defaults(),
188
+ )
189
+ .listen((_) {})
190
+ .asFuture<void>();
191
+
192
+ await _waitFor(
193
+ () =>
194
+ inferenceService.sessions.length == 2 &&
195
+ inferenceService.sessions.last.didGenerate,
196
+ );
197
+ final secondSession = inferenceService.sessions.last;
198
+ expect(secondSession.messages.single.text, contains('first answer'));
199
+ expect(secondSession.messages.single.text, contains('second'));
200
+ secondSession.responses.add(const gemma.TextResponse('second answer'));
201
+ await secondSession.responses.close();
202
+ await secondDone;
203
+
204
+ expect(inferenceService.createCount, 2);
205
+ },
206
+ );
207
+
208
test(
209
'passes system content once and never stages it as a chat turn',
210
() async {
@@ -689,10 +754,14 @@ class _FakeInferenceService implements OnDeviceInferenceService {
754
final List<String?> systemInstructions = [];
755
final List<bool?> supportImages = [];
756
int createCount = 0;
757
+ bool canReuseChatSession = true;
758
759
@override
760
bool get isLoaded => true;
761
762
+ @override
763
+ bool get supportsChatSessionReuse => canReuseChatSession;
764
+
765
@override
766
bool currentModelSupportsVision = false;
767