Fix surface creation and update notification bugs

packages/genui/lib/src/core/genui_manager.dart

@@ -158,14 +158,11 @@ class GenUiManager implements GenUiHost {
   void handleMessage(A2uiMessage message) {
     switch (message) {
       case SurfaceUpdate():
-        // No need for SurfaceAdded here because A2uiMessage will never generate
-        // those. We decide here if the surface is new or not, and generate a
-        // SurfaceAdded event if so.
         final String surfaceId = message.surfaceId;
         final ValueNotifier<UiDefinition?> notifier = getSurfaceNotifier(
           surfaceId,
         );
-        final isNew = notifier.value == null;
+
         UiDefinition uiDefinition =
             notifier.value ?? UiDefinition(surfaceId: surfaceId);
         final Map<String, Component> newComponents = Map.of(
@@ -176,26 +173,33 @@ class GenUiManager implements GenUiHost {
         }
         uiDefinition = uiDefinition.copyWith(components: newComponents);
         notifier.value = uiDefinition;
-        if (isNew) {
-          genUiLogger.info('Adding surface $surfaceId');
-          _surfaceUpdates.add(SurfaceAdded(surfaceId, uiDefinition));
-        } else {
+
+        // Notify UI ONLY if rendering has begun (i.e., rootComponentId is set)
+        if (uiDefinition.rootComponentId != null) {
           genUiLogger.info('Updating surface $surfaceId');
           _surfaceUpdates.add(SurfaceUpdated(surfaceId, uiDefinition));
+        } else {
+          genUiLogger.info(
+            'Caching components for surface $surfaceId (pre-rendering)',
+          );
         }
       case BeginRendering():
-        dataModelForSurface(message.surfaceId);
+        final String surfaceId = message.surfaceId;
+        dataModelForSurface(surfaceId);
         final ValueNotifier<UiDefinition?> notifier = getSurfaceNotifier(
-          message.surfaceId,
+          surfaceId,
         );
+
+        // Update the definition with the root component
         final UiDefinition uiDefinition =
-            notifier.value ?? UiDefinition(surfaceId: message.surfaceId);
+            notifier.value ?? UiDefinition(surfaceId: surfaceId);
         final UiDefinition newUiDefinition = uiDefinition.copyWith(
           rootComponentId: message.root,
         );
         notifier.value = newUiDefinition;
-        genUiLogger.info('Started rendering ${message.surfaceId}');
-        _surfaceUpdates.add(SurfaceUpdated(message.surfaceId, newUiDefinition));
+
+        genUiLogger.info('Creating and rendering surface $surfaceId');
+        _surfaceUpdates.add(SurfaceAdded(surfaceId, newUiDefinition));
       case DataModelUpdate():
         final String path = message.path ?? '/';
         genUiLogger.info(
@@ -205,6 +209,15 @@ class GenUiManager implements GenUiHost {
         );
         final DataModel dataModel = dataModelForSurface(message.surfaceId);
         dataModel.update(DataPath(path), message.contents);
+
+        // Notify UI of an update if the surface is already rendering
+        final ValueNotifier<UiDefinition?> notifier = getSurfaceNotifier(
+          message.surfaceId,
+        );
+        final UiDefinition? uiDefinition = notifier.value;
+        if (uiDefinition != null && uiDefinition.rootComponentId != null) {
+          _surfaceUpdates.add(SurfaceUpdated(message.surfaceId, uiDefinition));
+        }
       case SurfaceDeletion():
         final String surfaceId = message.surfaceId;
         if (_surfaces.containsKey(surfaceId)) {

packages/genui/test/core/genui_manager_test.dart

@@ -41,24 +41,19 @@ void main() {
         ),
       ];
 
-      final Future<GenUiUpdate> futureAdded = manager.surfaceUpdates.first;
       manager.handleMessage(
         SurfaceUpdate(surfaceId: surfaceId, components: components),
       );
-      final GenUiUpdate addedUpdate = await futureAdded;
-      expect(addedUpdate, isA<SurfaceAdded>());
-      expect(addedUpdate.surfaceId, surfaceId);
 
-      final Future<GenUiUpdate> futureUpdated = manager.surfaceUpdates.first;
+      final Future<GenUiUpdate> futureUpdate = manager.surfaceUpdates.first;
       manager.handleMessage(
         const BeginRendering(surfaceId: surfaceId, root: 'root'),
       );
-      final GenUiUpdate updatedUpdate = await futureUpdated;
+      final GenUiUpdate update = await futureUpdate;
 
-      expect(updatedUpdate, isA<SurfaceUpdated>());
-      expect(updatedUpdate.surfaceId, surfaceId);
-      final UiDefinition definition =
-          (updatedUpdate as SurfaceUpdated).definition;
+      expect(update, isA<SurfaceAdded>());
+      expect(update.surfaceId, surfaceId);
+      final UiDefinition definition = (update as SurfaceAdded).definition;
       expect(definition, isNotNull);
       expect(definition.rootComponentId, 'root');
       expect(manager.surfaces[surfaceId]!.value, isNotNull);
@@ -77,10 +72,6 @@ void main() {
             },
           ),
         ];
-        manager.handleMessage(
-          SurfaceUpdate(surfaceId: surfaceId, components: oldComponents),
-        );
-
         final newComponents = [
           const Component(
             id: 'root',
@@ -90,18 +81,22 @@ void main() {
           ),
         ];
 
-        final Future<GenUiUpdate> futureUpdate = manager.surfaceUpdates.first;
+        final Future<void> expectation = expectLater(
+          manager.surfaceUpdates,
+          emitsInOrder([isA<SurfaceAdded>(), isA<SurfaceUpdated>()]),
+        );
+
+        manager.handleMessage(
+          SurfaceUpdate(surfaceId: surfaceId, components: oldComponents),
+        );
+        manager.handleMessage(
+          const BeginRendering(surfaceId: surfaceId, root: 'root'),
+        );
         manager.handleMessage(
           SurfaceUpdate(surfaceId: surfaceId, components: newComponents),
         );
-        final GenUiUpdate update = await futureUpdate;
-
-        expect(update, isA<SurfaceUpdated>());
-        expect(update.surfaceId, surfaceId);
-        final UiDefinition updatedDefinition =
-            (update as SurfaceUpdated).definition;
-        expect(updatedDefinition.components['root'], newComponents[0]);
-        expect(manager.surfaces[surfaceId]!.value, updatedDefinition);
+
+        await expectation;
       },
     );
 

packages/genui/test/ui_tools_test.dart

@@ -64,6 +64,9 @@ void main() {
       );
 
       await tool.invoke(args);
+      genUiManager.handleMessage(
+        const BeginRendering(surfaceId: 'testSurface', root: 'root'),
+      );
 
       await future;
     });
@@ -99,7 +102,7 @@ void main() {
       final Future<void> future = expectLater(
         genUiManager.surfaceUpdates,
         emits(
-          isA<SurfaceUpdated>()
+          isA<SurfaceAdded>()
               .having((e) => e.surfaceId, surfaceIdKey, 'testSurface')
               .having(
                 (e) => e.definition.rootComponentId,