fix: handle a reply cut off part-way through an A2UI message

README.md

@@ -228,10 +228,10 @@ just as well as in the browser.
 nothing is parsed until something listens. Three kinds of event arrive:
 `GenUiText` carries a piece of prose, `GenUiSurface` carries a surface the model
 has just opened, once, and `GenUiError` carries a message the model got wrong,
-such as a component the catalog lacks or a surface created twice, already in the
-shape `a2uiErrorMessage` sends back to the model. The stream carries on after a
-`GenUiError`. A failure of the model call itself is an error on the stream,
-which then ends.
+such as a component the catalog lacks, a surface created twice, or a message cut
+off when the reply ended, already in the shape `a2uiErrorMessage` sends back to
+the model. The stream carries on after a `GenUiError`. A failure of the model
+call itself is an error on the stream, which then ends.
 
 `ReplyBuilder` is the fold most apps want: it turns the events into a `Reply`
 with `text`, `surfaces`, `errors`, `failure`, and `isComplete`, and rebuilds its

example/lib/chat.dart

@@ -6,6 +6,7 @@ import 'package:genui_jaspr_example/interaction.dart';
 import 'package:genui_jaspr_example/server/chat_path.dart';
 import 'package:jaspr/dom.dart';
 import 'package:jaspr/jaspr.dart';
+import 'package:universal_web/js_interop.dart';
 import 'package:universal_web/web.dart' as web;
 
 /// One entry in the transcript: something the user said, or a model's reply.
@@ -24,6 +25,10 @@ class Turn {
   final Stream<GenUiEvent>? events;
 }
 
+/// What the user is shown when a reply fails, whatever the cause.
+const String failedTurnMessage =
+    'The model stopped before finishing. Try again, or ask a different way.';
+
 /// Streams the model's reply to [prompt] as text chunks.
 ///
 /// The seam [ChatView] is tested at: production passes the genkit client,
@@ -177,7 +182,15 @@ class _ChatViewState extends State<ChatView> {
     setState(() {
       _busy = false;
       final failure = reply.failure;
-      if (failure != null) _error = '$failure';
+      if (failure != null) {
+        // The browser never learns why a call failed: the server turns every
+        // cause into the same 500. So the user gets a plain message, and the
+        // failure itself goes to the console for whoever is debugging.
+        _error = failedTurnMessage;
+        // coverage:ignore-start
+        web.console.error('Reply failed: $failure'.toJS);
+        // coverage:ignore-end
+      }
       // A turn with neither words nor a surface would render as an empty
       // bubble, which is what a failed request used to leave behind.
       if (reply.isEmpty) _turns.remove(turn);

example/lib/prompt.dart

@@ -16,4 +16,8 @@ const _role = '''
 You build user interfaces by emitting A2UI messages, a JSON protocol the client
 renders into real HTML. Reply with a short sentence for the user, then the JSON
 messages that build the interface. Keep the sentence to one or two lines: the
-interface carries the detail, not the prose.''';
+interface carries the detail, not the prose.
+
+Write every piece of text in your own words, in the sentence and in the
+interface alike. Do not quote descriptions, articles, or other reference text
+word for word.''';

example/lib/server/chat_agent.dart

@@ -1,3 +1,5 @@
+import 'dart:io';
+
 import 'package:genkit/genkit.dart';
 import 'package:genkit_shelf/genkit_shelf.dart';
 import 'package:genui_jaspr_example/prompt.dart';
@@ -34,8 +36,12 @@ Agent<dynamic> chatAgent(
 /// all come from `genkit_shelf`, so nothing here has to know how a frame looks.
 /// Mount it in front of the page: `main.server.dart` routes everything under
 /// [chatPath] here and lets the rest fall through to the rendered app.
+///
+/// A turn that fails is logged to stderr with its cause and stack trace.
+/// `genkit_shelf` tells the browser only "Internal server error", so without
+/// this the reason a reply broke would be written down nowhere.
 Handler chatHandler(Agent<dynamic> agent) {
-  final turn = shelfHandler(agent.action);
+  final turn = shelfHandler(_withFailureLogging(agent.action));
   final snapshot = shelfHandler(agent.getSnapshotDataAction);
   final abort = shelfHandler(agent.abortAgentAction);
 
@@ -46,3 +52,58 @@ Handler chatHandler(Agent<dynamic> agent) {
     _ => Response.notFound('No such route.'),
   };
 }
+
+/// The agent's turn [action], writing any failure to stderr.
+///
+/// The failure has to be caught here, around the action itself. A model call
+/// that breaks does not throw out of the agent: the turn ends with its error in
+/// [AgentOutput.error], and `genkit_shelf` then sends the browser a frame that
+/// hides the cause. The error still carries the original exception at this
+/// point, so this is the last place it can be written down. Genkit keeps the
+/// exception that caused a failure but not the stack trace it was thrown with,
+/// so a stack is logged only when that exception is an [Error], which carries
+/// its own. A model plugin's crash, such as the null check Gemini's recitation
+/// stop trips, is one.
+///
+/// Genkit's [Action] has no `copyWith`, so the wrapper copies each field.
+/// Anything it leaves out is lost to `genkit_shelf`.
+Action<AgentInput, AgentOutput, AgentStreamChunk, AgentInit>
+_withFailureLogging(
+  Action<AgentInput, AgentOutput, AgentStreamChunk, AgentInit> action,
+) {
+  return Action(
+    name: action.name,
+    actionType: action.actionType,
+    description: action.description,
+    inputSchema: action.inputSchema,
+    outputSchema: action.outputSchema,
+    streamSchema: action.streamSchema,
+    initSchema: action.initSchema,
+    metadata: action.metadata,
+    fn: (input, context) async {
+      final output = await action.fn(input, context);
+      final error = output.error;
+      if (error != null) {
+        final cause = error.details;
+        _logFailure(
+          '${error.status}: ${error.message}',
+          // With nothing underneath, Genkit fills the details in with the
+          // message itself, which would only repeat the summary.
+          cause: cause == error.message ? null : cause,
+          stackTrace: cause is Error ? cause.stackTrace : null,
+        );
+      }
+      return output;
+    },
+  );
+}
+
+void _logFailure(String summary, {Object? cause, StackTrace? stackTrace}) {
+  stderr.writeln(
+    [
+      'Chat turn failed: $summary',
+      if (cause != null) 'Cause: $cause',
+      ?stackTrace,
+    ].join('\n'),
+  );
+}

example/test/browser/chat_browser_test.dart

@@ -9,6 +9,10 @@ import 'package:jaspr_test/client_test.dart';
 /// A model reply carrying one A2UI message, fenced the way the parser expects.
 String fenced(String json) => '```json\n$json\n```\n';
 
+/// What the user is shown when a reply fails, whatever the cause.
+const plainFailureText =
+    'The model stopped before finishing. Try again, or ask a different way.';
+
 void main() {
   group('ChatView in a browser', () {
     testClient('typing and sending renders the streamed reply', (tester) async {
@@ -109,10 +113,39 @@ void main() {
       await tester.input(find.tag('input'), value: 'hi');
       await tester.click(find.tag('button'));
 
-      expect(find.textContaining('model unavailable'), findsOneComponent);
+      expect(find.text(plainFailureText), findsOneComponent);
+      // The raw failure is for developers, in the console, not for the user.
+      expect(find.textContaining('model unavailable'), findsNothing);
       // The user's own turn stays; the failed reply leaves no empty bubble.
       expect(find.text('hi'), findsOneComponent);
       expect(find.byType(Surface), findsNothing);
     });
+
+    testClient('a reply cut off mid-message keeps its prose, not the JSON', (
+      tester,
+    ) async {
+      // What Gemini does when its recitation filter stops a reply: the prose
+      // arrives, the message starts, and then the call fails.
+      tester.pumpComponent(
+        ChatView(
+          send: (prompt) async* {
+            yield 'Here is Paris.\n\n```json\n';
+            yield '{"version":"v0.9","updateComponents":{"surfaceId":"s",'
+                '"components":[{"id":"root","component":"Text",'
+                '"text":"The 12th-';
+            throw StateError('GenkitException: Internal server error');
+          },
+        ),
+      );
+
+      await tester.input(find.tag('input'), value: 'Tell me about Paris');
+      await tester.click(find.tag('button'));
+
+      expect(find.textContaining('Here is Paris.'), findsOneComponent);
+      expect(find.textContaining('ended part-way'), findsOneComponent);
+      expect(find.text(plainFailureText), findsOneComponent);
+      expect(find.textContaining('updateComponents'), findsNothing);
+      expect(find.textContaining('GenkitException'), findsNothing);
+    });
   });
 }

example/test/round_trip_test.dart

@@ -9,6 +9,7 @@ import 'package:genui_jaspr_example/server/chat_path.dart';
 import 'package:jaspr/server.dart';
 import 'package:jaspr_test/jaspr_test.dart';
 import 'package:shelf/shelf_io.dart' as shelf_io;
+import 'package:test/fake.dart';
 
 /// What a model would reply with, standing in for the model itself.
 final reply =
@@ -27,6 +28,23 @@ Here's a short form.
 ]}}
 ```''';
 
+/// The prompt that makes the stand-in model fail part-way through its reply,
+/// the way Gemini's plugin does when a recitation filter stops it.
+const failPartWay = 'fail part-way';
+
+/// Stands in for the server's stderr, keeping what it was sent.
+class CapturedStderr extends Fake implements Stdout {
+  final _written = StringBuffer();
+
+  /// Everything written so far.
+  String get written => _written.toString();
+
+  void clear() => _written.clear();
+
+  @override
+  void writeln([Object? object = '']) => _written.writeln(object);
+}
+
 void main() {
   group('the whole round trip over HTTP', () {
     late HttpServer server;
@@ -36,6 +54,10 @@ void main() {
     /// agent told it: the system prompt, the history, the user's turn.
     final requests = <ModelRequest>[];
 
+    /// What the server wrote to stderr. Shelf runs each request in the zone
+    /// the server was started in, so that is where this is installed.
+    final serverStderr = CapturedStderr();
+
     setUpAll(() async {
       Jaspr.initializeApp();
 
@@ -48,6 +70,23 @@ void main() {
           name: 'canned',
           fn: (request, context) async {
             requests.add(request);
+            if (request.messages.last.text == failPartWay) {
+              context.sendChunk(
+                ModelResponseChunk(content: [TextPart(text: 'Here is ')]),
+              );
+              // What genkit_google_genai throws on an empty candidate: a null
+              // check fails, and the plugin wraps it as INTERNAL.
+              try {
+                <String, String?>{}['role']!;
+              } on Object catch (error, stackTrace) {
+                throw GenkitException(
+                  'Google AI Error: $error',
+                  status: StatusCodes.INTERNAL,
+                  underlyingException: error,
+                  stackTrace: stackTrace,
+                );
+              }
+            }
             if (context.streamingRequested) {
               // Chunked awkwardly on purpose, including a split inside a fence.
               for (var i = 0; i < reply.length; i += 7) {
@@ -75,17 +114,23 @@ void main() {
           },
         );
 
-      server = await shelf_io.serve(
-        chatHandler(chatAgent(ai, model: modelRef('canned'))),
-        InternetAddress.loopbackIPv4,
-        0,
+      server = await IOOverrides.runZoned(
+        () => shelf_io.serve(
+          chatHandler(chatAgent(ai, model: modelRef('canned'))),
+          InternetAddress.loopbackIPv4,
+          0,
+        ),
+        stderr: () => serverStderr,
       );
       url = 'http://${server.address.host}:${server.port}/$chatPath';
     });
 
     tearDownAll(() => server.close(force: true));
 
-    setUp(requests.clear);
+    setUp(() {
+      requests.clear();
+      serverStderr.clear();
+    });
 
     /// The browser's view of the agent: one chat, its session kept by the
     /// server.
@@ -131,6 +176,15 @@ void main() {
       expect(system, contains('"TextField"'));
     });
 
+    test('the model is asked to write in its own words', () async {
+      // Gemini stops a reply that reproduces text it was trained on, part-way
+      // through, so the prompt steers it away from quoting stock descriptions.
+      await turn('tell me about Paris');
+
+      final system = textsOf(requests.single, Role.system).join();
+      expect(system, contains('in your own words'));
+    });
+
     test('the reply becomes a rendered surface', () async {
       final result = await turn('make me a form');
 
@@ -217,6 +271,24 @@ void main() {
       await agent.abort(chat.snapshotId!);
     });
 
+    test('a turn that fails is logged with its cause on the server', () async {
+      // The browser only hears "Internal server error". The cause, and where
+      // it came from, belong in the terminal of whoever runs the server.
+      await expectLater(
+        ask(openChat(), failPartWay).drain<void>(),
+        throwsA(anything),
+      );
+
+      expect(serverStderr.written, contains('Null check operator'));
+      expect(serverStderr.written, contains('round_trip_test.dart'));
+    });
+
+    test('a turn that succeeds logs nothing', () async {
+      await turn('make me a form');
+
+      expect(serverStderr.written, isEmpty);
+    });
+
     test('anything else under the path is not found', () async {
       final client = HttpClient();
       final request = await client.postUrl(Uri.parse('$url/nope'));

lib/src/transport/a2ui_parser_transformer.dart

@@ -11,6 +11,10 @@ import 'package:genui_jaspr/src/transport/generation_events.dart';
 /// network happens to break them, often in the middle of a JSON object. This
 /// buffers until a message is complete, so a surface can be updated as soon
 /// as one arrives rather than after the whole response.
+///
+/// A stream that ends part-way through a message, because a safety filter or
+/// a token limit stopped the model or the call failed, reports the fragment as
+/// an [A2uiValidationException] rather than showing it as prose.
 class A2uiParserTransformer
     extends StreamTransformerBase<String, GenerationEvent> {
   /// Creates an [A2uiParserTransformer].
@@ -56,16 +60,44 @@ class _ParserStream {
   }
 
   void _onDone() {
-    // Whatever is left is prose, unless it is only the whitespace a model puts
-    // after its last message, which the separator rule drops everywhere else.
-    if (_buffer.isNotEmpty &&
-        !(_lastEventWasMessage && _buffer.trim().isEmpty)) {
-      _emitText(_buffer);
-    }
+    final leftover = _buffer;
     _buffer = '';
+    if (_isUnfinishedMessage(leftover)) {
+      // The model stopped part-way through a message, cut off by a safety
+      // filter, a token limit, or a failed call. Printed as prose it would
+      // put raw JSON in front of the user, so it is reported as the broken
+      // message it is.
+      _controller.addError(
+        A2uiValidationException(
+          'The reply ended part-way through an A2UI message',
+          json: leftover,
+        ),
+      );
+    } else if (leftover.isNotEmpty &&
+        !(_lastEventWasMessage && leftover.trim().isEmpty)) {
+      // Whatever is left is prose, unless it is only the whitespace a model
+      // puts after its last message, which the separator rule drops
+      // everywhere else.
+      _emitText(leftover);
+    }
     unawaited(_controller.close());
   }
 
+  /// Whether [leftover], what the buffer still holds when the stream ends, is
+  /// the start of a message rather than prose.
+  ///
+  /// Besides a possible message, which opens with a fence or a brace, the
+  /// buffer can hold back a partial fence marker or the whitespace after a
+  /// message. Neither is a message, and both are left to the prose rule. A
+  /// `json` fence says it is a message outright. A brace or any other fence
+  /// counts only once it names a message, so a reply that ends on a stray
+  /// brace or an unrelated code block still shows it.
+  bool _isUnfinishedMessage(String leftover) {
+    if (leftover.startsWith('```json')) return true;
+    if (!leftover.startsWith('```') && !leftover.startsWith('{')) return false;
+    return _messageKeys.any((key) => leftover.contains('"$key"'));
+  }
+
   void _processBuffer() {
     while (_buffer.isNotEmpty) {
       if (_consumeFenced()) continue;