From 2dee88f1b431f284c61b1242b11f378533ae7cbf Mon Sep 17 00:00:00 2001
From: pckimlong <pochkimlong88@gmail.com>
Date: Sat, 10 Jan 2026 11:49:52 +0700
Subject: [PATCH 01/10] feat(kai_engine): add CoreMessageRepository with
 callback-based persistence

- Add CoreMessageRepository implementation with simplified callback API
- Add automatic timestamp-based sorting for InMemoryMessageRepository
- Add onInitial callback for loading messages from persistence
- Add onPut callback for persisting insert/update operations (upsert)
- Add optional onRemove callback for persisting deletions
- Add session caching and initialization state tracking
- Add comprehensive test coverage for both
---
 .../lib/src/message_repository_base.dart      |  82 +++-
 .../src/message_repository_base_test.dart     | 431 ++++++++++++++++++
 2 files changed, 510 insertions(+), 3 deletions(-)
 create mode 100644 packages/kai_engine/test/src/message_repository_base_test.dart

diff --git a/packages/kai_engine/lib/src/message_repository_base.dart b/packages/kai_engine/lib/src/message_repository_base.dart
index baf70f5..3226c60 100644
--- a/packages/kai_engine/lib/src/message_repository_base.dart
+++ b/packages/kai_engine/lib/src/message_repository_base.dart
@@ -20,13 +20,14 @@ abstract interface class MessageRepositoryBase<T> {
 }
 
 /// Abstract interface for a core message repository. No adapter required.
-abstract interface class CoreMessageRepositoryBase
-    extends MessageRepositoryBase<CoreMessage> {}
+abstract interface class CoreMessageRepositoryBase extends MessageRepositoryBase<CoreMessage> {}
 
 /// Prebuilt repository for memory persistence.
+/// Messages are automatically sorted by timestamp on initialization.
 final class InMemoryMessageRepository implements CoreMessageRepositoryBase {
   InMemoryMessageRepository({Iterable<CoreMessage>? initialMessages})
-    : _messages = initialMessages?.toList() ?? [];
+    : _messages = (initialMessages?.toList() ?? [])
+        ..sort((a, b) => a.timestamp.compareTo(b.timestamp));
 
   final List<CoreMessage> _messages;
 
@@ -62,3 +63,78 @@ final class InMemoryMessageRepository implements CoreMessageRepositoryBase {
     return Future.value(messages);
   }
 }
+
+/// Prebuilt repository for callback-based persistence.
+///
+/// Simplified API with only 2 required callbacks:
+/// - [onInitial]: Load initial messages from persistence
+/// - [onPut]: Persist insert or update operations (upsert)
+/// - [onRemove]: Optional callback for persisting deletions
+///
+/// Note: Callbacks are for persistence only. The repository maintains
+/// its own internal state for immediate read operations. Messages are
+/// automatically sorted by timestamp on initial load.
+final class CoreMessageRepository implements CoreMessageRepositoryBase {
+  final Future<Iterable<CoreMessage>> Function(ConversationSession) onInitial;
+  final Future<Iterable<CoreMessage>> Function(ConversationSession, Iterable<CoreMessage>) onPut;
+  final Future<void> Function(Iterable<CoreMessage>)? onRemove;
+
+  final List<CoreMessage> _messages = [];
+  ConversationSession? _cachedSession;
+  bool _isInitialized = false;
+
+  CoreMessageRepository({required this.onInitial, required this.onPut, this.onRemove});
+
+  @override
+  Future<Iterable<CoreMessage>> getMessages(ConversationSession session) async {
+    if (!_isInitialized || _cachedSession != session) {
+      _cachedSession = session;
+      final loaded = await onInitial(session);
+      _messages.clear();
+      _messages.addAll(loaded);
+      _messages.sort((a, b) => a.timestamp.compareTo(b.timestamp));
+      _isInitialized = true;
+    }
+    return _messages;
+  }
+
+  @override
+  Future<Iterable<CoreMessage>> saveMessages({
+    required ConversationSession session,
+    required Iterable<CoreMessage> messages,
+  }) async {
+    if (_cachedSession != null && _cachedSession != session) {
+      _cachedSession = session;
+      _messages.clear();
+    } else {
+      _cachedSession = session;
+    }
+    _isInitialized = true;
+    _messages.addAll(messages);
+    await onPut(session, messages);
+    return messages;
+  }
+
+  @override
+  Future<Iterable<CoreMessage>> updateMessages(Iterable<CoreMessage> messages) async {
+    if (_cachedSession == null) {
+      throw StateError('Session not initialized. Call getMessages or saveMessages first.');
+    }
+    for (var message in messages) {
+      final index = _messages.indexWhere((m) => m.messageId == message.messageId);
+      if (index != -1) {
+        _messages[index] = message;
+      }
+    }
+    await onPut(_cachedSession!, messages);
+    return messages;
+  }
+
+  @override
+  Future<void> removeMessages(Iterable<CoreMessage> messages) async {
+    _messages.removeWhere(
+      (message) => messages.map((e) => e.messageId).contains(message.messageId),
+    );
+    await onRemove?.call(messages);
+  }
+}
diff --git a/packages/kai_engine/test/src/message_repository_base_test.dart b/packages/kai_engine/test/src/message_repository_base_test.dart
new file mode 100644
index 0000000..cfccd5d
--- /dev/null
+++ b/packages/kai_engine/test/src/message_repository_base_test.dart
@@ -0,0 +1,431 @@
+import 'package:kai_engine/src/message_repository_base.dart';
+import 'package:kai_engine/src/models/conversation_session.dart';
+import 'package:kai_engine/src/models/core_message.dart';
+import 'package:test/test.dart';
+
+void main() {
+  group('InMemoryMessageRepository', () {
+    test('should sort messages by timestamp on initialization', () async {
+      final now = DateTime.now();
+      final msg1 = CoreMessage.user(messageId: '1', content: 'First');
+      final msg2 = CoreMessage.ai(messageId: '2', content: 'Second');
+      final msg3 = CoreMessage.user(messageId: '3', content: 'Third');
+
+      // Create messages with different timestamps
+      final messages = [
+        msg3.copyWith(timestamp: now.add(Duration(minutes: 2))),
+        msg1.copyWith(timestamp: now),
+        msg2.copyWith(timestamp: now.add(Duration(minutes: 1))),
+      ];
+
+      final repo = InMemoryMessageRepository(initialMessages: messages);
+      final session = ConversationSession.withCurrentTime(id: 'test');
+      final result = await repo.getMessages(session);
+
+      // Should be sorted by timestamp
+      expect(result.first.messageId, equals('1'));
+      expect(result.elementAt(1).messageId, equals('2'));
+      expect(result.last.messageId, equals('3'));
+    });
+  });
+
+  group('CoreMessageRepository', () {
+    late ConversationSession session;
+    late CoreMessage testMessage1;
+    late CoreMessage testMessage2;
+
+    setUp(() {
+      session = ConversationSession.withCurrentTime(id: 'test-session');
+      testMessage1 = CoreMessage.user(messageId: 'msg-1', content: 'Hello');
+      testMessage2 = CoreMessage.ai(messageId: 'msg-2', content: 'Hi there');
+    });
+
+    group('constructor', () {
+      test('should create instance with required callbacks', () {
+        final repo = CoreMessageRepository(
+          onInitial: (s) async => [],
+          onPut: (s, msgs) async => msgs,
+        );
+
+        expect(repo, isA<CoreMessageRepository>());
+      });
+
+      test('should create instance with optional onRemove callback', () {
+        final repo = CoreMessageRepository(
+          onInitial: (s) async => [],
+          onPut: (s, msgs) async => msgs,
+          onRemove: (msgs) async {},
+        );
+
+        expect(repo, isA<CoreMessageRepository>());
+      });
+    });
+
+    group('getMessages', () {
+      test('should call onInitial callback with session and populate internal state', () async {
+        var capturedSession;
+        final repo = CoreMessageRepository(
+          onInitial: (s) {
+            capturedSession = s;
+            return Future.value([testMessage1]);
+          },
+          onPut: (s, msgs) async => msgs,
+        );
+
+        final result = await repo.getMessages(session);
+
+        expect(capturedSession, equals(session));
+        expect(result, equals([testMessage1]));
+      });
+
+      test('should sort messages by timestamp on initial load', () async {
+        final now = DateTime.now();
+        final msg1 = CoreMessage.user(messageId: '1', content: 'First');
+        final msg2 = CoreMessage.ai(messageId: '2', content: 'Second');
+        final msg3 = CoreMessage.user(messageId: '3', content: 'Third');
+
+        // Return messages out of order
+        final repo = CoreMessageRepository(
+          onInitial: (s) async => [
+            msg3.copyWith(timestamp: now.add(Duration(minutes: 2))),
+            msg1.copyWith(timestamp: now),
+            msg2.copyWith(timestamp: now.add(Duration(minutes: 1))),
+          ],
+          onPut: (s, msgs) async => msgs,
+        );
+
+        final result = await repo.getMessages(session);
+
+        // Should be sorted by timestamp
+        expect(result.first.messageId, equals('1'));
+        expect(result.elementAt(1).messageId, equals('2'));
+        expect(result.last.messageId, equals('3'));
+      });
+
+      test('should populate internal state from onInitial callback', () async {
+        final repo = CoreMessageRepository(
+          onInitial: (s) async => [testMessage1, testMessage2],
+          onPut: (s, msgs) async => msgs,
+        );
+
+        await repo.getMessages(session);
+        final result = await repo.getMessages(session);
+
+        expect(result, equals([testMessage1, testMessage2]));
+      });
+
+      test('should cache state and not reload on subsequent getMessages calls', () async {
+        var callCount = 0;
+        final repo = CoreMessageRepository(
+          onInitial: (s) async {
+            callCount++;
+            return callCount == 1 ? [testMessage1] : [testMessage2];
+          },
+          onPut: (s, msgs) async => msgs,
+        );
+
+        final result1 = await repo.getMessages(session);
+        final result2 = await repo.getMessages(session);
+
+        expect(callCount, equals(1)); // Only called once
+        expect(result1, equals([testMessage1]));
+        expect(result2, equals([testMessage1])); // Returns cached state
+      });
+
+      test('should cache session for later use', () async {
+        final repo = CoreMessageRepository(
+          onInitial: (s) async => [testMessage1],
+          onPut: (s, msgs) async => msgs,
+        );
+
+        await repo.getMessages(session);
+
+        final updatedMessage = testMessage1.copyWith(content: 'Updated');
+        final result = await repo.updateMessages([updatedMessage]);
+
+        expect(result, contains(updatedMessage));
+      });
+    });
+
+    group('saveMessages', () {
+      test('should call onPut callback with session and messages for persistence', () async {
+        var capturedSession;
+        var capturedMessages;
+        final repo = CoreMessageRepository(
+          onInitial: (s) async => [],
+          onPut: (s, msgs) {
+            capturedSession = s;
+            capturedMessages = msgs;
+            return Future.value(msgs);
+          },
+        );
+
+        await repo.saveMessages(session: session, messages: [testMessage1]);
+
+        expect(capturedSession, equals(session));
+        expect(capturedMessages, equals([testMessage1]));
+      });
+
+      test('should add messages to internal state', () async {
+        final repo = CoreMessageRepository(
+          onInitial: (s) async => [],
+          onPut: (s, msgs) async => msgs,
+        );
+
+        await repo.saveMessages(session: session, messages: [testMessage1]);
+        final result = await repo.getMessages(session);
+
+        expect(result, equals([testMessage1]));
+      });
+
+      test('should accumulate messages in internal state across multiple saves', () async {
+        final repo = CoreMessageRepository(
+          onInitial: (s) async => [],
+          onPut: (s, msgs) async => msgs,
+        );
+
+        await repo.saveMessages(session: session, messages: [testMessage1]);
+        await repo.saveMessages(session: session, messages: [testMessage2]);
+        final result = await repo.getMessages(session);
+
+        expect(result, equals([testMessage1, testMessage2]));
+      });
+
+      test('should cache session for later use', () async {
+        final repo = CoreMessageRepository(
+          onInitial: (s) async => [],
+          onPut: (s, msgs) async => msgs,
+        );
+
+        await repo.saveMessages(session: session, messages: [testMessage1]);
+
+        final updatedMessage = testMessage1.copyWith(content: 'Updated');
+        final result = await repo.updateMessages([updatedMessage]);
+
+        expect(result, contains(updatedMessage));
+      });
+    });
+
+    group('updateMessages', () {
+      test('should call onPut callback with cached session and messages for persistence', () async {
+        var capturedSession;
+        var capturedMessages;
+        final repo = CoreMessageRepository(
+          onInitial: (s) async => [testMessage1],
+          onPut: (s, msgs) {
+            capturedSession = s;
+            capturedMessages = msgs;
+            return Future.value(msgs);
+          },
+        );
+
+        await repo.getMessages(session);
+        await repo.updateMessages([testMessage1]);
+
+        expect(capturedSession, equals(session));
+        expect(capturedMessages, equals([testMessage1]));
+      });
+
+      test('should update messages in internal state', () async {
+        final repo = CoreMessageRepository(
+          onInitial: (s) async => [testMessage1],
+          onPut: (s, msgs) async => msgs,
+        );
+
+        await repo.getMessages(session);
+        final updated = testMessage1.copyWith(content: 'Updated');
+        await repo.updateMessages([updated]);
+        final result = await repo.getMessages(session);
+
+        expect(result.first.content, equals('Updated'));
+      });
+
+      test('should throw StateError when session not initialized', () async {
+        final repo = CoreMessageRepository(
+          onInitial: (s) async => [],
+          onPut: (s, msgs) async => msgs,
+        );
+
+        expect(
+          () => repo.updateMessages([testMessage1]),
+          throwsA(
+            isA<StateError>().having(
+              (e) => e.message,
+              'message',
+              contains('Session not initialized'),
+            ),
+          ),
+        );
+      });
+
+      test('should use session from saveMessages if getMessages not called', () async {
+        final repo = CoreMessageRepository(
+          onInitial: (s) async => [],
+          onPut: (s, msgs) async => msgs,
+        );
+
+        await repo.saveMessages(session: session, messages: [testMessage1]);
+        final updated = testMessage1.copyWith(content: 'Updated');
+        await repo.updateMessages([updated]);
+        final result = await repo.getMessages(session);
+
+        expect(result.first.content, equals('Updated'));
+      });
+
+      test('should only update existing messages, not add new ones', () async {
+        final repo = CoreMessageRepository(
+          onInitial: (s) async => [testMessage1],
+          onPut: (s, msgs) async => msgs,
+        );
+
+        await repo.getMessages(session);
+        await repo.updateMessages([testMessage2]);
+        final result = await repo.getMessages(session);
+
+        expect(result, equals([testMessage1])); // testMessage2 not added
+      });
+    });
+
+    group('removeMessages', () {
+      test('should call onRemove callback when provided for persistence', () async {
+        var capturedMessages;
+        final repo = CoreMessageRepository(
+          onInitial: (s) async => [testMessage1, testMessage2],
+          onPut: (s, msgs) async => msgs,
+          onRemove: (msgs) {
+            capturedMessages = msgs;
+            return Future.value();
+          },
+        );
+
+        await repo.getMessages(session);
+        await repo.removeMessages([testMessage1]);
+
+        expect(capturedMessages, equals([testMessage1]));
+      });
+
+      test('should remove messages from internal state', () async {
+        final repo = CoreMessageRepository(
+          onInitial: (s) async => [testMessage1, testMessage2],
+          onPut: (s, msgs) async => msgs,
+        );
+
+        await repo.getMessages(session);
+        await repo.removeMessages([testMessage1]);
+        final result = await repo.getMessages(session);
+
+        expect(result, equals([testMessage2]));
+      });
+
+      test('should complete successfully when onRemove is null', () async {
+        final repo = CoreMessageRepository(
+          onInitial: (s) async => [testMessage1],
+          onPut: (s, msgs) async => msgs,
+        );
+
+        await repo.getMessages(session);
+        await expectLater(repo.removeMessages([testMessage1]), completes);
+      });
+
+      test('should handle multiple messages in removeMessages', () async {
+        final repo = CoreMessageRepository(
+          onInitial: (s) async => [testMessage1, testMessage2],
+          onPut: (s, msgs) async => msgs,
+        );
+
+        await repo.getMessages(session);
+        await repo.removeMessages([testMessage1, testMessage2]);
+        final result = await repo.getMessages(session);
+
+        expect(result, isEmpty);
+      });
+    });
+
+    group('integration scenarios', () {
+      test('should handle complete message lifecycle with persistence callbacks', () async {
+        final persistedMessages = <CoreMessage>[];
+
+        final repo = CoreMessageRepository(
+          onInitial: (s) async => persistedMessages.toList(),
+          onPut: (s, msgs) async {
+            // Simulate upsert in persistence layer
+            for (var msg in msgs) {
+              final index = persistedMessages.indexWhere((m) => m.messageId == msg.messageId);
+              if (index != -1) {
+                persistedMessages[index] = msg;
+              } else {
+                persistedMessages.add(msg);
+              }
+            }
+            return msgs;
+          },
+          onRemove: (msgs) async {
+            persistedMessages.removeWhere((m) => msgs.contains(m));
+          },
+        );
+
+        // Initial load
+        var current = await repo.getMessages(session);
+        expect(current, isEmpty);
+
+        // Save messages
+        await repo.saveMessages(session: session, messages: [testMessage1, testMessage2]);
+        current = await repo.getMessages(session);
+        expect(current.length, equals(2));
+
+        // Update message
+        final updated = testMessage1.copyWith(content: 'Updated content');
+        await repo.updateMessages([updated]);
+        current = await repo.getMessages(session);
+        expect(current.any((m) => m.content == 'Updated content'), isTrue);
+        expect(current.length, equals(2)); // Still 2 messages, one updated
+
+        // Remove message
+        await repo.removeMessages([testMessage2]);
+        current = await repo.getMessages(session);
+        expect(current.length, equals(1));
+        expect(current.first.messageId, equals(testMessage1.messageId));
+        expect(current.first.content, equals('Updated content'));
+      });
+
+      test('should handle append-only pattern without onRemove', () async {
+        final repo = CoreMessageRepository(
+          onInitial: (s) async => [],
+          onPut: (s, msgs) async => msgs,
+        );
+
+        await repo.saveMessages(session: session, messages: [testMessage1]);
+        await repo.saveMessages(session: session, messages: [testMessage2]);
+
+        final current = await repo.getMessages(session);
+        expect(current.length, equals(2));
+
+        // Remove removes from internal state even without onRemove callback
+        await repo.removeMessages([testMessage1]);
+        final afterRemove = await repo.getMessages(session);
+        expect(afterRemove.length, equals(1)); // 1 removed, 1 remains
+        expect(afterRemove.first.messageId, equals(testMessage2.messageId));
+      });
+
+      test('should maintain state independently of persistence layer', () async {
+        var persistenceCallCount = 0;
+
+        final repo = CoreMessageRepository(
+          onInitial: (s) async {
+            persistenceCallCount++;
+            return [];
+          },
+          onPut: (s, msgs) async => msgs,
+        );
+
+        // First load calls persistence
+        await repo.getMessages(session);
+        expect(persistenceCallCount, equals(1));
+
+        // Subsequent reads use internal state, don't call persistence
+        await repo.getMessages(session);
+        expect(persistenceCallCount, equals(1)); // Still 1, not 2
+      });
+    });
+  });
+}

From 49a07fadda639c565b675902e80ad8e1733a6c91 Mon Sep 17 00:00:00 2001
From: pckimlong <pochkimlong88@gmail.com>
Date: Sat, 10 Jan 2026 13:32:04 +0700
Subject: [PATCH 02/10] ``` refactor(kai_engine): add KaiChatController and
 improve code formatting

- Add KaiChatController as ready-to-use chat controller implementation
- Add comprehensive documentation with usage examples for basic, persistent storage, and custom context scenarios
- Add static create() factory method with async initialization support
- Add optional callbacks for onPromptsReady and onStateChange
- Add systemPrompt parameter to SimpleContextEngine for customization
- Update CoreMessageRepository callback
---
 .../lib/src/chat_controller_base.dart         | 213 +++++-
 .../lib/src/conversation_manager.dart         |  27 +-
 .../lib/src/message_repository_base.dart      |   6 +-
 .../kai_engine/lib/src/prompt_engine.dart     |  63 +-
 packages/kai_engine/lib/src/src.dart          |   4 +-
 .../test/src/chat_controller_base_test.dart   | 695 +++++++++++-------
 .../src/message_repository_base_test.dart     |  62 +-
 7 files changed, 676 insertions(+), 394 deletions(-)

diff --git a/packages/kai_engine/lib/src/chat_controller_base.dart b/packages/kai_engine/lib/src/chat_controller_base.dart
index f8cf028..ba08b6e 100644
--- a/packages/kai_engine/lib/src/chat_controller_base.dart
+++ b/packages/kai_engine/lib/src/chat_controller_base.dart
@@ -27,10 +27,9 @@ abstract base class ChatControllerBase<TEntity> {
        _cancelToken = CancelToken();
 
   /// Handle generation state updates
-  final _generationStateController =
-      BehaviorSubject<GenerationState<GenerationResult>>.seeded(
-        GenerationState.initial(),
-      );
+  final _generationStateController = BehaviorSubject<GenerationState<GenerationResult>>.seeded(
+    GenerationState.initial(),
+  );
 
   ContextEngine build();
 
@@ -74,9 +73,7 @@ abstract base class ChatControllerBase<TEntity> {
         // Exclude the current user message from history, by natural flow user message is happening after query input
         // but to trickly show message in UI faster, we trickly hide this
         histories: await _conversationManager.getMessages().then(
-          (messages) => messages
-              .whereNot((m) => m.messageId == userMessage.messageId)
-              .toIList(),
+          (messages) => messages.whereNot((m) => m.messageId == userMessage.messageId).toIList(),
         ),
       );
 
@@ -104,9 +101,7 @@ abstract base class ChatControllerBase<TEntity> {
       final insertedUserMessage = await insertedUserMessageFuture;
       userMessage = insertedUserMessage;
       if (contextResult.userMessage.content != insertedUserMessage.content) {
-        final revised = insertedUserMessage.copyWith(
-          content: contextResult.userMessage.content,
-        );
+        final revised = insertedUserMessage.copyWith(content: contextResult.userMessage.content);
         await _conversationManager.updateMessages(IList([revised]));
         userMessage = revised;
       }
@@ -124,9 +119,7 @@ abstract base class ChatControllerBase<TEntity> {
         onStateUpdate: (state) => _generationStateController.add(state),
       );
 
-      generationResult = await AIGenerationPhase(
-        _generationService,
-      ).execute(aiGenerationInput);
+      generationResult = await AIGenerationPhase(_generationService).execute(aiGenerationInput);
 
       // Save the AI-generated messages to the conversation
       if (generationResult.generatedMessages.isNotEmpty) {
@@ -137,16 +130,12 @@ abstract base class ChatControllerBase<TEntity> {
               final originalIds = generationResult.generatedMessages
                   .map((e) => e.messageId)
                   .toSet();
-              return items
-                  .where((e) => originalIds.contains(e.messageId))
-                  .toIList();
+              return items.where((e) => originalIds.contains(e.messageId)).toIList();
             });
 
         // Make sure generated message is updated to date with required metadata
         // before pass to other below service
-        generationResult = generationResult.copyWith(
-          generatedMessages: addedMessages,
-        );
+        generationResult = generationResult.copyWith(generatedMessages: addedMessages);
       }
 
       if (_postResponseEngine != null) {
@@ -161,9 +150,7 @@ abstract base class ChatControllerBase<TEntity> {
 
         await _postResponseEngine.execute(postResponseInput);
       }
-      final generationState = GenerationState<GenerationResult>.complete(
-        generationResult,
-      );
+      final generationState = GenerationState<GenerationResult>.complete(generationResult);
       _generationStateController.add(generationState);
       return generationState;
     } catch (error, stackTrace) {
@@ -184,10 +171,8 @@ abstract base class ChatControllerBase<TEntity> {
     _setState(GenerationState.error(KaiException.cancelled()));
   }
 
-  Stream<IList<CoreMessage>> get messagesStream =>
-      _conversationManager.messagesStream;
-  Future<IList<CoreMessage>> getAllMessages() =>
-      _conversationManager.getMessages();
+  Stream<IList<CoreMessage>> get messagesStream => _conversationManager.messagesStream;
+  Future<IList<CoreMessage>> getAllMessages() => _conversationManager.getMessages();
 
   Stream<GenerationState<GenerationResult>> get generationStateStream =>
       _generationStateController.stream;
@@ -203,8 +188,7 @@ abstract base class ChatControllerBase<TEntity> {
     onStateChange(state);
   }
 
-  void _setLoadingPhase(LoadingPhase phase) =>
-      _setState(GenerationState.loading(phase));
+  void _setLoadingPhase(LoadingPhase phase) => _setState(GenerationState.loading(phase));
 
   /// Callback to listen for state changes, this helpful for implementing custom logic or log the
   /// state provide more flexible. default to do nothing
@@ -214,3 +198,176 @@ abstract base class ChatControllerBase<TEntity> {
   /// this allow to track what being sent for generation. this helpful for logging etc
   void onPromptsReady(IList<CoreMessage> messages) {}
 }
+
+/// A ready-to-use chat controller implementation for AI-powered conversations.
+///
+/// `KaiChatController` is the primary entry point for integrating AI chat functionality
+/// into your application. It extends [ChatControllerBase] with sensible defaults
+/// and a convenient factory constructor for quick setup.
+///
+/// ## Features
+/// - **Simple initialization**: Use [create] factory for async setup with minimal configuration
+/// - **Flexible storage**: Supports both in-memory and persistent message storage
+/// - **Customizable context**: Configure system prompts or provide custom [ContextEngine]
+/// - **Extensible pipeline**: Add query preprocessing and post-response processing
+///
+/// ## Basic Usage
+/// ```dart
+/// final controller = await KaiChatController.create(
+///   generationService: myGenerationService,
+///   systemPrompt: 'You are a helpful assistant.',
+/// );
+///
+/// // Submit user input and get AI response
+/// final result = await controller.submit('Hello!');
+///
+/// // Listen to message updates
+/// controller.messagesStream.listen((messages) {
+///   print('Messages: $messages');
+/// });
+/// ```
+///
+/// ## With Persistent Storage
+/// ```dart
+/// final controller = await KaiChatController.create(
+///   sessionId: 'user-123-session-1',
+///   generationService: myGenerationService,
+///   repository: MyFirebaseMessageRepository(),
+///   systemPrompt: 'You are a helpful assistant.',
+/// );
+/// ```
+///
+/// ## With Custom Context Engine
+/// ```dart
+/// final controller = await KaiChatController.create(
+///   generationService: myGenerationService,
+///   contextEngine: ContextEngine.builder([
+///     PromptTemplate.system('You are an expert coder.'),
+///     PromptTemplate.buildParallel(UserProfileContext()),
+///     PromptTemplate.buildSequential(HistoryContext()),
+///     PromptTemplate.input(),
+///   ]),
+/// );
+/// ```
+///
+/// See also:
+/// - [ChatControllerBase] for the underlying orchestration logic
+/// - [ContextEngine] for customizing prompt construction
+/// - [ConversationManager] for message persistence details
+final class KaiChatController extends ChatControllerBase<CoreMessage> {
+  final ContextEngine _contextEngine;
+  final void Function(IList<CoreMessage> messages)? _onPromptsReady;
+  final void Function(GenerationState<GenerationResult> state)? _onStateChange;
+
+  KaiChatController._({
+    required ContextEngine contextEngine,
+    required super.conversationManager,
+    required super.queryEngine,
+    required super.generationService,
+    required super.postResponseEngine,
+    void Function(IList<CoreMessage> messages)? onPromptsReady,
+    void Function(GenerationState<GenerationResult> state)? onStateChange,
+  }) : _contextEngine = contextEngine,
+       _onPromptsReady = onPromptsReady,
+       _onStateChange = onStateChange;
+
+  @override
+  ContextEngine build() => _contextEngine;
+
+  /// Creates a new [KaiChatController] with async initialization.
+  ///
+  /// This factory method handles all the setup required for a functional chat controller,
+  /// including creating the conversation session and loading existing messages from storage.
+  ///
+  /// ## Parameters
+  ///
+  /// - **[sessionId]**: Unique identifier for the conversation session. Defaults to `'default'`.
+  ///   Use different session IDs to maintain separate conversation threads.
+  ///
+  /// - **[generationService]**: Required. The AI generation service that handles
+  ///   communication with the AI backend (e.g., OpenAI, Gemini, Claude).
+  ///
+  /// - **[repository]**: Optional message repository for persistent storage.
+  ///   If not provided, uses [InMemoryMessageRepository] (messages lost on app restart).
+  ///   Implement [CoreMessageRepository] for custom persistence (Firebase, SQLite, etc.).
+  ///
+  /// - **[contextEngine]**: Optional custom context engine for advanced prompt construction.
+  ///   Mutually exclusive with [systemPrompt] - provide one or the other, not both.
+  ///
+  /// - **[systemPrompt]**: Optional system prompt text for simple use cases.
+  ///   Creates a [SimpleContextEngine] internally. Cannot be used with [contextEngine].
+  ///
+  /// - **[queryEngine]**: Optional query preprocessor for input transformation,
+  ///   intent detection, or query augmentation before context building.
+  ///
+  /// - **[postEngine]**: Optional post-response processor for actions after AI generation,
+  ///   such as analytics, logging, or triggering side effects.
+  ///
+  /// - **[onPromptsReady]**: Optional callback invoked when prompts are ready for generation.
+  ///   Receives the list of messages that will be sent to the AI service.
+  ///
+  /// - **[onStateChange]**: Optional callback invoked when the generation state changes.
+  ///   Receives the current [GenerationState] for tracking progress.
+  ///
+  /// ## Returns
+  /// A fully initialized [KaiChatController] ready for use.
+  ///
+  /// ## Throws
+  /// - [AssertionError] if both [contextEngine] and [systemPrompt] are provided.
+  ///
+  /// ## Example
+  /// ```dart
+  /// final controller = await KaiChatController.create(
+  ///   sessionId: 'chat-${DateTime.now().millisecondsSinceEpoch}',
+  ///   generationService: GeminiGenerationService(apiKey: 'your-key'),
+  ///   repository: FirebaseMessageRepository(userId: currentUser.id),
+  ///   systemPrompt: 'You are a friendly cooking assistant.',
+  /// );
+  /// ```
+  static Future<KaiChatController> create({
+    String sessionId = 'default',
+    required GenerationServiceBase generationService,
+    CoreMessageRepository? repository,
+    ContextEngine? contextEngine,
+    String? systemPrompt,
+    QueryEngineBase? queryEngine,
+    PostResponseEngineBase? postEngine,
+    void Function(IList<CoreMessage> messages)? onPromptsReady,
+    void Function(GenerationState<GenerationResult> state)? onStateChange,
+  }) async {
+    assert(
+      contextEngine == null || systemPrompt == null,
+      'Cannot provide both contextEngine and systemPrompt. Use contextEngine for custom prompt building, or systemPrompt for simple cases.',
+    );
+
+    final session = ConversationSession.withCurrentTime(id: sessionId);
+
+    final conversation = await ConversationManager.create(
+      repository: repository ?? InMemoryMessageRepository(),
+      session: session,
+      messageAdapter: CoreMessageAdapter(),
+    );
+
+    return KaiChatController._(
+      contextEngine: contextEngine ?? SimpleContextEngine(systemPrompt: systemPrompt),
+      conversationManager: conversation,
+      queryEngine: queryEngine,
+      postResponseEngine: postEngine,
+      generationService: generationService,
+      onPromptsReady: onPromptsReady,
+      onStateChange: onStateChange,
+    );
+  }
+
+  @override
+  void onPromptsReady(IList<CoreMessage> messages) {
+    super.onPromptsReady(messages);
+    _onPromptsReady?.call(messages);
+  }
+
+  @override
+  void onStateChange(GenerationState<GenerationResult> state) {
+    super.onStateChange(state);
+    _onStateChange?.call(state);
+  }
+}
diff --git a/packages/kai_engine/lib/src/conversation_manager.dart b/packages/kai_engine/lib/src/conversation_manager.dart
index 72a864c..36b9274 100644
--- a/packages/kai_engine/lib/src/conversation_manager.dart
+++ b/packages/kai_engine/lib/src/conversation_manager.dart
@@ -18,9 +18,7 @@ class ConversationManager<T> {
   IList<CoreMessage> _messages = IList(const []);
   final BehaviorSubject<IList<CoreMessage>> _messagesController =
       BehaviorSubject<IList<CoreMessage>>();
-  final BehaviorSubject<bool> _loadingController = BehaviorSubject<bool>.seeded(
-    false,
-  );
+  final BehaviorSubject<bool> _loadingController = BehaviorSubject<bool>.seeded(false);
 
   ConversationManager._({
     required this.session,
@@ -59,9 +57,7 @@ class ConversationManager<T> {
   /// Adds a message to the conversation
   Future<IList<CoreMessage>> addMessages(IList<CoreMessage> messages) async {
     // Ensure system messages are ignored
-    final persistentMessages = messages
-        .where((m) => m.isMessageSavable)
-        .toList();
+    final persistentMessages = messages.where((m) => m.isMessageSavable).toList();
 
     // Store original state for rollback
     final originalMessages = _messages;
@@ -82,9 +78,7 @@ class ConversationManager<T> {
           .then((e) => e.map(_messageAdapter.toCoreMessage));
 
       _messages = _messages
-          .removeWhere(
-            (e) => persistentMessages.any((m) => m.messageId == e.messageId),
-          )
+          .removeWhere((e) => persistentMessages.any((m) => m.messageId == e.messageId))
           .addAll(result);
       _messagesController.add(_messages);
       return result.toIList();
@@ -97,9 +91,7 @@ class ConversationManager<T> {
   }
 
   Future<void> updateMessages(IList<CoreMessage> messages) async {
-    final persistentMessages = messages
-        .where((m) => m.isMessageSavable)
-        .toList();
+    final persistentMessages = messages.where((m) => m.isMessageSavable).toList();
 
     // Store original state for rollback
     final originalMessages = _messages;
@@ -117,9 +109,7 @@ class ConversationManager<T> {
       // Update repository
       final result = await _repository
           .updateMessages(
-            persistentMessages.map(
-              (e) => _messageAdapter.fromCoreMessage(e, session: session),
-            ),
+            persistentMessages.map((e) => _messageAdapter.fromCoreMessage(e, session: session)),
           )
           .then((e) => e.map(_messageAdapter.toCoreMessage));
 
@@ -154,9 +144,7 @@ class ConversationManager<T> {
 
       // Remove from repository
       await _repository.removeMessages(
-        messages.map(
-          (e) => _messageAdapter.fromCoreMessage(e, session: session),
-        ),
+        messages.map((e) => _messageAdapter.fromCoreMessage(e, session: session)),
       );
     } catch (error) {
       // Rollback optimistic update on failure
@@ -180,8 +168,7 @@ class ConversationManager<T> {
 }
 
 /// Prebuilt for in memory management of conversation messages
-final class InMemoryConversationManager
-    extends ConversationManager<CoreMessage> {
+final class InMemoryConversationManager extends ConversationManager<CoreMessage> {
   InMemoryConversationManager({
     required super.session,
     MessageRepositoryBase<CoreMessage>? repository,
diff --git a/packages/kai_engine/lib/src/message_repository_base.dart b/packages/kai_engine/lib/src/message_repository_base.dart
index 3226c60..42fed4a 100644
--- a/packages/kai_engine/lib/src/message_repository_base.dart
+++ b/packages/kai_engine/lib/src/message_repository_base.dart
@@ -75,9 +75,9 @@ final class InMemoryMessageRepository implements CoreMessageRepositoryBase {
 /// its own internal state for immediate read operations. Messages are
 /// automatically sorted by timestamp on initial load.
 final class CoreMessageRepository implements CoreMessageRepositoryBase {
-  final Future<Iterable<CoreMessage>> Function(ConversationSession) onInitial;
-  final Future<Iterable<CoreMessage>> Function(ConversationSession, Iterable<CoreMessage>) onPut;
-  final Future<void> Function(Iterable<CoreMessage>)? onRemove;
+  final Future<Iterable<CoreMessage>> Function(ConversationSession session) onInitial;
+  final Future<void> Function(ConversationSession session, Iterable<CoreMessage> messages) onPut;
+  final Future<void> Function(Iterable<CoreMessage> messages)? onRemove;
 
   final List<CoreMessage> _messages = [];
   ConversationSession? _cachedSession;
diff --git a/packages/kai_engine/lib/src/prompt_engine.dart b/packages/kai_engine/lib/src/prompt_engine.dart
index 6d32d3b..54d9a89 100644
--- a/packages/kai_engine/lib/src/prompt_engine.dart
+++ b/packages/kai_engine/lib/src/prompt_engine.dart
@@ -23,8 +23,7 @@ base class ContextEngineBuilder extends ContextEngine {
 /// system prompts, historical context, and user input according to a defined template structure.
 /// It supports both parallel and sequential context building strategies to optimize performance
 /// while maintaining logical ordering where required.
-abstract base class ContextEngine
-    extends KaiPhase<ContextEngineInput, ContextEngineOutput> {
+abstract base class ContextEngine extends KaiPhase<ContextEngineInput, ContextEngineOutput> {
   static ContextEngineBuilder builder(List<PromptTemplate> promptBuilder) {
     return ContextEngineBuilder(promptBuilder);
   }
@@ -42,10 +41,7 @@ abstract base class ContextEngine
       inputQuery: input.inputQuery,
       providedUserMessage: input.providedUserMessage,
     );
-    return ContextEngineOutput(
-      userMessage: result.userMessage,
-      prompts: result.prompts,
-    );
+    return ContextEngineOutput(userMessage: result.userMessage, prompts: result.prompts);
   }
 
   /// Generates a complete contextual prompt ready for AI consumption.
@@ -72,9 +68,7 @@ abstract base class ContextEngine
     CoreMessage? providedUserMessage,
   }) async {
     // Extract messageId for debug tracking
-    final userMessage =
-        providedUserMessage ??
-        CoreMessage.user(content: inputQuery.originalQuery);
+    final userMessage = providedUserMessage ?? CoreMessage.user(content: inputQuery.originalQuery);
     final messageId = userMessage.messageId;
 
     assert(
@@ -99,12 +93,7 @@ abstract base class ContextEngine
         .toList();
 
     // Process both concurrently
-    final parallelFuture = _buildParallelWithIndex(
-      parallelItems,
-      source,
-      inputQuery,
-      messageId,
-    );
+    final parallelFuture = _buildParallelWithIndex(parallelItems, source, inputQuery, messageId);
     final sequentialFuture = _buildSequentialWithIndex(
       sequentialItems,
       source,
@@ -173,9 +162,7 @@ abstract base class ContextEngine
         finalContexts.addAll(allResults[i]!);
       } else {
         // This should not happen - all template types should be handled above
-        throw StateError(
-          'Unhandled template type at index $i: ${template.runtimeType}',
-        );
+        throw StateError('Unhandled template type at index $i: ${template.runtimeType}');
       }
     }
 
@@ -196,8 +183,7 @@ abstract base class ContextEngine
   ///
   /// Returns:
   /// - A list of results with their original indices for proper reordering
-  Future<List<({int index, IList<CoreMessage> result})>>
-  _buildParallelWithIndex(
+  Future<List<({int index, IList<CoreMessage> result})>> _buildParallelWithIndex(
     List<({int index, PromptTemplate builder})> items,
     IList<CoreMessage> source,
     QueryContext inputQuery,
@@ -227,8 +213,7 @@ abstract base class ContextEngine
   ///
   /// Returns:
   /// - A list of results with their original indices for proper reordering
-  Future<List<({int index, IList<CoreMessage> result})>>
-  _buildSequentialWithIndex(
+  Future<List<({int index, IList<CoreMessage> result})>> _buildSequentialWithIndex(
     List<({int index, PromptTemplate builder})> items,
     IList<CoreMessage> source,
     QueryContext inputQuery,
@@ -240,11 +225,7 @@ abstract base class ContextEngine
     for (final item in items) {
       final sequentialBuilder = item.builder as _BuildSequentialPromptTemplate;
       final builder = sequentialBuilder.builder;
-      final context = await builder.build(
-        inputQuery,
-        messageId,
-        currentContext,
-      );
+      final context = await builder.build(inputQuery, messageId, currentContext);
 
       currentContext = context;
       results.add((index: item.index, result: context));
@@ -263,12 +244,16 @@ abstract base class ContextEngine
 ///
 /// This implementation is suitable for most basic conversational AI scenarios.
 final class SimpleContextEngine extends ContextEngine {
+  final String? systemPrompt;
+
   @override
   List<PromptTemplate> get promptBuilder => [
-    PromptTemplate.system("You're kai, a useful friendly personal assistant."),
+    PromptTemplate.system(systemPrompt ?? "You're kai, a useful friendly personal assistant."),
     PromptTemplate.buildSequential(HistoryContext()),
     PromptTemplate.input(),
   ];
+
+  SimpleContextEngine({this.systemPrompt});
 }
 
 @freezed
@@ -276,15 +261,11 @@ sealed class PromptTemplate with _$PromptTemplate {
   const PromptTemplate._();
 
   static PromptTemplate buildParallelFn(ParallelContextBuilderFn builder) {
-    return PromptTemplate.buildParallel(
-      _FunctionParallelContextBuilder(builder),
-    );
+    return PromptTemplate.buildParallel(_FunctionParallelContextBuilder(builder));
   }
 
   static PromptTemplate buildSequentialFn(SequentialContextBuilderFn builder) {
-    return PromptTemplate.buildSequential(
-      _FunctionSequentialContextBuilder(builder),
-    );
+    return PromptTemplate.buildSequential(_FunctionSequentialContextBuilder(builder));
   }
 
   /// Creates a system prompt template with fixed text content.
@@ -322,9 +303,8 @@ sealed class PromptTemplate with _$PromptTemplate {
   ///
   /// Parameters:
   /// - [builder]: The context builder that will generate the prompt content
-  const factory PromptTemplate.buildSequential(
-    SequentialContextBuilder builder,
-  ) = _BuildSequentialPromptTemplate;
+  const factory PromptTemplate.buildSequential(SequentialContextBuilder builder) =
+      _BuildSequentialPromptTemplate;
 
   /// Creates an input prompt template for the user's message.
   ///
@@ -340,11 +320,7 @@ sealed class PromptTemplate with _$PromptTemplate {
   /// - [input]: The user's raw input message
   /// - [messages]: The final list of messages in the conversation context after processing other templates include system prompts
   const factory PromptTemplate.input([
-    FutureOr<String> Function(
-      QueryContext input,
-      IList<CoreMessage> finalizedMessages,
-    )?
-    revision,
+    FutureOr<String> Function(QueryContext input, IList<CoreMessage> finalizedMessages)? revision,
   ]) = InputPromptTemplate;
 }
 
@@ -377,8 +353,7 @@ final class _FunctionParallelContextBuilder implements ParallelContextBuilder {
   }
 }
 
-final class _FunctionSequentialContextBuilder
-    implements SequentialContextBuilder {
+final class _FunctionSequentialContextBuilder implements SequentialContextBuilder {
   final SequentialContextBuilderFn _builder;
 
   const _FunctionSequentialContextBuilder(this._builder);
diff --git a/packages/kai_engine/lib/src/src.dart b/packages/kai_engine/lib/src/src.dart
index 72c857f..71a358e 100644
--- a/packages/kai_engine/lib/src/src.dart
+++ b/packages/kai_engine/lib/src/src.dart
@@ -8,9 +8,9 @@ export 'generation_service_base.dart';
 export 'kai_logger.dart';
 export 'message_adapter_base.dart';
 export 'message_repository_base.dart';
-export 'post_response_engine_base.dart';
-export 'prompt_engine.dart';
 export 'pipeline/kai_phase.dart';
 export 'pipeline/phase_types.dart';
+export 'post_response_engine_base.dart';
+export 'prompt_engine.dart';
 export 'query_engine_base.dart';
 export 'tool_schema.dart';
diff --git a/packages/kai_engine/test/src/chat_controller_base_test.dart b/packages/kai_engine/test/src/chat_controller_base_test.dart
index f73c28a..c824742 100644
--- a/packages/kai_engine/test/src/chat_controller_base_test.dart
+++ b/packages/kai_engine/test/src/chat_controller_base_test.dart
@@ -1,9 +1,9 @@
 import 'dart:async';
 
 import 'package:fast_immutable_collections/fast_immutable_collections.dart';
-import 'package:test/test.dart';
 import 'package:kai_engine/kai_engine.dart';
 import 'package:mocktail/mocktail.dart';
+import 'package:test/test.dart';
 
 GenerationResult createTestGenerationResult({
   required IList<CoreMessage> messages,
@@ -18,11 +18,9 @@ GenerationResult createTestGenerationResult({
   );
 }
 
-final class MockConversationManager extends Mock
-    implements ConversationManager<TestMessage> {
+final class MockConversationManager extends Mock implements ConversationManager<TestMessage> {
   @override
-  ConversationSession get session =>
-      const ConversationSession(id: 'test_session');
+  ConversationSession get session => const ConversationSession(id: 'test_session');
 }
 
 class MockGenerationService extends Mock implements GenerationServiceBase {}
@@ -137,10 +135,7 @@ final class TestContextEngine extends ContextEngine {
     final userMessage =
         providedUserMessage?.copyWith(content: inputQuery.processedQuery) ??
         CoreMessage.user(content: inputQuery.processedQuery);
-    return Future.value((
-      userMessage: userMessage,
-      prompts: IList([userMessage]),
-    ));
+    return Future.value((userMessage: userMessage, prompts: IList([userMessage])));
   }
 }
 
@@ -180,10 +175,7 @@ final class TestChatController extends ChatControllerBase<TestMessage> {
 
   @override
   GenerationExecuteConfig generativeConfigs(IList<CoreMessage> prompts) {
-    return const GenerationExecuteConfig(
-      tools: [],
-      config: {'temperature': 0.7},
-    );
+    return const GenerationExecuteConfig(tools: [], config: {'temperature': 0.7});
   }
 }
 
@@ -201,9 +193,7 @@ void main() {
     registerFallbackValue(CancelToken());
     registerFallbackValue(CoreMessage.user(content: 'test'));
     registerFallbackValue(
-      createTestGenerationResult(
-        messages: IList([CoreMessage.ai(content: 'test')]),
-      ),
+      createTestGenerationResult(messages: IList([CoreMessage.ai(content: 'test')])),
     );
     registerFallbackValue(<ToolSchema>[]);
     registerFallbackValue(<String, dynamic>{});
@@ -235,9 +225,7 @@ void main() {
         ),
         initialRequestMessageId: 'test',
         requestMessages: const IList.empty(),
-        result: createTestGenerationResult(
-          messages: IList([CoreMessage.ai(content: 'test')]),
-        ),
+        result: createTestGenerationResult(messages: IList([CoreMessage.ai(content: 'test')])),
         conversationManager: MockConversationManager(),
       ),
     );
@@ -303,11 +291,9 @@ void main() {
                 CoreMessage? providedUserMessage,
               }) async => (
                 userMessage:
-                    providedUserMessage ??
-                    CoreMessage.user(messageId: 'user-1', content: 'Hello'),
+                    providedUserMessage ?? CoreMessage.user(messageId: 'user-1', content: 'Hello'),
                 prompts: IList([
-                  providedUserMessage ??
-                      CoreMessage.user(messageId: 'user-1', content: 'Hello'),
+                  providedUserMessage ?? CoreMessage.user(messageId: 'user-1', content: 'Hello'),
                 ]),
               ),
         );
@@ -321,12 +307,9 @@ void main() {
             config: any(named: 'config'),
           ),
         ).thenAnswer((_) {
-          final streamController =
-              StreamController<GenerationState<GenerationResult>>();
+          final streamController = StreamController<GenerationState<GenerationResult>>();
           final result = createTestGenerationResult(
-            messages: IList([
-              CoreMessage.ai(messageId: 'ai-1', content: 'Hi there'),
-            ]),
+            messages: IList([CoreMessage.ai(messageId: 'ai-1', content: 'Hi there')]),
           );
 
           Future.microtask(() {
@@ -337,17 +320,12 @@ void main() {
           return streamController.stream;
         });
 
-        when(() => mockConversationManager.addMessages(any())).thenAnswer((
-          invocation,
-        ) async {
-          final messages =
-              invocation.positionalArguments[0] as IList<CoreMessage>;
+        when(() => mockConversationManager.addMessages(any())).thenAnswer((invocation) async {
+          final messages = invocation.positionalArguments[0] as IList<CoreMessage>;
           return messages;
         });
 
-        when(
-          () => mockConversationManager.updateMessages(any()),
-        ).thenAnswer((_) async {});
+        when(() => mockConversationManager.updateMessages(any())).thenAnswer((_) async {});
 
         controller = TestChatController(
           conversationManager: mockConversationManager,
@@ -381,14 +359,9 @@ void main() {
         fakeQueryEngine.reset();
         fakeQueryEngine.throwOnExecute = Exception('Query processing failed');
 
-        when(
-          () => mockConversationManager.removeMessages(any()),
-        ).thenAnswer((_) async {});
+        when(() => mockConversationManager.removeMessages(any())).thenAnswer((_) async {});
 
-        final result = await controller.submit(
-          'Hello',
-          revertInputOnError: true,
-        );
+        final result = await controller.submit('Hello', revertInputOnError: true);
 
         expect(result, isA<GenerationErrorState<GenerationResult>>());
 
@@ -423,14 +396,9 @@ void main() {
           testContextEngine: failingContextEngine,
         );
 
-        when(
-          () => mockConversationManager.removeMessages(any()),
-        ).thenAnswer((_) async {});
+        when(() => mockConversationManager.removeMessages(any())).thenAnswer((_) async {});
 
-        final result = await failingController.submit(
-          'Hello',
-          revertInputOnError: true,
-        );
+        final result = await failingController.submit('Hello', revertInputOnError: true);
 
         expect(result, isA<GenerationErrorState<GenerationResult>>());
 
@@ -447,8 +415,7 @@ void main() {
             config: any(named: 'config'),
           ),
         ).thenAnswer((_) {
-          final streamController =
-              StreamController<GenerationState<GenerationResult>>();
+          final streamController = StreamController<GenerationState<GenerationResult>>();
 
           Future.microtask(() {
             streamController.addError(Exception('Generation failed'));
@@ -457,14 +424,9 @@ void main() {
           return streamController.stream;
         });
 
-        when(
-          () => mockConversationManager.removeMessages(any()),
-        ).thenAnswer((_) async {});
+        when(() => mockConversationManager.removeMessages(any())).thenAnswer((_) async {});
 
-        final result = await controller.submit(
-          'Hello',
-          revertInputOnError: true,
-        );
+        final result = await controller.submit('Hello', revertInputOnError: true);
 
         expect(result, isA<GenerationErrorState<GenerationResult>>());
 
@@ -481,8 +443,7 @@ void main() {
             config: any(named: 'config'),
           ),
         ).thenAnswer((_) {
-          final streamController =
-              StreamController<GenerationState<GenerationResult>>();
+          final streamController = StreamController<GenerationState<GenerationResult>>();
 
           Future.microtask(() {
             streamController.close();
@@ -491,14 +452,9 @@ void main() {
           return streamController.stream;
         });
 
-        when(
-          () => mockConversationManager.removeMessages(any()),
-        ).thenAnswer((_) async {});
+        when(() => mockConversationManager.removeMessages(any())).thenAnswer((_) async {});
 
-        final result = await controller.submit(
-          'Hello',
-          revertInputOnError: true,
-        );
+        final result = await controller.submit('Hello', revertInputOnError: true);
 
         expect(result, isA<GenerationErrorState<GenerationResult>>());
 
@@ -506,97 +462,70 @@ void main() {
       });
 
       test('should handle post response engine failure gracefully', () async {
-        fakePostResponseEngine.throwOnExecute = Exception(
-          'Post response failed',
-        );
+        fakePostResponseEngine.throwOnExecute = Exception('Post response failed');
 
         final result = await controller.submit('Hello');
 
         expect(result, isA<GenerationErrorState<GenerationResult>>());
       });
 
-      test(
-        'should revert input on error when revertInputOnError is true',
-        () async {
-          fakeQueryEngine.reset();
-          fakeQueryEngine.throwOnExecute = Exception('Query failed');
+      test('should revert input on error when revertInputOnError is true', () async {
+        fakeQueryEngine.reset();
+        fakeQueryEngine.throwOnExecute = Exception('Query failed');
 
-          when(
-            () => mockConversationManager.removeMessages(any()),
-          ).thenAnswer((_) async {});
+        when(() => mockConversationManager.removeMessages(any())).thenAnswer((_) async {});
 
-          final result = await controller.submit(
-            'Hello',
-            revertInputOnError: true,
-          );
+        final result = await controller.submit('Hello', revertInputOnError: true);
 
-          expect(result, isA<GenerationErrorState<GenerationResult>>());
+        expect(result, isA<GenerationErrorState<GenerationResult>>());
 
-          verify(() => mockConversationManager.removeMessages(any())).called(1);
-        },
-      );
+        verify(() => mockConversationManager.removeMessages(any())).called(1);
+      });
 
-      test(
-        'should revert persisted message on error when revertInputOnError is true',
-        () async {
-          when(
-            () => mockGenerationService.stream(
-              any(),
-              cancelToken: any(named: 'cancelToken'),
-              tools: any(named: 'tools'),
-              toolingConfig: any(named: 'toolingConfig'),
-              config: any(named: 'config'),
-            ),
-          ).thenThrow(Exception('Generation failed'));
+      test('should revert persisted message on error when revertInputOnError is true', () async {
+        when(
+          () => mockGenerationService.stream(
+            any(),
+            cancelToken: any(named: 'cancelToken'),
+            tools: any(named: 'tools'),
+            toolingConfig: any(named: 'toolingConfig'),
+            config: any(named: 'config'),
+          ),
+        ).thenThrow(Exception('Generation failed'));
 
-          when(
-            () => mockConversationManager.removeMessages(any()),
-          ).thenAnswer((_) async {});
+        when(() => mockConversationManager.removeMessages(any())).thenAnswer((_) async {});
 
-          final result = await controller.submit(
-            'Hello',
-            revertInputOnError: true,
-          );
+        final result = await controller.submit('Hello', revertInputOnError: true);
 
-          expect(result, isA<GenerationErrorState<GenerationResult>>());
+        expect(result, isA<GenerationErrorState<GenerationResult>>());
 
-          verify(() => mockConversationManager.removeMessages(any())).called(1);
-        },
-      );
+        verify(() => mockConversationManager.removeMessages(any())).called(1);
+      });
 
-      test(
-        'should not revert input on error when revertInputOnError is false',
-        () async {
-          fakeQueryEngine.reset();
-          fakeQueryEngine.throwOnExecute = Exception('Query failed');
+      test('should not revert input on error when revertInputOnError is false', () async {
+        fakeQueryEngine.reset();
+        fakeQueryEngine.throwOnExecute = Exception('Query failed');
 
-          final result = await controller.submit(
-            'Hello',
-            revertInputOnError: false,
-          );
+        final result = await controller.submit('Hello', revertInputOnError: false);
 
-          expect(result, isA<GenerationErrorState<GenerationResult>>());
+        expect(result, isA<GenerationErrorState<GenerationResult>>());
 
-          verifyNever(() => mockConversationManager.removeMessages(any()));
-        },
-      );
+        verifyNever(() => mockConversationManager.removeMessages(any()));
+      });
 
-      test(
-        'should pass custom generative configs to generation service',
-        () async {
-          await controller.submit('Hello');
+      test('should pass custom generative configs to generation service', () async {
+        await controller.submit('Hello');
 
-          verify(
-            () => mockGenerationService.stream(
-              any(),
-              cancelToken: any(named: 'cancelToken'),
-              tools: [],
-              toolingConfig: any(named: 'toolingConfig'),
-              config: {'temperature': 0.7},
-            ),
-          ).called(1);
-        },
-      );
+        verify(
+          () => mockGenerationService.stream(
+            any(),
+            cancelToken: any(named: 'cancelToken'),
+            tools: [],
+            toolingConfig: any(named: 'toolingConfig'),
+            config: {'temperature': 0.7},
+          ),
+        ).called(1);
+      });
 
       test(
         'should handle stream error states and revert input when revertInputOnError is true',
@@ -610,14 +539,11 @@ void main() {
               config: any(named: 'config'),
             ),
           ).thenAnswer((_) {
-            final streamController =
-                StreamController<GenerationState<GenerationResult>>();
+            final streamController = StreamController<GenerationState<GenerationResult>>();
 
             Future.microtask(() {
               streamController.add(
-                GenerationState.error(
-                  KaiException.exception('Stream generation failed'),
-                ),
+                GenerationState.error(KaiException.exception('Stream generation failed')),
               );
               streamController.close();
             });
@@ -625,14 +551,9 @@ void main() {
             return streamController.stream;
           });
 
-          when(
-            () => mockConversationManager.removeMessages(any()),
-          ).thenAnswer((_) async {});
+          when(() => mockConversationManager.removeMessages(any())).thenAnswer((_) async {});
 
-          final result = await controller.submit(
-            'Hello',
-            revertInputOnError: true,
-          );
+          final result = await controller.submit('Hello', revertInputOnError: true);
 
           expect(result, isA<GenerationErrorState<GenerationResult>>());
 
@@ -652,14 +573,11 @@ void main() {
               config: any(named: 'config'),
             ),
           ).thenAnswer((_) {
-            final streamController =
-                StreamController<GenerationState<GenerationResult>>();
+            final streamController = StreamController<GenerationState<GenerationResult>>();
 
             Future.microtask(() {
               streamController.add(
-                GenerationState.error(
-                  KaiException.exception('Stream generation failed'),
-                ),
+                GenerationState.error(KaiException.exception('Stream generation failed')),
               );
               streamController.close();
             });
@@ -667,10 +585,7 @@ void main() {
             return streamController.stream;
           });
 
-          final result = await controller.submit(
-            'Hello',
-            revertInputOnError: false,
-          );
+          final result = await controller.submit('Hello', revertInputOnError: false);
 
           expect(result, isA<GenerationErrorState<GenerationResult>>());
 
@@ -680,9 +595,7 @@ void main() {
 
       test('should emit states through stream during submission', () async {
         final states = <GenerationState<GenerationResult>>[];
-        final subscription = controller.generationStateStream.listen(
-          states.add,
-        );
+        final subscription = controller.generationStateStream.listen(states.add);
 
         await controller.submit('Hello');
 
@@ -693,112 +606,96 @@ void main() {
         await subscription.cancel();
       });
 
-      test(
-        'should not have duplicate user input in prompt sent to stream',
-        () async {
-          IList<CoreMessage>? capturedPrompts;
-
-          final localMockConversationManager = MockConversationManager();
-          final localMockGenerationService = MockGenerationService();
-          final localFakeQueryEngine = FakeQueryEngine();
-          final localFakePostResponseEngine = FakePostResponseEngine();
-
-          when(() => localMockConversationManager.getMessages()).thenAnswer(
-            (_) async => IList([
-              CoreMessage.user(content: 'Previous message'),
-              CoreMessage.ai(content: 'Previous response'),
-            ]),
-          );
+      test('should not have duplicate user input in prompt sent to stream', () async {
+        IList<CoreMessage>? capturedPrompts;
 
-          when(
-            () => localMockConversationManager.addMessages(any()),
-          ).thenAnswer((invocation) async {
-            final messages =
-                invocation.positionalArguments[0] as IList<CoreMessage>;
-            return messages;
-          });
+        final localMockConversationManager = MockConversationManager();
+        final localMockGenerationService = MockGenerationService();
+        final localFakeQueryEngine = FakeQueryEngine();
+        final localFakePostResponseEngine = FakePostResponseEngine();
 
-          when(
-            () => localMockConversationManager.updateMessages(any()),
-          ).thenAnswer((_) async {});
+        when(() => localMockConversationManager.getMessages()).thenAnswer(
+          (_) async => IList([
+            CoreMessage.user(content: 'Previous message'),
+            CoreMessage.ai(content: 'Previous response'),
+          ]),
+        );
 
-          localFakeQueryEngine.onExecute = (input) async => QueryContext(
-            originalQuery: input.rawInput,
-            processedQuery: input.rawInput,
-            session: const ConversationSession(id: 'test_session'),
-          );
+        when(() => localMockConversationManager.addMessages(any())).thenAnswer((invocation) async {
+          final messages = invocation.positionalArguments[0] as IList<CoreMessage>;
+          return messages;
+        });
 
-          when(
-            () => localMockGenerationService.stream(
-              captureAny(that: isA<IList<CoreMessage>>()),
-              cancelToken: any(named: 'cancelToken'),
-              tools: any(named: 'tools'),
-              toolingConfig: any(named: 'toolingConfig'),
-              config: any(named: 'config'),
-            ),
-          ).thenAnswer((invocation) {
-            capturedPrompts =
-                invocation.positionalArguments[0] as IList<CoreMessage>;
-
-            final streamController =
-                StreamController<GenerationState<GenerationResult>>();
-            final result = createTestGenerationResult(
-              messages: IList([
-                CoreMessage.ai(messageId: 'ai-1', content: 'Hi there'),
-              ]),
-            );
+        when(() => localMockConversationManager.updateMessages(any())).thenAnswer((_) async {});
 
-            Future.microtask(() {
-              streamController.add(GenerationState.complete(result));
-              streamController.close();
-            });
+        localFakeQueryEngine.onExecute = (input) async => QueryContext(
+          originalQuery: input.rawInput,
+          processedQuery: input.rawInput,
+          session: const ConversationSession(id: 'test_session'),
+        );
 
-            return streamController.stream;
-          });
+        when(
+          () => localMockGenerationService.stream(
+            captureAny(that: isA<IList<CoreMessage>>()),
+            cancelToken: any(named: 'cancelToken'),
+            tools: any(named: 'tools'),
+            toolingConfig: any(named: 'toolingConfig'),
+            config: any(named: 'config'),
+          ),
+        ).thenAnswer((invocation) {
+          capturedPrompts = invocation.positionalArguments[0] as IList<CoreMessage>;
 
-          final localTestContextEngine = TestContextEngine(
-            generateFunction:
-                ({
-                  required IList<CoreMessage> source,
-                  required QueryContext inputQuery,
-                  CoreMessage? providedUserMessage,
-                }) async => (
-                  userMessage:
-                      providedUserMessage ??
-                      CoreMessage.user(messageId: 'user-1', content: 'Hello'),
-                  prompts: IList([
-                    CoreMessage.user(content: 'Previous message'),
-                    CoreMessage.ai(content: 'Previous response'),
-                    providedUserMessage ??
-                        CoreMessage.user(messageId: 'user-1', content: 'Hello'),
-                  ]),
-                ),
+          final streamController = StreamController<GenerationState<GenerationResult>>();
+          final result = createTestGenerationResult(
+            messages: IList([CoreMessage.ai(messageId: 'ai-1', content: 'Hi there')]),
           );
 
-          final localController = TestChatController(
-            conversationManager: localMockConversationManager,
-            generationService: localMockGenerationService,
-            queryEngine: localFakeQueryEngine,
-            postResponseEngine: localFakePostResponseEngine,
-            testContextEngine: localTestContextEngine,
-          );
+          Future.microtask(() {
+            streamController.add(GenerationState.complete(result));
+            streamController.close();
+          });
+
+          return streamController.stream;
+        });
 
-          await localController.submit('Hello');
+        final localTestContextEngine = TestContextEngine(
+          generateFunction:
+              ({
+                required IList<CoreMessage> source,
+                required QueryContext inputQuery,
+                CoreMessage? providedUserMessage,
+              }) async => (
+                userMessage:
+                    providedUserMessage ?? CoreMessage.user(messageId: 'user-1', content: 'Hello'),
+                prompts: IList([
+                  CoreMessage.user(content: 'Previous message'),
+                  CoreMessage.ai(content: 'Previous response'),
+                  providedUserMessage ?? CoreMessage.user(messageId: 'user-1', content: 'Hello'),
+                ]),
+              ),
+        );
 
-          expect(capturedPrompts, isNotNull);
-          expect(capturedPrompts!.length, equals(3));
+        final localController = TestChatController(
+          conversationManager: localMockConversationManager,
+          generationService: localMockGenerationService,
+          queryEngine: localFakeQueryEngine,
+          postResponseEngine: localFakePostResponseEngine,
+          testContextEngine: localTestContextEngine,
+        );
 
-          expect(capturedPrompts![0].content, equals('Previous message'));
-          expect(capturedPrompts![1].content, equals('Previous response'));
-          expect(capturedPrompts![2].content, equals('Hello'));
+        await localController.submit('Hello');
 
-          final helloMessages = capturedPrompts!
-              .where((msg) => msg.content == 'Hello')
-              .toList();
-          expect(helloMessages.length, equals(1));
-          expect(helloMessages[0].messageId, isNotEmpty);
-        },
-      );
+        expect(capturedPrompts, isNotNull);
+        expect(capturedPrompts!.length, equals(3));
+
+        expect(capturedPrompts![0].content, equals('Previous message'));
+        expect(capturedPrompts![1].content, equals('Previous response'));
+        expect(capturedPrompts![2].content, equals('Hello'));
+
+        final helloMessages = capturedPrompts!.where((msg) => msg.content == 'Hello').toList();
+        expect(helloMessages.length, equals(1));
+        expect(helloMessages[0].messageId, isNotEmpty);
+      });
     });
 
     group('streams', () {
@@ -813,4 +710,292 @@ void main() {
       });
     });
   });
+
+  group('KaiChatController', () {
+    late MockGenerationService mockGenerationService;
+
+    setUp(() {
+      mockGenerationService = MockGenerationService();
+
+      when(
+        () => mockGenerationService.stream(
+          any(),
+          cancelToken: any(named: 'cancelToken'),
+          tools: any(named: 'tools'),
+          toolingConfig: any(named: 'toolingConfig'),
+          config: any(named: 'config'),
+        ),
+      ).thenAnswer((_) {
+        final streamController = StreamController<GenerationState<GenerationResult>>();
+        final result = createTestGenerationResult(
+          messages: IList([CoreMessage.ai(messageId: 'ai-1', content: 'Response')]),
+        );
+
+        Future.microtask(() {
+          streamController.add(GenerationState.complete(result));
+          streamController.close();
+        });
+
+        return streamController.stream;
+      });
+    });
+
+    group('create', () {
+      test(
+        'should throw AssertionError when both contextEngine and systemPrompt are provided',
+        () async {
+          expect(
+            () => KaiChatController.create(
+              generationService: mockGenerationService,
+              contextEngine: TestContextEngine(),
+              systemPrompt: 'You are helpful',
+            ),
+            throwsA(isA<AssertionError>()),
+          );
+        },
+      );
+
+      test(
+        'should create controller with default values when no optional parameters provided',
+        () async {
+          final controller = await KaiChatController.create(
+            generationService: mockGenerationService,
+          );
+
+          expect(controller, isNotNull);
+          expect(controller, isA<KaiChatController>());
+
+          controller.dispose();
+        },
+      );
+
+      test('should create controller with custom sessionId', () async {
+        final controller = await KaiChatController.create(
+          sessionId: 'custom-session-123',
+          generationService: mockGenerationService,
+        );
+
+        expect(controller, isNotNull);
+
+        final messages = await controller.getAllMessages();
+        expect(messages, isNotNull);
+
+        controller.dispose();
+      });
+
+      test('should create controller with systemPrompt', () async {
+        final controller = await KaiChatController.create(
+          generationService: mockGenerationService,
+          systemPrompt: 'You are a helpful assistant.',
+        );
+
+        expect(controller, isNotNull);
+
+        final result = await controller.submit('Hello');
+        expect(result, isA<GenerationCompleteState<GenerationResult>>());
+
+        controller.dispose();
+      });
+
+      test('should create controller with custom contextEngine', () async {
+        final customContextEngine = TestContextEngine();
+
+        final controller = await KaiChatController.create(
+          generationService: mockGenerationService,
+          contextEngine: customContextEngine,
+        );
+
+        expect(controller, isNotNull);
+
+        final result = await controller.submit('Hello');
+        expect(result, isA<GenerationCompleteState<GenerationResult>>());
+
+        controller.dispose();
+      });
+
+      test('should create controller with queryEngine', () async {
+        final queryEngine = FakeQueryEngine();
+
+        final controller = await KaiChatController.create(
+          generationService: mockGenerationService,
+          queryEngine: queryEngine,
+        );
+
+        expect(controller, isNotNull);
+        expect(queryEngine.executeCalled, isFalse);
+
+        final result = await controller.submit('Hello');
+        expect(result, isA<GenerationCompleteState<GenerationResult>>());
+
+        expect(queryEngine.executeCalled, isTrue);
+
+        controller.dispose();
+      });
+
+      test('should create controller with postEngine', () async {
+        final postEngine = FakePostResponseEngine();
+
+        final controller = await KaiChatController.create(
+          generationService: mockGenerationService,
+          postEngine: postEngine,
+        );
+
+        expect(controller, isNotNull);
+        expect(postEngine.executeCalled, isFalse);
+
+        final result = await controller.submit('Hello');
+        expect(result, isA<GenerationCompleteState<GenerationResult>>());
+
+        expect(postEngine.executeCalled, isTrue);
+
+        controller.dispose();
+      });
+
+      test('should use InMemoryMessageRepository when repository not provided', () async {
+        final controller = await KaiChatController.create(generationService: mockGenerationService);
+
+        expect(controller, isNotNull);
+
+        final messages = await controller.getAllMessages();
+        expect(messages, isA<IList<CoreMessage>>());
+
+        controller.dispose();
+      });
+
+      test('should create controller with all parameters', () async {
+        final queryEngine = FakeQueryEngine();
+        final postEngine = FakePostResponseEngine();
+        final customContextEngine = TestContextEngine();
+
+        final controller = await KaiChatController.create(
+          sessionId: 'full-test-session',
+          generationService: mockGenerationService,
+          contextEngine: customContextEngine,
+          queryEngine: queryEngine,
+          postEngine: postEngine,
+        );
+
+        expect(controller, isNotNull);
+
+        final result = await controller.submit('Hello');
+        expect(result, isA<GenerationCompleteState<GenerationResult>>());
+
+        expect(queryEngine.executeCalled, isTrue);
+        expect(postEngine.executeCalled, isTrue);
+
+        controller.dispose();
+      });
+
+      test('should handle empty systemPrompt', () async {
+        final controller = await KaiChatController.create(
+          generationService: mockGenerationService,
+          systemPrompt: '',
+        );
+
+        expect(controller, isNotNull);
+
+        final result = await controller.submit('Hello');
+        expect(result, isA<GenerationCompleteState<GenerationResult>>());
+
+        controller.dispose();
+      });
+
+      test('should create multiple controllers with different sessionIds', () async {
+        final controller1 = await KaiChatController.create(
+          sessionId: 'session-1',
+          generationService: mockGenerationService,
+        );
+
+        final controller2 = await KaiChatController.create(
+          sessionId: 'session-2',
+          generationService: mockGenerationService,
+        );
+
+        expect(controller1, isNotNull);
+        expect(controller2, isNotNull);
+        expect(controller1, isNot(equals(controller2)));
+
+        controller1.dispose();
+        controller2.dispose();
+      });
+    });
+
+    group('submit', () {
+      test('should handle successful submission with default configuration', () async {
+        final controller = await KaiChatController.create(
+          generationService: mockGenerationService,
+          systemPrompt: 'You are helpful',
+        );
+
+        final result = await controller.submit('Hello, how are you?');
+
+        expect(result, isA<GenerationCompleteState<GenerationResult>>());
+
+        final messages = await controller.getAllMessages();
+        expect(messages.length, greaterThan(0));
+
+        controller.dispose();
+      });
+
+      test('should emit loading and complete states', () async {
+        final controller = await KaiChatController.create(generationService: mockGenerationService);
+
+        final states = <GenerationState<GenerationResult>>[];
+        final subscription = controller.generationStateStream.listen(states.add);
+
+        await controller.submit('Test message');
+
+        expect(states.length, greaterThan(1));
+        expect(states.any((s) => s is GenerationLoadingState), isTrue);
+        expect(states.any((s) => s is GenerationCompleteState), isTrue);
+
+        await subscription.cancel();
+        controller.dispose();
+      });
+
+      test('should handle error and revert input when revertInputOnError is true', () async {
+        when(
+          () => mockGenerationService.stream(
+            any(),
+            cancelToken: any(named: 'cancelToken'),
+            tools: any(named: 'tools'),
+            toolingConfig: any(named: 'toolingConfig'),
+            config: any(named: 'config'),
+          ),
+        ).thenAnswer((_) {
+          final streamController = StreamController<GenerationState<GenerationResult>>();
+
+          Future.microtask(() {
+            streamController.addError(Exception('Generation failed'));
+          });
+
+          return streamController.stream;
+        });
+
+        final controller = await KaiChatController.create(generationService: mockGenerationService);
+
+        final result = await controller.submit('Test message', revertInputOnError: true);
+
+        expect(result, isA<GenerationErrorState<GenerationResult>>());
+
+        controller.dispose();
+      });
+    });
+
+    group('messagesStream', () {
+      test('should emit messages through stream', () async {
+        final controller = await KaiChatController.create(generationService: mockGenerationService);
+
+        final messagesList = <IList<CoreMessage>>[];
+        final subscription = controller.messagesStream.listen(messagesList.add);
+
+        await controller.submit('Hello');
+
+        expect(messagesList.length, greaterThan(0));
+
+        await subscription.cancel();
+        controller.dispose();
+      });
+    });
+  });
 }
diff --git a/packages/kai_engine/test/src/message_repository_base_test.dart b/packages/kai_engine/test/src/message_repository_base_test.dart
index cfccd5d..ebf4967 100644
--- a/packages/kai_engine/test/src/message_repository_base_test.dart
+++ b/packages/kai_engine/test/src/message_repository_base_test.dart
@@ -42,10 +42,7 @@ void main() {
 
     group('constructor', () {
       test('should create instance with required callbacks', () {
-        final repo = CoreMessageRepository(
-          onInitial: (s) async => [],
-          onPut: (s, msgs) async => msgs,
-        );
+        final repo = CoreMessageRepository(onInitial: (s) async => [], onPut: (s, msgs) async {});
 
         expect(repo, isA<CoreMessageRepository>());
       });
@@ -53,7 +50,7 @@ void main() {
       test('should create instance with optional onRemove callback', () {
         final repo = CoreMessageRepository(
           onInitial: (s) async => [],
-          onPut: (s, msgs) async => msgs,
+          onPut: (s, msgs) async {},
           onRemove: (msgs) async {},
         );
 
@@ -69,7 +66,7 @@ void main() {
             capturedSession = s;
             return Future.value([testMessage1]);
           },
-          onPut: (s, msgs) async => msgs,
+          onPut: (s, msgs) async {},
         );
 
         final result = await repo.getMessages(session);
@@ -91,7 +88,7 @@ void main() {
             msg1.copyWith(timestamp: now),
             msg2.copyWith(timestamp: now.add(Duration(minutes: 1))),
           ],
-          onPut: (s, msgs) async => msgs,
+          onPut: (s, msgs) async {},
         );
 
         final result = await repo.getMessages(session);
@@ -105,7 +102,7 @@ void main() {
       test('should populate internal state from onInitial callback', () async {
         final repo = CoreMessageRepository(
           onInitial: (s) async => [testMessage1, testMessage2],
-          onPut: (s, msgs) async => msgs,
+          onPut: (s, msgs) async {},
         );
 
         await repo.getMessages(session);
@@ -121,7 +118,7 @@ void main() {
             callCount++;
             return callCount == 1 ? [testMessage1] : [testMessage2];
           },
-          onPut: (s, msgs) async => msgs,
+          onPut: (s, msgs) async {},
         );
 
         final result1 = await repo.getMessages(session);
@@ -135,7 +132,7 @@ void main() {
       test('should cache session for later use', () async {
         final repo = CoreMessageRepository(
           onInitial: (s) async => [testMessage1],
-          onPut: (s, msgs) async => msgs,
+          onPut: (s, msgs) async {},
         );
 
         await repo.getMessages(session);
@@ -167,10 +164,7 @@ void main() {
       });
 
       test('should add messages to internal state', () async {
-        final repo = CoreMessageRepository(
-          onInitial: (s) async => [],
-          onPut: (s, msgs) async => msgs,
-        );
+        final repo = CoreMessageRepository(onInitial: (s) async => [], onPut: (s, msgs) async {});
 
         await repo.saveMessages(session: session, messages: [testMessage1]);
         final result = await repo.getMessages(session);
@@ -179,10 +173,7 @@ void main() {
       });
 
       test('should accumulate messages in internal state across multiple saves', () async {
-        final repo = CoreMessageRepository(
-          onInitial: (s) async => [],
-          onPut: (s, msgs) async => msgs,
-        );
+        final repo = CoreMessageRepository(onInitial: (s) async => [], onPut: (s, msgs) async {});
 
         await repo.saveMessages(session: session, messages: [testMessage1]);
         await repo.saveMessages(session: session, messages: [testMessage2]);
@@ -192,10 +183,7 @@ void main() {
       });
 
       test('should cache session for later use', () async {
-        final repo = CoreMessageRepository(
-          onInitial: (s) async => [],
-          onPut: (s, msgs) async => msgs,
-        );
+        final repo = CoreMessageRepository(onInitial: (s) async => [], onPut: (s, msgs) async {});
 
         await repo.saveMessages(session: session, messages: [testMessage1]);
 
@@ -229,7 +217,7 @@ void main() {
       test('should update messages in internal state', () async {
         final repo = CoreMessageRepository(
           onInitial: (s) async => [testMessage1],
-          onPut: (s, msgs) async => msgs,
+          onPut: (s, msgs) async {},
         );
 
         await repo.getMessages(session);
@@ -241,10 +229,7 @@ void main() {
       });
 
       test('should throw StateError when session not initialized', () async {
-        final repo = CoreMessageRepository(
-          onInitial: (s) async => [],
-          onPut: (s, msgs) async => msgs,
-        );
+        final repo = CoreMessageRepository(onInitial: (s) async => [], onPut: (s, msgs) async {});
 
         expect(
           () => repo.updateMessages([testMessage1]),
@@ -259,10 +244,7 @@ void main() {
       });
 
       test('should use session from saveMessages if getMessages not called', () async {
-        final repo = CoreMessageRepository(
-          onInitial: (s) async => [],
-          onPut: (s, msgs) async => msgs,
-        );
+        final repo = CoreMessageRepository(onInitial: (s) async => [], onPut: (s, msgs) async {});
 
         await repo.saveMessages(session: session, messages: [testMessage1]);
         final updated = testMessage1.copyWith(content: 'Updated');
@@ -275,7 +257,7 @@ void main() {
       test('should only update existing messages, not add new ones', () async {
         final repo = CoreMessageRepository(
           onInitial: (s) async => [testMessage1],
-          onPut: (s, msgs) async => msgs,
+          onPut: (s, msgs) async {},
         );
 
         await repo.getMessages(session);
@@ -291,7 +273,7 @@ void main() {
         var capturedMessages;
         final repo = CoreMessageRepository(
           onInitial: (s) async => [testMessage1, testMessage2],
-          onPut: (s, msgs) async => msgs,
+          onPut: (s, msgs) async {},
           onRemove: (msgs) {
             capturedMessages = msgs;
             return Future.value();
@@ -307,7 +289,7 @@ void main() {
       test('should remove messages from internal state', () async {
         final repo = CoreMessageRepository(
           onInitial: (s) async => [testMessage1, testMessage2],
-          onPut: (s, msgs) async => msgs,
+          onPut: (s, msgs) async {},
         );
 
         await repo.getMessages(session);
@@ -320,7 +302,7 @@ void main() {
       test('should complete successfully when onRemove is null', () async {
         final repo = CoreMessageRepository(
           onInitial: (s) async => [testMessage1],
-          onPut: (s, msgs) async => msgs,
+          onPut: (s, msgs) async {},
         );
 
         await repo.getMessages(session);
@@ -330,7 +312,7 @@ void main() {
       test('should handle multiple messages in removeMessages', () async {
         final repo = CoreMessageRepository(
           onInitial: (s) async => [testMessage1, testMessage2],
-          onPut: (s, msgs) async => msgs,
+          onPut: (s, msgs) async {},
         );
 
         await repo.getMessages(session);
@@ -357,7 +339,6 @@ void main() {
                 persistedMessages.add(msg);
               }
             }
-            return msgs;
           },
           onRemove: (msgs) async {
             persistedMessages.removeWhere((m) => msgs.contains(m));
@@ -389,10 +370,7 @@ void main() {
       });
 
       test('should handle append-only pattern without onRemove', () async {
-        final repo = CoreMessageRepository(
-          onInitial: (s) async => [],
-          onPut: (s, msgs) async => msgs,
-        );
+        final repo = CoreMessageRepository(onInitial: (s) async => [], onPut: (s, msgs) async {});
 
         await repo.saveMessages(session: session, messages: [testMessage1]);
         await repo.saveMessages(session: session, messages: [testMessage2]);
@@ -415,7 +393,7 @@ void main() {
             persistenceCallCount++;
             return [];
           },
-          onPut: (s, msgs) async => msgs,
+          onPut: (s, msgs) async {},
         );
 
         // First load calls persistence

From b552cdd36c6aaa205da0b0051512ced2ba52e225 Mon Sep 17 00:00:00 2001
From: pckimlong <pochkimlong88@gmail.com>
Date: Sat, 10 Jan 2026 16:03:21 +0700
Subject: [PATCH 03/10] ``` refactor(chat_ui): migrate KaiChatView to
 controller-based architecture with enhanced customization

- Replace direct message/state props with ChatControllerBase for reactive state management
- Add StreamBuilder integration for messages and generation state streams
- Add onMessageSent and onError callbacks for better event handling
- Add messageItemBuilder, transientItemBuilder, and functionCallBuilder for granular customization
- Add showTimeLabel parameter to display message timestamps
- Rename
---
 .../lib/src/kai_chat_view.dart                | 159 ++++++++-----
 .../lib/src/widgets/kai_composer.dart         | 131 ++++++++++-
 .../lib/src/widgets/kai_message_bubble.dart   |  81 +++++--
 .../lib/src/widgets/kai_message_list.dart     | 211 +++++++++++++++---
 packages/kai_engine_chat_ui/pubspec.yaml      |   2 +
 5 files changed, 471 insertions(+), 113 deletions(-)

diff --git a/packages/kai_engine_chat_ui/lib/src/kai_chat_view.dart b/packages/kai_engine_chat_ui/lib/src/kai_chat_view.dart
index e9cb649..480dc23 100644
--- a/packages/kai_engine_chat_ui/lib/src/kai_chat_view.dart
+++ b/packages/kai_engine_chat_ui/lib/src/kai_chat_view.dart
@@ -1,3 +1,4 @@
+import 'package:fast_immutable_collections/fast_immutable_collections.dart';
 import 'package:flutter/material.dart';
 import 'package:kai_engine/kai_engine.dart';
 
@@ -7,43 +8,65 @@ import 'widgets/kai_message_list.dart';
 
 /// A batteries-included chat view: message list + composer.
 ///
-/// You control state by providing [messages],
-/// [generationState], and callbacks.
+/// Provide a [ChatControllerBase] to manage messages and generation state.
+/// The view listens to the controller's streams and rebuilds efficiently.
 class KaiChatView extends StatefulWidget {
   const KaiChatView({
     super.key,
-    required this.messages,
-    required this.onSend,
-    this.generationState,
-    this.onCancel,
+    required this.controller,
+    this.onMessageSent,
+    this.onError,
     this.messageFilter,
     this.messageBuilder,
+    this.messageItemBuilder,
     this.transientBuilder,
+    this.transientItemBuilder,
+    this.functionCallBuilder,
     this.onMessageTap,
     this.reverse = true,
-    this.controller,
+    this.scrollController,
     this.composerPadding = const EdgeInsets.fromLTRB(16, 0, 16, 12),
     this.listPadding,
     this.composerBuilder,
+    this.showTimeLabel = false,
   });
 
-  final List<CoreMessage> messages;
-  final GenerationState<GenerationResult>? generationState;
-  final Future<void> Function(String text) onSend;
-  final VoidCallback? onCancel;
+  final ChatControllerBase controller;
 
+  /// Called immediately when user sends a message (before awaiting result).
+  /// Useful for scrolling to bottom, clearing focus, etc.
+  final VoidCallback? onMessageSent;
+
+  /// Called when an error occurs during message submission.
+  /// Provides [BuildContext] for showing dialogs and [KaiException] for error details.
+  final void Function(BuildContext context, KaiException error)? onError;
+
+  /// Additional filter applied on top of the default filter.
+  /// Default filter already excludes non-displayable and background context messages.
+  /// Use this to add custom filtering rules without rewriting the defaults.
+  ///
+  /// Example - Filter out function call messages:
+  /// ```dart
+  /// KaiChatView(
+  ///   messageFilter: (message) => !message.content.contains('"functionCall":'),
+  ///   // ...
+  /// )
+  /// ```
   final KaiMessageFilter? messageFilter;
   final KaiMessageBuilder? messageBuilder;
+  final KaiMessageItemBuilder? messageItemBuilder;
   final KaiTransientBuilder? transientBuilder;
+  final KaiTransientItemBuilder? transientItemBuilder;
+  final KaiFunctionCallBuilder? functionCallBuilder;
   final void Function(CoreMessage message)? onMessageTap;
 
   final bool reverse;
-  final ScrollController? controller;
+  final ScrollController? scrollController;
   final EdgeInsetsGeometry composerPadding;
   final EdgeInsetsGeometry? listPadding;
 
-  final Widget Function(BuildContext context, KaiComposer composer)?
-  composerBuilder;
+  final Widget Function(BuildContext context, KaiComposer composer)? composerBuilder;
+  final bool showTimeLabel;
 
   @override
   State<KaiChatView> createState() => _KaiChatViewState();
@@ -56,8 +79,8 @@ class _KaiChatViewState extends State<KaiChatView> {
   @override
   void initState() {
     super.initState();
-    if (widget.controller != null) {
-      _scrollController = widget.controller!;
+    if (widget.scrollController != null) {
+      _scrollController = widget.scrollController!;
     } else {
       _scrollController = ScrollController();
       _ownsController = true;
@@ -70,11 +93,9 @@ class _KaiChatViewState extends State<KaiChatView> {
     super.dispose();
   }
 
-  void scrollToBottom() {
+  void _scrollToBottom() {
     if (!_scrollController.hasClients) return;
-    final target = widget.reverse
-        ? 0.0
-        : _scrollController.position.maxScrollExtent;
+    final target = widget.reverse ? 0.0 : _scrollController.position.maxScrollExtent;
     _scrollController.animateTo(
       target,
       duration: const Duration(milliseconds: 220),
@@ -82,41 +103,77 @@ class _KaiChatViewState extends State<KaiChatView> {
     );
   }
 
+  Future<void> _handleSend(String text) async {
+    widget.onMessageSent?.call();
+    _scrollToBottom();
+
+    final result = await widget.controller.submit(text);
+    if (result.isError && widget.onError != null && mounted) {
+      final error = (result as GenerationErrorState).exception;
+      widget.onError!(context, error);
+    }
+  }
+
+  void _handleCancel() {
+    widget.controller.cancel();
+  }
+
   @override
   Widget build(BuildContext context) {
     final theme = context.kaiChatTheme;
-    final composer = KaiComposer(
-      onSend: (text) async {
-        await widget.onSend(text);
-        scrollToBottom();
-      },
-      onCancel: widget.onCancel,
-      isGenerating: widget.generationState?.isGenerating ?? false,
-    );
 
-    return Column(
-      children: [
-        Expanded(
-          child: KaiMessageList(
-            controller: _scrollController,
-            messages: widget.messages,
-            generationState: widget.generationState,
-            reverse: widget.reverse,
-            messageFilter: widget.messageFilter ?? defaultKaiMessageFilter,
-            messageBuilder: widget.messageBuilder,
-            transientBuilder: widget.transientBuilder,
-            onMessageTap: widget.onMessageTap,
-            padding: widget.listPadding ?? theme.listPadding,
-          ),
-        ),
-        SafeArea(
-          top: false,
-          child: Padding(
-            padding: widget.composerPadding,
-            child: widget.composerBuilder?.call(context, composer) ?? composer,
-          ),
-        ),
-      ],
+    return StreamBuilder<GenerationState<GenerationResult>>(
+      stream: widget.controller.generationStateStream,
+      initialData: GenerationState.initial(),
+      builder: (context, stateSnapshot) {
+        final generationState = stateSnapshot.data ?? GenerationState.initial();
+        final isGenerating = generationState.isGenerating;
+
+        final composer = KaiComposer(
+          onSend: _handleSend,
+          onCancel: _handleCancel,
+          isGenerating: isGenerating,
+        );
+
+        return Column(
+          children: [
+            Expanded(
+              child: StreamBuilder<IList<CoreMessage>>(
+                stream: widget.controller.messagesStream,
+                initialData: const IListConst([]),
+                builder: (context, messagesSnapshot) {
+                  final messages = messagesSnapshot.data ?? const IListConst([]);
+                  final combinedFilter = widget.messageFilter != null
+                      ? (CoreMessage m) => defaultKaiMessageFilter(m) && widget.messageFilter!(m)
+                      : defaultKaiMessageFilter;
+                  return KaiMessageList(
+                    controller: _scrollController,
+                    messages: messages,
+                    generationState: generationState,
+                    reverse: widget.reverse,
+                    messageFilter: combinedFilter,
+                    messageBuilder: widget.messageBuilder,
+                    messageItemBuilder: widget.messageItemBuilder,
+                    transientBuilder: widget.transientBuilder,
+                    transientItemBuilder: widget.transientItemBuilder,
+                    functionCallBuilder: widget.functionCallBuilder,
+                    onMessageTap: widget.onMessageTap,
+                    padding: widget.listPadding ?? theme.listPadding,
+                    showTimeLabel: widget.showTimeLabel,
+                  );
+                },
+              ),
+            ),
+            SafeArea(
+              top: false,
+              child: Padding(
+                padding: widget.composerPadding,
+                child: widget.composerBuilder?.call(context, composer) ?? composer,
+              ),
+            ),
+          ],
+        );
+      },
     );
   }
 }
diff --git a/packages/kai_engine_chat_ui/lib/src/widgets/kai_composer.dart b/packages/kai_engine_chat_ui/lib/src/widgets/kai_composer.dart
index 3bfdb04..a485ba0 100644
--- a/packages/kai_engine_chat_ui/lib/src/widgets/kai_composer.dart
+++ b/packages/kai_engine_chat_ui/lib/src/widgets/kai_composer.dart
@@ -1,5 +1,105 @@
 import 'package:flutter/material.dart';
 
+/// State data provided to [KaiComposerBuilder].
+class KaiComposerState {
+  const KaiComposerState({
+    required this.isGenerating,
+    required this.isSending,
+    required this.canSend,
+    required this.text,
+  });
+
+  /// Whether the AI is currently generating a response.
+  final bool isGenerating;
+
+  /// Whether a send operation is in progress.
+  final bool isSending;
+
+  /// Whether the user can send a message (has text and not busy).
+  final bool canSend;
+
+  /// Current text in the input field.
+  final String text;
+}
+
+/// Callbacks provided to [KaiComposerBuilder].
+class KaiComposerCallbacks {
+  const KaiComposerCallbacks({required this.onSend, this.onCancel, required this.onTextChanged});
+
+  /// Send the current text in the input field.
+  final Future<void> Function() onSend;
+
+  /// Cancel the current generation (optional).
+  final VoidCallback? onCancel;
+
+  /// Called when the text field content changes.
+  final void Function(String text) onTextChanged;
+}
+
+/// Builder that creates a custom composer widget.
+///
+/// This builder receives the [state] and [callbacks] to build a custom composer.
+/// You can create your own text field, send button, and any additional UI elements.
+///
+/// Example - Create a custom composer with attachment button:
+/// ```dart
+/// KaiComposer(
+///   onSend: (text) => controller.submit(text),
+///   onCancel: () => controller.cancel(),
+///   isGenerating: isGenerating,
+///   builder: (context, state, callbacks) {
+///     return Row(
+///       children: [
+///         IconButton(
+///           icon: Icon(Icons.attach_file),
+///           onPressed: () => _handleAttachment(),
+///         ),
+///         Expanded(
+///           child: TextField(
+///             decoration: InputDecoration(
+///               hintText: 'Type a message...',
+///               border: OutlineInputBorder(),
+///             ),
+///             onChanged: callbacks.onTextChanged,
+///             onSubmitted: (_) => callbacks.onSend(),
+///           ),
+///         ),
+///         if (state.isGenerating)
+///           IconButton(
+///             icon: Icon(Icons.stop),
+///             onPressed: callbacks.onCancel,
+///           )
+///         else
+///           IconButton(
+///             icon: Icon(Icons.send),
+///             onPressed: state.canSend ? callbacks.onSend : null,
+///           ),
+///       ],
+///     );
+///   },
+/// )
+/// ```
+///
+/// Example - Use with KaiChatView:
+/// ```dart
+/// KaiChatView(
+///   controller: myController,
+///   composerBuilder: (context, composer) {
+///     return KaiComposer(
+///       onSend: composer.onSend,
+///       onCancel: composer.onCancel,
+///       isGenerating: composer.isGenerating,
+///       builder: (context, state, callbacks) {
+///         // Custom composer UI
+///         return YourCustomComposer(state: state, callbacks: callbacks);
+///       },
+///     );
+///   },
+/// )
+/// ```
+typedef KaiComposerBuilder =
+    Widget Function(BuildContext context, KaiComposerState state, KaiComposerCallbacks callbacks);
+
 class KaiComposer extends StatefulWidget {
   const KaiComposer({
     super.key,
@@ -11,6 +111,7 @@ class KaiComposer extends StatefulWidget {
     this.autofocus = false,
     this.maxLines = 6,
     this.onError,
+    this.builder,
   });
 
   final Future<void> Function(String text) onSend;
@@ -21,6 +122,7 @@ class KaiComposer extends StatefulWidget {
   final bool autofocus;
   final int maxLines;
   final void Function(Object error, StackTrace stackTrace)? onError;
+  final KaiComposerBuilder? builder;
 
   @override
   State<KaiComposer> createState() => _KaiComposerState();
@@ -53,10 +155,12 @@ class _KaiComposerState extends State<KaiComposer> {
     final text = _controller.text.trim();
     if (text.isEmpty) return;
 
+    // Clear text immediately for responsive UX (like Resonate pattern)
+    _controller.clear();
     setState(() => _sending = true);
+
     try {
       await widget.onSend(text);
-      if (mounted) _controller.clear();
     } catch (e, s) {
       widget.onError?.call(e, s);
     } finally {
@@ -66,10 +170,28 @@ class _KaiComposerState extends State<KaiComposer> {
 
   @override
   Widget build(BuildContext context) {
-    final colors = Theme.of(context).colorScheme;
     final isBusy = widget.isGenerating || _sending;
     final canSend = _controller.text.trim().isNotEmpty && !isBusy;
 
+    final state = KaiComposerState(
+      isGenerating: widget.isGenerating,
+      isSending: _sending,
+      canSend: canSend,
+      text: _controller.text,
+    );
+
+    final callbacks = KaiComposerCallbacks(
+      onSend: _handleSend,
+      onCancel: widget.onCancel,
+      onTextChanged: (text) => setState(() {}),
+    );
+
+    if (widget.builder != null) {
+      return widget.builder!(context, state, callbacks);
+    }
+
+    final colors = Theme.of(context).colorScheme;
+
     return Row(
       children: [
         Expanded(
@@ -87,10 +209,7 @@ class _KaiComposerState extends State<KaiComposer> {
                 borderRadius: BorderRadius.circular(16),
                 borderSide: BorderSide.none,
               ),
-              contentPadding: const EdgeInsets.symmetric(
-                horizontal: 14,
-                vertical: 12,
-              ),
+              contentPadding: const EdgeInsets.symmetric(horizontal: 14, vertical: 12),
             ),
             onSubmitted: (_) => _handleSend(),
             onChanged: (_) => setState(() {}),
diff --git a/packages/kai_engine_chat_ui/lib/src/widgets/kai_message_bubble.dart b/packages/kai_engine_chat_ui/lib/src/widgets/kai_message_bubble.dart
index 3423f02..2549b8e 100644
--- a/packages/kai_engine_chat_ui/lib/src/widgets/kai_message_bubble.dart
+++ b/packages/kai_engine_chat_ui/lib/src/widgets/kai_message_bubble.dart
@@ -1,10 +1,10 @@
 import 'package:flutter/material.dart';
+import 'package:gpt_markdown/gpt_markdown.dart';
 import 'package:kai_engine/kai_engine.dart';
 
 import '../kai_chat_theme.dart';
 
-typedef KaiMessageContentBuilder =
-    Widget Function(BuildContext context, CoreMessage message);
+typedef KaiMessageContentBuilder = Widget Function(BuildContext context, CoreMessage message);
 
 /// Default bubble used by [KaiMessageList] / [KaiChatView].
 class KaiMessageBubble extends StatelessWidget {
@@ -15,6 +15,7 @@ class KaiMessageBubble extends StatelessWidget {
     this.onTap,
     this.isUserMessage,
     this.semanticLabel,
+    this.showTimeLabel = false,
   });
 
   final CoreMessage message;
@@ -22,6 +23,7 @@ class KaiMessageBubble extends StatelessWidget {
   final VoidCallback? onTap;
   final bool? isUserMessage;
   final String? semanticLabel;
+  final bool showTimeLabel;
 
   @override
   Widget build(BuildContext context) {
@@ -36,19 +38,34 @@ class KaiMessageBubble extends StatelessWidget {
         ? (theme.userTextColor ?? colors.onPrimary)
         : (theme.aiTextColor ?? colors.onSurface);
 
+    final textStyle = Theme.of(
+      context,
+    ).textTheme.bodyMedium?.copyWith(color: textColor, height: 1.35);
+
     final content =
         contentBuilder?.call(context, message) ??
-        Text(
-          message.content,
-          style: Theme.of(
-            context,
-          ).textTheme.bodyMedium?.copyWith(color: textColor, height: 1.35),
-        );
+        (isUser
+            ? Text(message.content, style: textStyle)
+            : GptMarkdown(message.content, style: textStyle));
+
+    final child = Column(
+      mainAxisSize: MainAxisSize.min,
+      crossAxisAlignment: isUser ? CrossAxisAlignment.end : CrossAxisAlignment.start,
+      children: [
+        content,
+        if (showTimeLabel)
+          Padding(
+            padding: const EdgeInsets.only(top: 6),
+            child: Text(
+              _formatTime(message.timestamp),
+              style: TextStyle(fontSize: 12, color: textColor.withValues(alpha: 0.6)),
+            ),
+          ),
+      ],
+    );
 
     return Row(
-      mainAxisAlignment: isUser
-          ? MainAxisAlignment.end
-          : MainAxisAlignment.start,
+      mainAxisAlignment: isUser ? MainAxisAlignment.end : MainAxisAlignment.start,
       children: [
         Flexible(
           child: Semantics(
@@ -69,25 +86,22 @@ class KaiMessageBubble extends StatelessWidget {
                     padding: theme.bubblePadding,
                     decoration: BoxDecoration(
                       color: bubbleColor,
-                      borderRadius: BorderRadius.circular(theme.bubbleRadius)
-                          .copyWith(
-                            bottomLeft: isUser
-                                ? Radius.circular(theme.bubbleRadius)
-                                : const Radius.circular(6),
-                            bottomRight: isUser
-                                ? const Radius.circular(6)
-                                : Radius.circular(theme.bubbleRadius),
-                          ),
+                      borderRadius: BorderRadius.circular(theme.bubbleRadius).copyWith(
+                        bottomLeft: isUser
+                            ? Radius.circular(theme.bubbleRadius)
+                            : const Radius.circular(6),
+                        bottomRight: isUser
+                            ? const Radius.circular(6)
+                            : Radius.circular(theme.bubbleRadius),
+                      ),
                       border: Border.all(
-                        color: colors.outlineVariant.withValues(
-                          alpha: isUser ? 0.15 : 0.35,
-                        ),
+                        color: colors.outlineVariant.withValues(alpha: isUser ? 0.15 : 0.35),
                         width: 1,
                       ),
                     ),
                     child: DefaultTextStyle.merge(
                       style: TextStyle(color: textColor),
-                      child: content,
+                      child: child,
                     ),
                   ),
                 ),
@@ -98,4 +112,23 @@ class KaiMessageBubble extends StatelessWidget {
       ],
     );
   }
+
+  String _formatTime(DateTime timestamp) {
+    final now = DateTime.now();
+    final difference = now.difference(timestamp);
+
+    if (difference.inMinutes < 1) {
+      return 'Just now';
+    } else if (difference.inMinutes < 60) {
+      return '${difference.inMinutes}m ago';
+    } else if (difference.inHours < 24) {
+      return '${difference.inHours}h ago';
+    } else if (difference.inDays == 1) {
+      return 'Yesterday';
+    } else if (difference.inDays < 7) {
+      return '${difference.inDays}d ago';
+    } else {
+      return '${timestamp.day}/${timestamp.month}/${timestamp.year}';
+    }
+  }
 }
diff --git a/packages/kai_engine_chat_ui/lib/src/widgets/kai_message_list.dart b/packages/kai_engine_chat_ui/lib/src/widgets/kai_message_list.dart
index fb09c8b..a862bac 100644
--- a/packages/kai_engine_chat_ui/lib/src/widgets/kai_message_list.dart
+++ b/packages/kai_engine_chat_ui/lib/src/widgets/kai_message_list.dart
@@ -1,3 +1,4 @@
+import 'package:fast_immutable_collections/fast_immutable_collections.dart';
 import 'package:flutter/material.dart';
 import 'package:kai_engine/kai_engine.dart';
 
@@ -6,17 +7,147 @@ import 'kai_message_bubble.dart';
 import 'kai_streaming_text.dart';
 import 'kai_typing_indicator.dart';
 
-typedef KaiMessageBuilder =
-    Widget Function(BuildContext context, CoreMessage message);
+/// Builder that constructs a custom message widget.
+///
+/// This builder receives the [context] and [message] as parameters.
+/// You can use this to create a completely custom message layout.
+///
+/// Example - Create a simple text message:
+/// ```dart
+/// KaiChatView(
+///   messageBuilder: (context, message) {
+///     return Text(message.content);
+///   },
+///   // ...
+/// )
+/// ```
+typedef KaiMessageBuilder = Widget Function(BuildContext context, CoreMessage message);
+
+/// Builder that provides a default child widget for customization.
+///
+/// This builder receives the [defaultChild] which is the pre-built [KaiMessageBubble].
+/// You can wrap it with additional widgets (e.g., avatar, timestamp, status indicators)
+/// or use it as a reference to build your own custom layout.
+///
+/// Example - Add an avatar next to messages:
+/// ```dart
+/// KaiChatView(
+///   messageItemBuilder: (context, message, defaultChild) {
+///     return Row(
+///       children: [
+///         CircleAvatar(
+///           backgroundImage: NetworkImage(message.type == CoreMessageType.user
+///               ? 'https://api.dicebear.com/7.x/avataaars/svg?seed=user'
+///               : 'https://api.dicebear.com/7.x/bottts/svg?seed=ai'),
+///         ),
+///         SizedBox(width: 8),
+///         Expanded(child: defaultChild),
+///       ],
+///     );
+///   },
+///   // ...
+/// )
+/// ```
+///
+/// Example - Add a status indicator:
+/// ```dart
+/// KaiChatView(
+///   messageItemBuilder: (context, message, defaultChild) {
+///     return Column(
+///       crossAxisAlignment: CrossAxisAlignment.start,
+///       children: [
+///         defaultChild,
+///         if (message.metadata['status'] != null)
+///           Text(
+///             message.metadata['status'],
+///             style: TextStyle(fontSize: 12, color: Colors.grey),
+///           ),
+///       ],
+///     );
+///   },
+///   // ...
+/// )
+/// ```
+typedef KaiMessageItemBuilder =
+    Widget Function(BuildContext context, CoreMessage message, Widget defaultChild);
+
+/// Builder that constructs a custom transient state widget.
+///
+/// This builder receives the [context] and [state] as parameters.
+/// You can use this to create a completely custom transient state layout.
+///
+/// Example - Create a simple loading indicator:
+/// ```dart
+/// KaiChatView(
+///   transientBuilder: (context, state) {
+///     return CircularProgressIndicator();
+///   },
+///   // ...
+/// )
+/// ```
 typedef KaiTransientBuilder =
-    Widget Function(
-      BuildContext context,
-      GenerationState<GenerationResult> state,
-    );
+    Widget Function(BuildContext context, GenerationState<GenerationResult> state);
+
+/// Builder that provides a default child widget for transient state customization.
+///
+/// Similar to [KaiMessageItemBuilder], this allows you to customize the transient
+/// state (loading, streaming, function calling) while using the default bubble as a base.
+///
+/// Example - Add a loading animation:
+/// ```dart
+/// KaiChatView(
+///   transientItemBuilder: (context, state, defaultChild) {
+///     return Column(
+///       children: [
+///         defaultChild,
+///         LinearProgressIndicator(),
+///       ],
+///     );
+///   },
+///   // ...
+/// )
+/// ```
+typedef KaiTransientItemBuilder =
+    Widget Function(BuildContext context, GenerationState<GenerationResult> state, Widget child);
+
+/// Filter that determines whether a message should be displayed.
+///
+/// This filter receives the [message] as a parameter.
+/// You can use this to filter out certain types of messages.
+///
+/// Example - Filter out background context messages:
+/// ```dart
+/// KaiChatView(
+///   messageFilter: (message) => !message.isBackgroundContext,
+///   // ...
+/// )
+/// ```
 typedef KaiMessageFilter = bool Function(CoreMessage message);
 
-bool defaultKaiMessageFilter(CoreMessage m) =>
-    m.isDisplayable && !m.isBackgroundContext;
+/// Builder that constructs a custom function call widget.
+///
+/// This builder receives the [message] as a parameter.
+/// If provided, function messages will be displayed in the chat.
+/// If null, function messages are filtered out by default.
+///
+/// Example - Display function call with name and result:
+/// ```dart
+/// KaiChatView(
+///   functionCallBuilder: (context, message) {
+///     return Card(
+///       child: ListTile(
+///         leading: Icon(Icons.code),
+///         title: Text('Function: ${message.extensions["name"]}'),
+///         subtitle: Text(message.content),
+///       ),
+///     );
+///   },
+///   // ...
+/// )
+/// ```
+typedef KaiFunctionCallBuilder = Widget Function(BuildContext context, CoreMessage message);
+
+bool defaultKaiMessageFilter(CoreMessage m) => m.isDisplayable && !m.isBackgroundContext;
 
 class KaiMessageList extends StatelessWidget {
   const KaiMessageList({
@@ -27,27 +158,40 @@ class KaiMessageList extends StatelessWidget {
     this.reverse = true,
     this.messageFilter = defaultKaiMessageFilter,
     this.messageBuilder,
+    this.messageItemBuilder,
     this.transientBuilder,
+    this.transientItemBuilder,
+    this.functionCallBuilder,
     this.onMessageTap,
     this.padding,
     this.physics,
+    this.showTimeLabel = false,
   });
 
-  final List<CoreMessage> messages;
+  final IList<CoreMessage> messages;
   final GenerationState<GenerationResult>? generationState;
   final ScrollController? controller;
   final bool reverse;
   final KaiMessageFilter messageFilter;
   final KaiMessageBuilder? messageBuilder;
+  final KaiMessageItemBuilder? messageItemBuilder;
   final KaiTransientBuilder? transientBuilder;
+  final KaiTransientItemBuilder? transientItemBuilder;
+  final KaiFunctionCallBuilder? functionCallBuilder;
   final void Function(CoreMessage message)? onMessageTap;
   final EdgeInsetsGeometry? padding;
   final ScrollPhysics? physics;
+  final bool showTimeLabel;
 
   @override
   Widget build(BuildContext context) {
     final theme = context.kaiChatTheme;
-    final filtered = messages.where(messageFilter).toList(growable: false);
+    final effectiveFilter = functionCallBuilder != null
+        ? (CoreMessage m) =>
+              (m.isDisplayable && !m.isBackgroundContext) ||
+              (m.type == CoreMessageType.function && !m.isBackgroundContext)
+        : messageFilter;
+    final filtered = messages.where(effectiveFilter).toIList();
     final transient = generationState != null && generationState!.isGenerating;
     final totalCount = filtered.length + (transient ? 1 : 0);
     final transientIndex = reverse ? 0 : filtered.length;
@@ -62,9 +206,9 @@ class KaiMessageList extends StatelessWidget {
         final isTransientIndex = transient && index == transientIndex;
         if (isTransientIndex) {
           final state = generationState!;
-          final built =
-              transientBuilder?.call(context, state) ??
-              _defaultTransient(context, state);
+          final defaultChild =
+              transientBuilder?.call(context, state) ?? _defaultTransient(context, state);
+          final built = transientItemBuilder?.call(context, state, defaultChild) ?? defaultChild;
           return Padding(
             padding: EdgeInsets.only(bottom: theme.itemSpacing),
             child: built,
@@ -76,13 +220,24 @@ class KaiMessageList extends StatelessWidget {
             : index;
         final message = filtered[messageIndex];
 
+        if (message.type == CoreMessageType.function && functionCallBuilder != null) {
+          return Padding(
+            padding: EdgeInsets.only(bottom: theme.itemSpacing),
+            child: functionCallBuilder!(context, message),
+          );
+        }
+
+        final defaultChild = KaiMessageBubble(
+          message: message,
+          isUserMessage: message.type == CoreMessageType.user,
+          onTap: onMessageTap == null ? null : () => onMessageTap!(message),
+          showTimeLabel: showTimeLabel,
+        );
+
         final bubble =
+            messageItemBuilder?.call(context, message, defaultChild) ??
             messageBuilder?.call(context, message) ??
-            KaiMessageBubble(
-              message: message,
-              isUserMessage: message.type == CoreMessageType.user,
-              onTap: onMessageTap == null ? null : () => onMessageTap!(message),
-            );
+            defaultChild;
 
         return Padding(
           padding: EdgeInsets.only(bottom: theme.itemSpacing),
@@ -92,27 +247,19 @@ class KaiMessageList extends StatelessWidget {
     );
   }
 
-  Widget _defaultTransient(
-    BuildContext context,
-    GenerationState<GenerationResult> state,
-  ) {
+  Widget _defaultTransient(BuildContext context, GenerationState<GenerationResult> state) {
     return switch (state) {
-      GenerationLoadingState() ||
-      GenerationFunctionCallingState() => KaiMessageBubble(
+      GenerationLoadingState() || GenerationFunctionCallingState() => KaiMessageBubble(
         message: CoreMessage.ai(content: ''),
         isUserMessage: false,
-        contentBuilder: (_, __) => const Padding(
-          padding: EdgeInsets.symmetric(vertical: 2),
-          child: KaiTypingIndicator(),
-        ),
+        contentBuilder: (_, __) =>
+            const Padding(padding: EdgeInsets.symmetric(vertical: 2), child: KaiTypingIndicator()),
       ),
       GenerationStreamingTextState(text: final text) => KaiMessageBubble(
         message: CoreMessage.ai(content: text),
         isUserMessage: false,
-        contentBuilder: (context, _) => KaiStreamingText(
-          text: text,
-          style: Theme.of(context).textTheme.bodyMedium,
-        ),
+        contentBuilder: (context, _) =>
+            KaiStreamingText(text: text, style: Theme.of(context).textTheme.bodyMedium),
       ),
       _ => const SizedBox.shrink(),
     };
diff --git a/packages/kai_engine_chat_ui/pubspec.yaml b/packages/kai_engine_chat_ui/pubspec.yaml
index 1cddcc6..5515f23 100644
--- a/packages/kai_engine_chat_ui/pubspec.yaml
+++ b/packages/kai_engine_chat_ui/pubspec.yaml
@@ -13,6 +13,8 @@ environment:
 dependencies:
   flutter:
     sdk: flutter
+  fast_immutable_collections: ^11.0.2
+  gpt_markdown: ^1.1.2
   kai_engine: ^0.1.1
 
 dev_dependencies:

From 5eb3fc013c770364c23b218c9adf3e98ecdf6013 Mon Sep 17 00:00:00 2001
From: pckimlong <pochkimlong88@gmail.com>
Date: Sat, 10 Jan 2026 17:36:48 +0700
Subject: [PATCH 04/10] ``` refactor(adapters): improve code formatting and add
 timestamp preservation in FirebaseAiContentAdapter

- Format multi-line getters and factory methods to single lines in CoreMessage
- Format long class declarations and method calls to single lines
- Add timestamp preservation logic when converting Content to CoreMessage
- Extract contentJson variable for better readability
- Add timestamp parsing from stored content with fallback handling
- Store timestamp in extensions for future round-trip conversions
---
 .../lib/src/models/core_message.dart          |  9 ++--
 .../lib/src/adapters.dart                     | 48 ++++++++++---------
 2 files changed, 28 insertions(+), 29 deletions(-)

diff --git a/packages/kai_engine/lib/src/models/core_message.dart b/packages/kai_engine/lib/src/models/core_message.dart
index 38d641e..a9128d1 100644
--- a/packages/kai_engine/lib/src/models/core_message.dart
+++ b/packages/kai_engine/lib/src/models/core_message.dart
@@ -84,14 +84,11 @@ sealed class CoreMessage with _$CoreMessage {
     @Default(<String, dynamic>{}) Map<String, dynamic> extensions,
   }) = _CoreMessage;
 
-  bool get isDisplayable =>
-      type == CoreMessageType.user || type == CoreMessageType.ai;
-  bool get isMessageSavable =>
-      !isBackgroundContext && type != CoreMessageType.system;
+  bool get isDisplayable => type == CoreMessageType.user || type == CoreMessageType.ai;
+  bool get isMessageSavable => !isBackgroundContext && type != CoreMessageType.system;
 
   // This allow for simple use case where adapter is not require, user can directly use [CoreMessage] as main model
-  factory CoreMessage.fromJson(Map<String, dynamic> json) =>
-      _$CoreMessageFromJson(json);
+  factory CoreMessage.fromJson(Map<String, dynamic> json) => _$CoreMessageFromJson(json);
 
   CoreMessage copyWithExtensions(Map<String, dynamic> updated) {
     final newExtensions = Map<String, dynamic>.from(extensions);
diff --git a/packages/kai_engine_firebase_ai/lib/src/adapters.dart b/packages/kai_engine_firebase_ai/lib/src/adapters.dart
index bcc1bff..3dccc74 100644
--- a/packages/kai_engine_firebase_ai/lib/src/adapters.dart
+++ b/packages/kai_engine_firebase_ai/lib/src/adapters.dart
@@ -5,8 +5,7 @@ import 'package:firebase_ai/src/content.dart'; // Access parsePart function
 import 'package:kai_engine/kai_engine.dart';
 
 /// Adapter for Firebase AI content to Core Message
-class FirebaseAiContentAdapter
-    implements GenerativeMessageAdapterBase<Content> {
+class FirebaseAiContentAdapter implements GenerativeMessageAdapterBase<Content> {
   const FirebaseAiContentAdapter();
 
   @override
@@ -19,7 +18,7 @@ class FirebaseAiContentAdapter
       CoreMessageType.unknown => null,
     };
 
-    // Check if we have stored original Content
+    // Check if we have stored original Content with timestamp
     if (message.extensions.containsKey('_originalContent')) {
       final originalContentJson = message.extensions['_originalContent'];
       if (originalContentJson != null) {
@@ -33,8 +32,7 @@ class FirebaseAiContentAdapter
           final reconstructedParts = <Part>[];
           final originalJson = originalContentJson as Map<String, dynamic>;
           final originalParts =
-              (originalJson['parts'] as List?)?.cast<Map<String, dynamic>>() ??
-              [];
+              (originalJson['parts'] as List?)?.cast<Map<String, dynamic>>() ?? [];
 
           for (int i = 0; i < parsedContent.parts.length; i++) {
             final parsedPart = parsedContent.parts[i];
@@ -42,13 +40,10 @@ class FirebaseAiContentAdapter
             // If this was parsed as UnknownPart but we have original JSON, try to reconstruct manually
             // Also handle InlineDataPart with null willContinue field
             if ((parsedPart is UnknownPart ||
-                    (parsedPart is InlineDataPart &&
-                        parsedPart.willContinue == null)) &&
+                    (parsedPart is InlineDataPart && parsedPart.willContinue == null)) &&
                 i < originalParts.length) {
               final originalPartJson = originalParts[i];
-              final reconstructedPart = _reconstructPartFromJson(
-                originalPartJson,
-              );
+              final reconstructedPart = _reconstructPartFromJson(originalPartJson);
               reconstructedParts.add(reconstructedPart ?? parsedPart);
             } else {
               reconstructedParts.add(parsedPart);
@@ -86,7 +81,8 @@ class FirebaseAiContentAdapter
     final combinedText = textParts.map((part) => part.text).join('\n').trim();
 
     // Store the entire Content object as JSON for perfect reconstruction
-    extensions['_originalContent'] = content.toJson();
+    final contentJson = content.toJson();
+    extensions['_originalContent'] = contentJson;
 
     final messageType = switch (content.role) {
       'user' => CoreMessageType.user,
@@ -96,11 +92,22 @@ class FirebaseAiContentAdapter
       _ => CoreMessageType.unknown,
     };
 
-    return CoreMessage.create(
-      type: messageType,
-      content: combinedText,
-      extensions: extensions,
-    );
+    // Try to preserve timestamp from stored content or extensions
+    DateTime? timestamp;
+    if (contentJson.containsKey('_timestamp')) {
+      try {
+        timestamp = DateTime.parse(contentJson['_timestamp'] as String);
+      } catch (e) {
+        // If parsing fails, use current time
+      }
+    }
+
+    // Store timestamp in extensions for future round-trip conversions
+    if (timestamp != null) {
+      extensions['_timestamp'] = timestamp.toIso8601String();
+    }
+
+    return CoreMessage.create(type: messageType, content: combinedText, extensions: extensions);
   }
 
   /// Helper method to manually reconstruct parts that Firebase AI's parseContent doesn't handle
@@ -111,9 +118,7 @@ class FirebaseAiContentAdapter
       final response = fr['response'];
       return FunctionResponse(
         fr['name'] as String,
-        response is Map
-            ? Map<String, Object?>.from(response)
-            : <String, Object?>{},
+        response is Map ? Map<String, Object?>.from(response) : <String, Object?>{},
         id: fr['id'] as String?,
       );
     }
@@ -144,10 +149,7 @@ class FirebaseAiContentAdapter
     // Handle FileData
     if (partJson.containsKey('fileData')) {
       final fileData = partJson['fileData'] as Map<String, dynamic>;
-      return FileData(
-        fileData['mimeType'] as String,
-        fileData['fileUri'] as String,
-      );
+      return FileData(fileData['mimeType'] as String, fileData['fileUri'] as String);
     }
 
     // If we can't reconstruct it, return null

From fbee811e31a06e93a96024b274266feed0848d7c Mon Sep 17 00:00:00 2001
From: pckimlong <pochkimlong88@gmail.com>
Date: Sat, 10 Jan 2026 17:39:59 +0700
Subject: [PATCH 05/10] ``` chore: bump version to 0.2.0 across all packages

- Update kai_engine from 0.1.1 to 0.2.0
- Update kai_engine_chat_ui from 0.1.1 to 0.2.0
- Update kai_engine_firebase_ai from 0.1.1 to 0.2.0
- Update prompt_block from 0.1.1 to 0.2.0
```
---
 packages/kai_engine/pubspec.yaml             | 2 +-
 packages/kai_engine_chat_ui/pubspec.yaml     | 2 +-
 packages/kai_engine_firebase_ai/pubspec.yaml | 2 +-
 packages/prompt_block/pubspec.yaml           | 2 +-
 4 files changed, 4 insertions(+), 4 deletions(-)

diff --git a/packages/kai_engine/pubspec.yaml b/packages/kai_engine/pubspec.yaml
index 19d16eb..1afb339 100644
--- a/packages/kai_engine/pubspec.yaml
+++ b/packages/kai_engine/pubspec.yaml
@@ -1,6 +1,6 @@
 name: kai_engine
 description: "A modular, extensible AI chat engine built with a pipeline-based architecture."
-version: 0.1.1
+version: 0.2.0
 homepage: https://github.com/pckimlong/kai_engine
 repository: https://github.com/pckimlong/kai_engine
 issue_tracker: https://github.com/pckimlong/kai_engine/issues
diff --git a/packages/kai_engine_chat_ui/pubspec.yaml b/packages/kai_engine_chat_ui/pubspec.yaml
index 5515f23..7e22378 100644
--- a/packages/kai_engine_chat_ui/pubspec.yaml
+++ b/packages/kai_engine_chat_ui/pubspec.yaml
@@ -1,6 +1,6 @@
 name: kai_engine_chat_ui
 description: Flutter chat UI widgets for kai_engine (message list, composer, streaming UI).
-version: 0.1.1
+version: 0.2.0
 homepage: https://github.com/pckimlong/kai_engine
 repository: https://github.com/pckimlong/kai_engine
 issue_tracker: https://github.com/pckimlong/kai_engine/issues
diff --git a/packages/kai_engine_firebase_ai/pubspec.yaml b/packages/kai_engine_firebase_ai/pubspec.yaml
index fdc522d..a0f919f 100644
--- a/packages/kai_engine_firebase_ai/pubspec.yaml
+++ b/packages/kai_engine_firebase_ai/pubspec.yaml
@@ -1,6 +1,6 @@
 name: kai_engine_firebase_ai
 description: "A Firebase AI adapter for the Kai Engine."
-version: 0.1.1
+version: 0.2.0
 homepage: https://github.com/pckimlong/kai_engine
 repository: https://github.com/pckimlong/kai_engine
 issue_tracker: https://github.com/pckimlong/kai_engine/issues
diff --git a/packages/prompt_block/pubspec.yaml b/packages/prompt_block/pubspec.yaml
index 956f683..7e77ad2 100644
--- a/packages/prompt_block/pubspec.yaml
+++ b/packages/prompt_block/pubspec.yaml
@@ -1,6 +1,6 @@
 name: prompt_block
 description: "A package for creating and managing prompt blocks in AI applications."
-version: 0.1.1
+version: 0.2.0
 homepage: https://github.com/pckimlong/kai_engine
 repository: https://github.com/pckimlong/kai_engine
 issue_tracker: https://github.com/pckimlong/kai_engine/issues

From d947542331c642c60a5a828ad82732455df78719 Mon Sep 17 00:00:00 2001
From: pckimlong <pochkimlong88@gmail.com>
Date: Sat, 10 Jan 2026 17:56:36 +0700
Subject: [PATCH 06/10] refactor: simplify timestamp formatting and fix error
 state handling

- Change _formatTime to display HH:mm format instead of relative time
- Use _setState instead of direct stream add in error handling for consistency
---
 .../lib/src/chat_controller_base.dart           |  2 +-
 .../lib/src/widgets/kai_message_bubble.dart     | 17 +----------------
 2 files changed, 2 insertions(+), 17 deletions(-)

diff --git a/packages/kai_engine/lib/src/chat_controller_base.dart b/packages/kai_engine/lib/src/chat_controller_base.dart
index ba08b6e..eedf143 100644
--- a/packages/kai_engine/lib/src/chat_controller_base.dart
+++ b/packages/kai_engine/lib/src/chat_controller_base.dart
@@ -161,7 +161,7 @@ abstract base class ChatControllerBase<TEntity> {
       final errorState = GenerationState<GenerationResult>.error(
         KaiException.exception(error.toString(), stackTrace),
       );
-      _generationStateController.add(errorState);
+      _setState(errorState);
       return errorState;
     }
   }
diff --git a/packages/kai_engine_chat_ui/lib/src/widgets/kai_message_bubble.dart b/packages/kai_engine_chat_ui/lib/src/widgets/kai_message_bubble.dart
index 2549b8e..b20c75f 100644
--- a/packages/kai_engine_chat_ui/lib/src/widgets/kai_message_bubble.dart
+++ b/packages/kai_engine_chat_ui/lib/src/widgets/kai_message_bubble.dart
@@ -114,21 +114,6 @@ class KaiMessageBubble extends StatelessWidget {
   }
 
   String _formatTime(DateTime timestamp) {
-    final now = DateTime.now();
-    final difference = now.difference(timestamp);
-
-    if (difference.inMinutes < 1) {
-      return 'Just now';
-    } else if (difference.inMinutes < 60) {
-      return '${difference.inMinutes}m ago';
-    } else if (difference.inHours < 24) {
-      return '${difference.inHours}h ago';
-    } else if (difference.inDays == 1) {
-      return 'Yesterday';
-    } else if (difference.inDays < 7) {
-      return '${difference.inDays}d ago';
-    } else {
-      return '${timestamp.day}/${timestamp.month}/${timestamp.year}';
-    }
+    return '${timestamp.hour.toString().padLeft(2, '0')}:${timestamp.minute.toString().padLeft(2, '0')}';
   }
 }

From 1b8a6182310a1cefa09213a712f0e9875cadbe89 Mon Sep 17 00:00:00 2001
From: pckimlong <pochkimlong88@gmail.com>
Date: Sat, 10 Jan 2026 18:15:08 +0700
Subject: [PATCH 07/10] refactor(chat_ui): adjust spacing, add inline composer
 button and empty state support
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

- Reduce default bubble widths (mobile: 320→280, desktop: 480→400)
- Reduce bubble and list padding for tighter layout
- Change bottom item spacing from theme.itemSpacing to 4px for first item
- Adjust bubble corner radius from 6 to 4 for sharper corners
- Add inline send/stop button option to KaiComposer with customization (icon, size, color)
- Add avatarBuilder support for custom message avatars
- Add em
---
 .../lib/src/kai_chat_theme.dart               | 29 ++----
 .../lib/src/kai_chat_view.dart                |  9 +-
 .../lib/src/widgets/kai_composer.dart         | 99 ++++++++++++++++++-
 .../lib/src/widgets/kai_message_bubble.dart   | 12 ++-
 .../lib/src/widgets/kai_message_list.dart     | 75 +++++++++++++-
 5 files changed, 196 insertions(+), 28 deletions(-)

diff --git a/packages/kai_engine_chat_ui/lib/src/kai_chat_theme.dart b/packages/kai_engine_chat_ui/lib/src/kai_chat_theme.dart
index 63e0235..39d92ac 100644
--- a/packages/kai_engine_chat_ui/lib/src/kai_chat_theme.dart
+++ b/packages/kai_engine_chat_ui/lib/src/kai_chat_theme.dart
@@ -17,14 +17,11 @@ class KaiChatTheme extends ThemeExtension<KaiChatTheme> {
   final Color? aiTextColor;
 
   const KaiChatTheme({
-    this.maxBubbleWidthMobile = 320,
-    this.maxBubbleWidthDesktop = 480,
+    this.maxBubbleWidthMobile = 280,
+    this.maxBubbleWidthDesktop = 400,
     this.bubbleRadius = 16,
-    this.bubblePadding = const EdgeInsets.symmetric(
-      horizontal: 12,
-      vertical: 10,
-    ),
-    this.listPadding = const EdgeInsets.symmetric(horizontal: 16, vertical: 12),
+    this.bubblePadding = const EdgeInsets.symmetric(horizontal: 12, vertical: 8),
+    this.listPadding = const EdgeInsets.symmetric(horizontal: 16, vertical: 4),
     this.itemSpacing = 12,
     this.userBubbleColor,
     this.aiBubbleColor,
@@ -47,8 +44,7 @@ class KaiChatTheme extends ThemeExtension<KaiChatTheme> {
   }) {
     return KaiChatTheme(
       maxBubbleWidthMobile: maxBubbleWidthMobile ?? this.maxBubbleWidthMobile,
-      maxBubbleWidthDesktop:
-          maxBubbleWidthDesktop ?? this.maxBubbleWidthDesktop,
+      maxBubbleWidthDesktop: maxBubbleWidthDesktop ?? this.maxBubbleWidthDesktop,
       bubbleRadius: bubbleRadius ?? this.bubbleRadius,
       bubblePadding: bubblePadding ?? this.bubblePadding,
       listPadding: listPadding ?? this.listPadding,
@@ -64,16 +60,8 @@ class KaiChatTheme extends ThemeExtension<KaiChatTheme> {
   KaiChatTheme lerp(ThemeExtension<KaiChatTheme>? other, double t) {
     if (other is! KaiChatTheme) return this;
     return KaiChatTheme(
-      maxBubbleWidthMobile: lerpDouble(
-        maxBubbleWidthMobile,
-        other.maxBubbleWidthMobile,
-        t,
-      )!,
-      maxBubbleWidthDesktop: lerpDouble(
-        maxBubbleWidthDesktop,
-        other.maxBubbleWidthDesktop,
-        t,
-      )!,
+      maxBubbleWidthMobile: lerpDouble(maxBubbleWidthMobile, other.maxBubbleWidthMobile, t)!,
+      maxBubbleWidthDesktop: lerpDouble(maxBubbleWidthDesktop, other.maxBubbleWidthDesktop, t)!,
       bubbleRadius: lerpDouble(bubbleRadius, other.bubbleRadius, t)!,
       bubblePadding: EdgeInsets.lerp(bubblePadding, other.bubblePadding, t)!,
       listPadding: EdgeInsets.lerp(listPadding, other.listPadding, t)!,
@@ -87,6 +75,5 @@ class KaiChatTheme extends ThemeExtension<KaiChatTheme> {
 }
 
 extension KaiChatThemeX on BuildContext {
-  KaiChatTheme get kaiChatTheme =>
-      Theme.of(this).extension<KaiChatTheme>() ?? const KaiChatTheme();
+  KaiChatTheme get kaiChatTheme => Theme.of(this).extension<KaiChatTheme>() ?? const KaiChatTheme();
 }
diff --git a/packages/kai_engine_chat_ui/lib/src/kai_chat_view.dart b/packages/kai_engine_chat_ui/lib/src/kai_chat_view.dart
index 480dc23..e4c8d59 100644
--- a/packages/kai_engine_chat_ui/lib/src/kai_chat_view.dart
+++ b/packages/kai_engine_chat_ui/lib/src/kai_chat_view.dart
@@ -22,13 +22,15 @@ class KaiChatView extends StatefulWidget {
     this.transientBuilder,
     this.transientItemBuilder,
     this.functionCallBuilder,
+    this.emptyStateBuilder,
     this.onMessageTap,
     this.reverse = true,
     this.scrollController,
-    this.composerPadding = const EdgeInsets.fromLTRB(16, 0, 16, 12),
+    this.composerPadding = const EdgeInsets.fromLTRB(16, 12, 16, 12),
     this.listPadding,
     this.composerBuilder,
     this.showTimeLabel = false,
+    this.avatarBuilder,
   });
 
   final ChatControllerBase controller;
@@ -58,7 +60,9 @@ class KaiChatView extends StatefulWidget {
   final KaiTransientBuilder? transientBuilder;
   final KaiTransientItemBuilder? transientItemBuilder;
   final KaiFunctionCallBuilder? functionCallBuilder;
+  final KaiEmptyStateBuilder? emptyStateBuilder;
   final void Function(CoreMessage message)? onMessageTap;
+  final KaiAvatarBuilder? avatarBuilder;
 
   final bool reverse;
   final ScrollController? scrollController;
@@ -157,9 +161,12 @@ class _KaiChatViewState extends State<KaiChatView> {
                     transientBuilder: widget.transientBuilder,
                     transientItemBuilder: widget.transientItemBuilder,
                     functionCallBuilder: widget.functionCallBuilder,
+                    emptyStateBuilder: widget.emptyStateBuilder,
+                    onEmptyStateSubmit: _handleSend,
                     onMessageTap: widget.onMessageTap,
                     padding: widget.listPadding ?? theme.listPadding,
                     showTimeLabel: widget.showTimeLabel,
+                    avatarBuilder: widget.avatarBuilder,
                   );
                 },
               ),
diff --git a/packages/kai_engine_chat_ui/lib/src/widgets/kai_composer.dart b/packages/kai_engine_chat_ui/lib/src/widgets/kai_composer.dart
index a485ba0..d84031e 100644
--- a/packages/kai_engine_chat_ui/lib/src/widgets/kai_composer.dart
+++ b/packages/kai_engine_chat_ui/lib/src/widgets/kai_composer.dart
@@ -97,6 +97,19 @@ class KaiComposerCallbacks {
 ///   },
 /// )
 /// ```
+///
+/// Example - Customize inline button:
+/// ```dart
+/// KaiComposer(
+///   onSend: (text) => controller.submit(text),
+///   isGenerating: isGenerating,
+///   inlineButton: true,
+///   sendIcon: Icons.send,
+///   stopIcon: Icons.close,
+///   buttonSize: 40,
+///   buttonColor: Colors.blue,
+/// )
+/// ```
 typedef KaiComposerBuilder =
     Widget Function(BuildContext context, KaiComposerState state, KaiComposerCallbacks callbacks);
 
@@ -109,9 +122,14 @@ class KaiComposer extends StatefulWidget {
     this.hintText = 'Type a message…',
     this.controller,
     this.autofocus = false,
-    this.maxLines = 6,
+    this.maxLines = 8,
     this.onError,
     this.builder,
+    this.inlineButton = true,
+    this.sendIcon,
+    this.stopIcon,
+    this.buttonSize = 36,
+    this.buttonColor,
   });
 
   final Future<void> Function(String text) onSend;
@@ -123,6 +141,11 @@ class KaiComposer extends StatefulWidget {
   final int maxLines;
   final void Function(Object error, StackTrace stackTrace)? onError;
   final KaiComposerBuilder? builder;
+  final bool inlineButton;
+  final IconData? sendIcon;
+  final IconData? stopIcon;
+  final double buttonSize;
+  final Color? buttonColor;
 
   @override
   State<KaiComposer> createState() => _KaiComposerState();
@@ -192,6 +215,80 @@ class _KaiComposerState extends State<KaiComposer> {
 
     final colors = Theme.of(context).colorScheme;
 
+    if (widget.inlineButton) {
+      final effectiveButtonColor = widget.buttonColor ?? colors.primary;
+      final effectiveSendIcon = widget.sendIcon ?? Icons.arrow_upward;
+      final effectiveStopIcon = widget.stopIcon ?? Icons.stop;
+
+      return Row(
+        children: [
+          Expanded(
+            child: TextField(
+              controller: _controller,
+              autofocus: widget.autofocus,
+              maxLines: widget.maxLines,
+              minLines: 1,
+              textCapitalization: TextCapitalization.sentences,
+              decoration: InputDecoration(
+                hintText: widget.hintText,
+                filled: true,
+                fillColor: colors.surfaceContainerHighest,
+                border: OutlineInputBorder(
+                  borderRadius: BorderRadius.circular(16),
+                  borderSide: BorderSide.none,
+                ),
+                contentPadding: const EdgeInsets.symmetric(horizontal: 14, vertical: 12),
+                suffixIcon: Padding(
+                  padding: const EdgeInsets.only(right: 8),
+                  child: Semantics(
+                    button: true,
+                    label: widget.isGenerating ? 'Stop generating' : 'Send message',
+                    child: Tooltip(
+                      message: widget.isGenerating ? 'Stop' : 'Send message',
+                      child: Material(
+                        color: Colors.transparent,
+                        child: Container(
+                          width: widget.buttonSize,
+                          height: widget.buttonSize,
+                          decoration: BoxDecoration(
+                            color: effectiveButtonColor,
+                            shape: BoxShape.circle,
+                            border: Border.all(
+                              color: colors.outlineVariant.withValues(alpha: 0.3),
+                              width: 1,
+                            ),
+                          ),
+                          child: InkWell(
+                            onTap: widget.isGenerating
+                                ? widget.onCancel
+                                : canSend
+                                ? _handleSend
+                                : null,
+                            borderRadius: BorderRadius.circular(widget.buttonSize / 2),
+                            child: Center(
+                              child: Icon(
+                                widget.isGenerating ? effectiveStopIcon : effectiveSendIcon,
+                                color: effectiveButtonColor == colors.primary
+                                    ? colors.onPrimary
+                                    : colors.onPrimaryContainer,
+                                size: widget.buttonSize * 0.56,
+                              ),
+                            ),
+                          ),
+                        ),
+                      ),
+                    ),
+                  ),
+                ),
+              ),
+              onSubmitted: (_) => _handleSend(),
+              onChanged: (_) => setState(() {}),
+            ),
+          ),
+        ],
+      );
+    }
+
     return Row(
       children: [
         Expanded(
diff --git a/packages/kai_engine_chat_ui/lib/src/widgets/kai_message_bubble.dart b/packages/kai_engine_chat_ui/lib/src/widgets/kai_message_bubble.dart
index b20c75f..5ffab90 100644
--- a/packages/kai_engine_chat_ui/lib/src/widgets/kai_message_bubble.dart
+++ b/packages/kai_engine_chat_ui/lib/src/widgets/kai_message_bubble.dart
@@ -3,6 +3,7 @@ import 'package:gpt_markdown/gpt_markdown.dart';
 import 'package:kai_engine/kai_engine.dart';
 
 import '../kai_chat_theme.dart';
+import 'kai_message_list.dart';
 
 typedef KaiMessageContentBuilder = Widget Function(BuildContext context, CoreMessage message);
 
@@ -16,6 +17,7 @@ class KaiMessageBubble extends StatelessWidget {
     this.isUserMessage,
     this.semanticLabel,
     this.showTimeLabel = false,
+    this.avatarBuilder,
   });
 
   final CoreMessage message;
@@ -24,6 +26,7 @@ class KaiMessageBubble extends StatelessWidget {
   final bool? isUserMessage;
   final String? semanticLabel;
   final bool showTimeLabel;
+  final KaiAvatarBuilder? avatarBuilder;
 
   @override
   Widget build(BuildContext context) {
@@ -66,7 +69,12 @@ class KaiMessageBubble extends StatelessWidget {
 
     return Row(
       mainAxisAlignment: isUser ? MainAxisAlignment.end : MainAxisAlignment.start,
+      crossAxisAlignment: CrossAxisAlignment.end,
       children: [
+        if (!isUser && avatarBuilder != null) ...[
+          avatarBuilder!(context, message),
+          const SizedBox(width: 8),
+        ],
         Flexible(
           child: Semantics(
             container: true,
@@ -89,9 +97,9 @@ class KaiMessageBubble extends StatelessWidget {
                       borderRadius: BorderRadius.circular(theme.bubbleRadius).copyWith(
                         bottomLeft: isUser
                             ? Radius.circular(theme.bubbleRadius)
-                            : const Radius.circular(6),
+                            : const Radius.circular(4),
                         bottomRight: isUser
-                            ? const Radius.circular(6)
+                            ? const Radius.circular(4)
                             : Radius.circular(theme.bubbleRadius),
                       ),
                       border: Border.all(
diff --git a/packages/kai_engine_chat_ui/lib/src/widgets/kai_message_list.dart b/packages/kai_engine_chat_ui/lib/src/widgets/kai_message_list.dart
index a862bac..c14496f 100644
--- a/packages/kai_engine_chat_ui/lib/src/widgets/kai_message_list.dart
+++ b/packages/kai_engine_chat_ui/lib/src/widgets/kai_message_list.dart
@@ -7,6 +7,26 @@ import 'kai_message_bubble.dart';
 import 'kai_streaming_text.dart';
 import 'kai_typing_indicator.dart';
 
+/// Builder that constructs a custom avatar widget.
+///
+/// This builder receives the [context] and [message] as parameters.
+/// You can use this to create a custom avatar for messages.
+///
+/// Example - Create a simple avatar:
+/// ```dart
+/// KaiChatView(
+///   avatarBuilder: (context, message) {
+///     return CircleAvatar(
+///       radius: 16,
+///       backgroundColor: Colors.blue,
+///       child: Icon(Icons.smart_toy, size: 18),
+///     );
+///   },
+///   // ...
+/// )
+/// ```
+typedef KaiAvatarBuilder = Widget Function(BuildContext context, CoreMessage message);
+
 /// Builder that constructs a custom message widget.
 ///
 /// This builder receives the [context] and [message] as parameters.
@@ -147,6 +167,42 @@ typedef KaiMessageFilter = bool Function(CoreMessage message);
 /// ```
 typedef KaiFunctionCallBuilder = Widget Function(BuildContext context, CoreMessage message);
 
+/// Builder that constructs a custom empty state widget.
+///
+/// This builder receives the [context] and [onSubmit] callback as parameters.
+/// The [onSubmit] callback allows you to trigger message submission from the empty state.
+/// If provided, this widget is shown when there are no messages in the chat.
+/// If null, an empty list is shown.
+///
+/// Example - Show a welcome message with a quick action button:
+/// ```dart
+/// KaiChatView(
+///   emptyStateBuilder: (context, onSubmit) {
+///     return Center(
+///       child: Column(
+///         mainAxisAlignment: MainAxisAlignment.center,
+///         children: [
+///           Icon(Icons.chat_bubble_outline, size: 64, color: Colors.grey),
+///           SizedBox(height: 16),
+///           Text(
+///             'Start a conversation!',
+///             style: TextStyle(fontSize: 18, color: Colors.grey),
+///           ),
+///           SizedBox(height: 16),
+///           ElevatedButton(
+///             onPressed: () => onSubmit('Hello, I need help!'),
+///             child: Text('Say Hello'),
+///           ),
+///         ],
+///       ),
+///     );
+///   },
+///   // ...
+/// )
+/// ```
+typedef KaiEmptyStateBuilder =
+    Widget Function(BuildContext context, void Function(String text) onSubmit);
+
 bool defaultKaiMessageFilter(CoreMessage m) => m.isDisplayable && !m.isBackgroundContext;
 
 class KaiMessageList extends StatelessWidget {
@@ -162,10 +218,13 @@ class KaiMessageList extends StatelessWidget {
     this.transientBuilder,
     this.transientItemBuilder,
     this.functionCallBuilder,
+    this.emptyStateBuilder,
+    this.onEmptyStateSubmit,
     this.onMessageTap,
     this.padding,
     this.physics,
     this.showTimeLabel = false,
+    this.avatarBuilder,
   });
 
   final IList<CoreMessage> messages;
@@ -178,10 +237,13 @@ class KaiMessageList extends StatelessWidget {
   final KaiTransientBuilder? transientBuilder;
   final KaiTransientItemBuilder? transientItemBuilder;
   final KaiFunctionCallBuilder? functionCallBuilder;
+  final KaiEmptyStateBuilder? emptyStateBuilder;
+  final void Function(String text)? onEmptyStateSubmit;
   final void Function(CoreMessage message)? onMessageTap;
   final EdgeInsetsGeometry? padding;
   final ScrollPhysics? physics;
   final bool showTimeLabel;
+  final KaiAvatarBuilder? avatarBuilder;
 
   @override
   Widget build(BuildContext context) {
@@ -196,6 +258,10 @@ class KaiMessageList extends StatelessWidget {
     final totalCount = filtered.length + (transient ? 1 : 0);
     final transientIndex = reverse ? 0 : filtered.length;
 
+    if (totalCount == 0 && emptyStateBuilder != null) {
+      return emptyStateBuilder!(context, onEmptyStateSubmit ?? (_) {});
+    }
+
     return ListView.builder(
       controller: controller,
       reverse: reverse,
@@ -210,7 +276,7 @@ class KaiMessageList extends StatelessWidget {
               transientBuilder?.call(context, state) ?? _defaultTransient(context, state);
           final built = transientItemBuilder?.call(context, state, defaultChild) ?? defaultChild;
           return Padding(
-            padding: EdgeInsets.only(bottom: theme.itemSpacing),
+            padding: EdgeInsets.only(bottom: index == 0 ? 4 : theme.itemSpacing),
             child: built,
           );
         }
@@ -222,7 +288,7 @@ class KaiMessageList extends StatelessWidget {
 
         if (message.type == CoreMessageType.function && functionCallBuilder != null) {
           return Padding(
-            padding: EdgeInsets.only(bottom: theme.itemSpacing),
+            padding: EdgeInsets.only(bottom: index == 0 ? 4 : theme.itemSpacing),
             child: functionCallBuilder!(context, message),
           );
         }
@@ -232,6 +298,7 @@ class KaiMessageList extends StatelessWidget {
           isUserMessage: message.type == CoreMessageType.user,
           onTap: onMessageTap == null ? null : () => onMessageTap!(message),
           showTimeLabel: showTimeLabel,
+          avatarBuilder: avatarBuilder,
         );
 
         final bubble =
@@ -240,7 +307,7 @@ class KaiMessageList extends StatelessWidget {
             defaultChild;
 
         return Padding(
-          padding: EdgeInsets.only(bottom: theme.itemSpacing),
+          padding: EdgeInsets.only(bottom: index == 0 ? 4 : theme.itemSpacing),
           child: bubble,
         );
       },
@@ -254,12 +321,14 @@ class KaiMessageList extends StatelessWidget {
         isUserMessage: false,
         contentBuilder: (_, __) =>
             const Padding(padding: EdgeInsets.symmetric(vertical: 2), child: KaiTypingIndicator()),
+        avatarBuilder: avatarBuilder,
       ),
       GenerationStreamingTextState(text: final text) => KaiMessageBubble(
         message: CoreMessage.ai(content: text),
         isUserMessage: false,
         contentBuilder: (context, _) =>
             KaiStreamingText(text: text, style: Theme.of(context).textTheme.bodyMedium),
+        avatarBuilder: avatarBuilder,
       ),
       _ => const SizedBox.shrink(),
     };

From 08dc1ebd214b7147535d2026baf68323c6b3664d Mon Sep 17 00:00:00 2001
From: pckimlong <pochkimlong88@gmail.com>
Date: Sat, 11 Apr 2026 16:41:26 +0700
Subject: [PATCH 08/10] format

---
 .../lib/src/widgets/kai_message_bubble.dart                   | 4 +---
 1 file changed, 1 insertion(+), 3 deletions(-)

diff --git a/packages/kai_engine_chat_ui/lib/src/widgets/kai_message_bubble.dart b/packages/kai_engine_chat_ui/lib/src/widgets/kai_message_bubble.dart
index 5ffab90..8d5006a 100644
--- a/packages/kai_engine_chat_ui/lib/src/widgets/kai_message_bubble.dart
+++ b/packages/kai_engine_chat_ui/lib/src/widgets/kai_message_bubble.dart
@@ -41,9 +41,7 @@ class KaiMessageBubble extends StatelessWidget {
         ? (theme.userTextColor ?? colors.onPrimary)
         : (theme.aiTextColor ?? colors.onSurface);
 
-    final textStyle = Theme.of(
-      context,
-    ).textTheme.bodyMedium?.copyWith(color: textColor, height: 1.35);
+    final textStyle = Theme.of(context).textTheme.bodyMedium?.copyWith(color: textColor);
 
     final content =
         contentBuilder?.call(context, message) ??

From 0ca72e83ccae543f2e871b4f12e3ab201fdbc00e Mon Sep 17 00:00:00 2001
From: pckimlong <pochkimlong88@gmail.com>
Date: Sat, 11 Apr 2026 16:45:05 +0700
Subject: [PATCH 09/10] format

---
 .../lib/src/chat_controller_base.dart         |  43 +-
 .../lib/src/conversation_manager.dart         |  27 +-
 .../lib/src/message_repository_base.dart      |  30 +-
 .../lib/src/models/core_message.dart          |   9 +-
 .../kai_engine/lib/src/prompt_engine.dart     |  61 +-
 .../test/src/chat_controller_base_test.dart   | 552 +++++++++++-------
 .../src/message_repository_base_test.dart     | 381 ++++++------
 7 files changed, 684 insertions(+), 419 deletions(-)

diff --git a/packages/kai_engine/lib/src/chat_controller_base.dart b/packages/kai_engine/lib/src/chat_controller_base.dart
index eedf143..7d88b45 100644
--- a/packages/kai_engine/lib/src/chat_controller_base.dart
+++ b/packages/kai_engine/lib/src/chat_controller_base.dart
@@ -27,9 +27,10 @@ abstract base class ChatControllerBase<TEntity> {
        _cancelToken = CancelToken();
 
   /// Handle generation state updates
-  final _generationStateController = BehaviorSubject<GenerationState<GenerationResult>>.seeded(
-    GenerationState.initial(),
-  );
+  final _generationStateController =
+      BehaviorSubject<GenerationState<GenerationResult>>.seeded(
+        GenerationState.initial(),
+      );
 
   ContextEngine build();
 
@@ -73,7 +74,9 @@ abstract base class ChatControllerBase<TEntity> {
         // Exclude the current user message from history, by natural flow user message is happening after query input
         // but to trickly show message in UI faster, we trickly hide this
         histories: await _conversationManager.getMessages().then(
-          (messages) => messages.whereNot((m) => m.messageId == userMessage.messageId).toIList(),
+          (messages) => messages
+              .whereNot((m) => m.messageId == userMessage.messageId)
+              .toIList(),
         ),
       );
 
@@ -101,7 +104,9 @@ abstract base class ChatControllerBase<TEntity> {
       final insertedUserMessage = await insertedUserMessageFuture;
       userMessage = insertedUserMessage;
       if (contextResult.userMessage.content != insertedUserMessage.content) {
-        final revised = insertedUserMessage.copyWith(content: contextResult.userMessage.content);
+        final revised = insertedUserMessage.copyWith(
+          content: contextResult.userMessage.content,
+        );
         await _conversationManager.updateMessages(IList([revised]));
         userMessage = revised;
       }
@@ -119,7 +124,9 @@ abstract base class ChatControllerBase<TEntity> {
         onStateUpdate: (state) => _generationStateController.add(state),
       );
 
-      generationResult = await AIGenerationPhase(_generationService).execute(aiGenerationInput);
+      generationResult = await AIGenerationPhase(
+        _generationService,
+      ).execute(aiGenerationInput);
 
       // Save the AI-generated messages to the conversation
       if (generationResult.generatedMessages.isNotEmpty) {
@@ -130,12 +137,16 @@ abstract base class ChatControllerBase<TEntity> {
               final originalIds = generationResult.generatedMessages
                   .map((e) => e.messageId)
                   .toSet();
-              return items.where((e) => originalIds.contains(e.messageId)).toIList();
+              return items
+                  .where((e) => originalIds.contains(e.messageId))
+                  .toIList();
             });
 
         // Make sure generated message is updated to date with required metadata
         // before pass to other below service
-        generationResult = generationResult.copyWith(generatedMessages: addedMessages);
+        generationResult = generationResult.copyWith(
+          generatedMessages: addedMessages,
+        );
       }
 
       if (_postResponseEngine != null) {
@@ -150,7 +161,9 @@ abstract base class ChatControllerBase<TEntity> {
 
         await _postResponseEngine.execute(postResponseInput);
       }
-      final generationState = GenerationState<GenerationResult>.complete(generationResult);
+      final generationState = GenerationState<GenerationResult>.complete(
+        generationResult,
+      );
       _generationStateController.add(generationState);
       return generationState;
     } catch (error, stackTrace) {
@@ -171,8 +184,10 @@ abstract base class ChatControllerBase<TEntity> {
     _setState(GenerationState.error(KaiException.cancelled()));
   }
 
-  Stream<IList<CoreMessage>> get messagesStream => _conversationManager.messagesStream;
-  Future<IList<CoreMessage>> getAllMessages() => _conversationManager.getMessages();
+  Stream<IList<CoreMessage>> get messagesStream =>
+      _conversationManager.messagesStream;
+  Future<IList<CoreMessage>> getAllMessages() =>
+      _conversationManager.getMessages();
 
   Stream<GenerationState<GenerationResult>> get generationStateStream =>
       _generationStateController.stream;
@@ -188,7 +203,8 @@ abstract base class ChatControllerBase<TEntity> {
     onStateChange(state);
   }
 
-  void _setLoadingPhase(LoadingPhase phase) => _setState(GenerationState.loading(phase));
+  void _setLoadingPhase(LoadingPhase phase) =>
+      _setState(GenerationState.loading(phase));
 
   /// Callback to listen for state changes, this helpful for implementing custom logic or log the
   /// state provide more flexible. default to do nothing
@@ -349,7 +365,8 @@ final class KaiChatController extends ChatControllerBase<CoreMessage> {
     );
 
     return KaiChatController._(
-      contextEngine: contextEngine ?? SimpleContextEngine(systemPrompt: systemPrompt),
+      contextEngine:
+          contextEngine ?? SimpleContextEngine(systemPrompt: systemPrompt),
       conversationManager: conversation,
       queryEngine: queryEngine,
       postResponseEngine: postEngine,
diff --git a/packages/kai_engine/lib/src/conversation_manager.dart b/packages/kai_engine/lib/src/conversation_manager.dart
index 36b9274..72a864c 100644
--- a/packages/kai_engine/lib/src/conversation_manager.dart
+++ b/packages/kai_engine/lib/src/conversation_manager.dart
@@ -18,7 +18,9 @@ class ConversationManager<T> {
   IList<CoreMessage> _messages = IList(const []);
   final BehaviorSubject<IList<CoreMessage>> _messagesController =
       BehaviorSubject<IList<CoreMessage>>();
-  final BehaviorSubject<bool> _loadingController = BehaviorSubject<bool>.seeded(false);
+  final BehaviorSubject<bool> _loadingController = BehaviorSubject<bool>.seeded(
+    false,
+  );
 
   ConversationManager._({
     required this.session,
@@ -57,7 +59,9 @@ class ConversationManager<T> {
   /// Adds a message to the conversation
   Future<IList<CoreMessage>> addMessages(IList<CoreMessage> messages) async {
     // Ensure system messages are ignored
-    final persistentMessages = messages.where((m) => m.isMessageSavable).toList();
+    final persistentMessages = messages
+        .where((m) => m.isMessageSavable)
+        .toList();
 
     // Store original state for rollback
     final originalMessages = _messages;
@@ -78,7 +82,9 @@ class ConversationManager<T> {
           .then((e) => e.map(_messageAdapter.toCoreMessage));
 
       _messages = _messages
-          .removeWhere((e) => persistentMessages.any((m) => m.messageId == e.messageId))
+          .removeWhere(
+            (e) => persistentMessages.any((m) => m.messageId == e.messageId),
+          )
           .addAll(result);
       _messagesController.add(_messages);
       return result.toIList();
@@ -91,7 +97,9 @@ class ConversationManager<T> {
   }
 
   Future<void> updateMessages(IList<CoreMessage> messages) async {
-    final persistentMessages = messages.where((m) => m.isMessageSavable).toList();
+    final persistentMessages = messages
+        .where((m) => m.isMessageSavable)
+        .toList();
 
     // Store original state for rollback
     final originalMessages = _messages;
@@ -109,7 +117,9 @@ class ConversationManager<T> {
       // Update repository
       final result = await _repository
           .updateMessages(
-            persistentMessages.map((e) => _messageAdapter.fromCoreMessage(e, session: session)),
+            persistentMessages.map(
+              (e) => _messageAdapter.fromCoreMessage(e, session: session),
+            ),
           )
           .then((e) => e.map(_messageAdapter.toCoreMessage));
 
@@ -144,7 +154,9 @@ class ConversationManager<T> {
 
       // Remove from repository
       await _repository.removeMessages(
-        messages.map((e) => _messageAdapter.fromCoreMessage(e, session: session)),
+        messages.map(
+          (e) => _messageAdapter.fromCoreMessage(e, session: session),
+        ),
       );
     } catch (error) {
       // Rollback optimistic update on failure
@@ -168,7 +180,8 @@ class ConversationManager<T> {
 }
 
 /// Prebuilt for in memory management of conversation messages
-final class InMemoryConversationManager extends ConversationManager<CoreMessage> {
+final class InMemoryConversationManager
+    extends ConversationManager<CoreMessage> {
   InMemoryConversationManager({
     required super.session,
     MessageRepositoryBase<CoreMessage>? repository,
diff --git a/packages/kai_engine/lib/src/message_repository_base.dart b/packages/kai_engine/lib/src/message_repository_base.dart
index 42fed4a..98a0e42 100644
--- a/packages/kai_engine/lib/src/message_repository_base.dart
+++ b/packages/kai_engine/lib/src/message_repository_base.dart
@@ -20,7 +20,8 @@ abstract interface class MessageRepositoryBase<T> {
 }
 
 /// Abstract interface for a core message repository. No adapter required.
-abstract interface class CoreMessageRepositoryBase extends MessageRepositoryBase<CoreMessage> {}
+abstract interface class CoreMessageRepositoryBase
+    extends MessageRepositoryBase<CoreMessage> {}
 
 /// Prebuilt repository for memory persistence.
 /// Messages are automatically sorted by timestamp on initialization.
@@ -75,15 +76,24 @@ final class InMemoryMessageRepository implements CoreMessageRepositoryBase {
 /// its own internal state for immediate read operations. Messages are
 /// automatically sorted by timestamp on initial load.
 final class CoreMessageRepository implements CoreMessageRepositoryBase {
-  final Future<Iterable<CoreMessage>> Function(ConversationSession session) onInitial;
-  final Future<void> Function(ConversationSession session, Iterable<CoreMessage> messages) onPut;
+  final Future<Iterable<CoreMessage>> Function(ConversationSession session)
+  onInitial;
+  final Future<void> Function(
+    ConversationSession session,
+    Iterable<CoreMessage> messages,
+  )
+  onPut;
   final Future<void> Function(Iterable<CoreMessage> messages)? onRemove;
 
   final List<CoreMessage> _messages = [];
   ConversationSession? _cachedSession;
   bool _isInitialized = false;
 
-  CoreMessageRepository({required this.onInitial, required this.onPut, this.onRemove});
+  CoreMessageRepository({
+    required this.onInitial,
+    required this.onPut,
+    this.onRemove,
+  });
 
   @override
   Future<Iterable<CoreMessage>> getMessages(ConversationSession session) async {
@@ -116,12 +126,18 @@ final class CoreMessageRepository implements CoreMessageRepositoryBase {
   }
 
   @override
-  Future<Iterable<CoreMessage>> updateMessages(Iterable<CoreMessage> messages) async {
+  Future<Iterable<CoreMessage>> updateMessages(
+    Iterable<CoreMessage> messages,
+  ) async {
     if (_cachedSession == null) {
-      throw StateError('Session not initialized. Call getMessages or saveMessages first.');
+      throw StateError(
+        'Session not initialized. Call getMessages or saveMessages first.',
+      );
     }
     for (var message in messages) {
-      final index = _messages.indexWhere((m) => m.messageId == message.messageId);
+      final index = _messages.indexWhere(
+        (m) => m.messageId == message.messageId,
+      );
       if (index != -1) {
         _messages[index] = message;
       }
diff --git a/packages/kai_engine/lib/src/models/core_message.dart b/packages/kai_engine/lib/src/models/core_message.dart
index a9128d1..38d641e 100644
--- a/packages/kai_engine/lib/src/models/core_message.dart
+++ b/packages/kai_engine/lib/src/models/core_message.dart
@@ -84,11 +84,14 @@ sealed class CoreMessage with _$CoreMessage {
     @Default(<String, dynamic>{}) Map<String, dynamic> extensions,
   }) = _CoreMessage;
 
-  bool get isDisplayable => type == CoreMessageType.user || type == CoreMessageType.ai;
-  bool get isMessageSavable => !isBackgroundContext && type != CoreMessageType.system;
+  bool get isDisplayable =>
+      type == CoreMessageType.user || type == CoreMessageType.ai;
+  bool get isMessageSavable =>
+      !isBackgroundContext && type != CoreMessageType.system;
 
   // This allow for simple use case where adapter is not require, user can directly use [CoreMessage] as main model
-  factory CoreMessage.fromJson(Map<String, dynamic> json) => _$CoreMessageFromJson(json);
+  factory CoreMessage.fromJson(Map<String, dynamic> json) =>
+      _$CoreMessageFromJson(json);
 
   CoreMessage copyWithExtensions(Map<String, dynamic> updated) {
     final newExtensions = Map<String, dynamic>.from(extensions);
diff --git a/packages/kai_engine/lib/src/prompt_engine.dart b/packages/kai_engine/lib/src/prompt_engine.dart
index 54d9a89..c3a9616 100644
--- a/packages/kai_engine/lib/src/prompt_engine.dart
+++ b/packages/kai_engine/lib/src/prompt_engine.dart
@@ -23,7 +23,8 @@ base class ContextEngineBuilder extends ContextEngine {
 /// system prompts, historical context, and user input according to a defined template structure.
 /// It supports both parallel and sequential context building strategies to optimize performance
 /// while maintaining logical ordering where required.
-abstract base class ContextEngine extends KaiPhase<ContextEngineInput, ContextEngineOutput> {
+abstract base class ContextEngine
+    extends KaiPhase<ContextEngineInput, ContextEngineOutput> {
   static ContextEngineBuilder builder(List<PromptTemplate> promptBuilder) {
     return ContextEngineBuilder(promptBuilder);
   }
@@ -41,7 +42,10 @@ abstract base class ContextEngine extends KaiPhase<ContextEngineInput, ContextEn
       inputQuery: input.inputQuery,
       providedUserMessage: input.providedUserMessage,
     );
-    return ContextEngineOutput(userMessage: result.userMessage, prompts: result.prompts);
+    return ContextEngineOutput(
+      userMessage: result.userMessage,
+      prompts: result.prompts,
+    );
   }
 
   /// Generates a complete contextual prompt ready for AI consumption.
@@ -68,7 +72,9 @@ abstract base class ContextEngine extends KaiPhase<ContextEngineInput, ContextEn
     CoreMessage? providedUserMessage,
   }) async {
     // Extract messageId for debug tracking
-    final userMessage = providedUserMessage ?? CoreMessage.user(content: inputQuery.originalQuery);
+    final userMessage =
+        providedUserMessage ??
+        CoreMessage.user(content: inputQuery.originalQuery);
     final messageId = userMessage.messageId;
 
     assert(
@@ -93,7 +99,12 @@ abstract base class ContextEngine extends KaiPhase<ContextEngineInput, ContextEn
         .toList();
 
     // Process both concurrently
-    final parallelFuture = _buildParallelWithIndex(parallelItems, source, inputQuery, messageId);
+    final parallelFuture = _buildParallelWithIndex(
+      parallelItems,
+      source,
+      inputQuery,
+      messageId,
+    );
     final sequentialFuture = _buildSequentialWithIndex(
       sequentialItems,
       source,
@@ -162,7 +173,9 @@ abstract base class ContextEngine extends KaiPhase<ContextEngineInput, ContextEn
         finalContexts.addAll(allResults[i]!);
       } else {
         // This should not happen - all template types should be handled above
-        throw StateError('Unhandled template type at index $i: ${template.runtimeType}');
+        throw StateError(
+          'Unhandled template type at index $i: ${template.runtimeType}',
+        );
       }
     }
 
@@ -183,7 +196,8 @@ abstract base class ContextEngine extends KaiPhase<ContextEngineInput, ContextEn
   ///
   /// Returns:
   /// - A list of results with their original indices for proper reordering
-  Future<List<({int index, IList<CoreMessage> result})>> _buildParallelWithIndex(
+  Future<List<({int index, IList<CoreMessage> result})>>
+  _buildParallelWithIndex(
     List<({int index, PromptTemplate builder})> items,
     IList<CoreMessage> source,
     QueryContext inputQuery,
@@ -213,7 +227,8 @@ abstract base class ContextEngine extends KaiPhase<ContextEngineInput, ContextEn
   ///
   /// Returns:
   /// - A list of results with their original indices for proper reordering
-  Future<List<({int index, IList<CoreMessage> result})>> _buildSequentialWithIndex(
+  Future<List<({int index, IList<CoreMessage> result})>>
+  _buildSequentialWithIndex(
     List<({int index, PromptTemplate builder})> items,
     IList<CoreMessage> source,
     QueryContext inputQuery,
@@ -225,7 +240,11 @@ abstract base class ContextEngine extends KaiPhase<ContextEngineInput, ContextEn
     for (final item in items) {
       final sequentialBuilder = item.builder as _BuildSequentialPromptTemplate;
       final builder = sequentialBuilder.builder;
-      final context = await builder.build(inputQuery, messageId, currentContext);
+      final context = await builder.build(
+        inputQuery,
+        messageId,
+        currentContext,
+      );
 
       currentContext = context;
       results.add((index: item.index, result: context));
@@ -248,7 +267,9 @@ final class SimpleContextEngine extends ContextEngine {
 
   @override
   List<PromptTemplate> get promptBuilder => [
-    PromptTemplate.system(systemPrompt ?? "You're kai, a useful friendly personal assistant."),
+    PromptTemplate.system(
+      systemPrompt ?? "You're kai, a useful friendly personal assistant.",
+    ),
     PromptTemplate.buildSequential(HistoryContext()),
     PromptTemplate.input(),
   ];
@@ -261,11 +282,15 @@ sealed class PromptTemplate with _$PromptTemplate {
   const PromptTemplate._();
 
   static PromptTemplate buildParallelFn(ParallelContextBuilderFn builder) {
-    return PromptTemplate.buildParallel(_FunctionParallelContextBuilder(builder));
+    return PromptTemplate.buildParallel(
+      _FunctionParallelContextBuilder(builder),
+    );
   }
 
   static PromptTemplate buildSequentialFn(SequentialContextBuilderFn builder) {
-    return PromptTemplate.buildSequential(_FunctionSequentialContextBuilder(builder));
+    return PromptTemplate.buildSequential(
+      _FunctionSequentialContextBuilder(builder),
+    );
   }
 
   /// Creates a system prompt template with fixed text content.
@@ -303,8 +328,9 @@ sealed class PromptTemplate with _$PromptTemplate {
   ///
   /// Parameters:
   /// - [builder]: The context builder that will generate the prompt content
-  const factory PromptTemplate.buildSequential(SequentialContextBuilder builder) =
-      _BuildSequentialPromptTemplate;
+  const factory PromptTemplate.buildSequential(
+    SequentialContextBuilder builder,
+  ) = _BuildSequentialPromptTemplate;
 
   /// Creates an input prompt template for the user's message.
   ///
@@ -320,7 +346,11 @@ sealed class PromptTemplate with _$PromptTemplate {
   /// - [input]: The user's raw input message
   /// - [messages]: The final list of messages in the conversation context after processing other templates include system prompts
   const factory PromptTemplate.input([
-    FutureOr<String> Function(QueryContext input, IList<CoreMessage> finalizedMessages)? revision,
+    FutureOr<String> Function(
+      QueryContext input,
+      IList<CoreMessage> finalizedMessages,
+    )?
+    revision,
   ]) = InputPromptTemplate;
 }
 
@@ -353,7 +383,8 @@ final class _FunctionParallelContextBuilder implements ParallelContextBuilder {
   }
 }
 
-final class _FunctionSequentialContextBuilder implements SequentialContextBuilder {
+final class _FunctionSequentialContextBuilder
+    implements SequentialContextBuilder {
   final SequentialContextBuilderFn _builder;
 
   const _FunctionSequentialContextBuilder(this._builder);
diff --git a/packages/kai_engine/test/src/chat_controller_base_test.dart b/packages/kai_engine/test/src/chat_controller_base_test.dart
index c824742..561fa0d 100644
--- a/packages/kai_engine/test/src/chat_controller_base_test.dart
+++ b/packages/kai_engine/test/src/chat_controller_base_test.dart
@@ -18,9 +18,11 @@ GenerationResult createTestGenerationResult({
   );
 }
 
-final class MockConversationManager extends Mock implements ConversationManager<TestMessage> {
+final class MockConversationManager extends Mock
+    implements ConversationManager<TestMessage> {
   @override
-  ConversationSession get session => const ConversationSession(id: 'test_session');
+  ConversationSession get session =>
+      const ConversationSession(id: 'test_session');
 }
 
 class MockGenerationService extends Mock implements GenerationServiceBase {}
@@ -135,7 +137,10 @@ final class TestContextEngine extends ContextEngine {
     final userMessage =
         providedUserMessage?.copyWith(content: inputQuery.processedQuery) ??
         CoreMessage.user(content: inputQuery.processedQuery);
-    return Future.value((userMessage: userMessage, prompts: IList([userMessage])));
+    return Future.value((
+      userMessage: userMessage,
+      prompts: IList([userMessage]),
+    ));
   }
 }
 
@@ -175,7 +180,10 @@ final class TestChatController extends ChatControllerBase<TestMessage> {
 
   @override
   GenerationExecuteConfig generativeConfigs(IList<CoreMessage> prompts) {
-    return const GenerationExecuteConfig(tools: [], config: {'temperature': 0.7});
+    return const GenerationExecuteConfig(
+      tools: [],
+      config: {'temperature': 0.7},
+    );
   }
 }
 
@@ -193,7 +201,9 @@ void main() {
     registerFallbackValue(CancelToken());
     registerFallbackValue(CoreMessage.user(content: 'test'));
     registerFallbackValue(
-      createTestGenerationResult(messages: IList([CoreMessage.ai(content: 'test')])),
+      createTestGenerationResult(
+        messages: IList([CoreMessage.ai(content: 'test')]),
+      ),
     );
     registerFallbackValue(<ToolSchema>[]);
     registerFallbackValue(<String, dynamic>{});
@@ -225,7 +235,9 @@ void main() {
         ),
         initialRequestMessageId: 'test',
         requestMessages: const IList.empty(),
-        result: createTestGenerationResult(messages: IList([CoreMessage.ai(content: 'test')])),
+        result: createTestGenerationResult(
+          messages: IList([CoreMessage.ai(content: 'test')]),
+        ),
         conversationManager: MockConversationManager(),
       ),
     );
@@ -291,9 +303,11 @@ void main() {
                 CoreMessage? providedUserMessage,
               }) async => (
                 userMessage:
-                    providedUserMessage ?? CoreMessage.user(messageId: 'user-1', content: 'Hello'),
+                    providedUserMessage ??
+                    CoreMessage.user(messageId: 'user-1', content: 'Hello'),
                 prompts: IList([
-                  providedUserMessage ?? CoreMessage.user(messageId: 'user-1', content: 'Hello'),
+                  providedUserMessage ??
+                      CoreMessage.user(messageId: 'user-1', content: 'Hello'),
                 ]),
               ),
         );
@@ -307,9 +321,12 @@ void main() {
             config: any(named: 'config'),
           ),
         ).thenAnswer((_) {
-          final streamController = StreamController<GenerationState<GenerationResult>>();
+          final streamController =
+              StreamController<GenerationState<GenerationResult>>();
           final result = createTestGenerationResult(
-            messages: IList([CoreMessage.ai(messageId: 'ai-1', content: 'Hi there')]),
+            messages: IList([
+              CoreMessage.ai(messageId: 'ai-1', content: 'Hi there'),
+            ]),
           );
 
           Future.microtask(() {
@@ -320,12 +337,17 @@ void main() {
           return streamController.stream;
         });
 
-        when(() => mockConversationManager.addMessages(any())).thenAnswer((invocation) async {
-          final messages = invocation.positionalArguments[0] as IList<CoreMessage>;
+        when(() => mockConversationManager.addMessages(any())).thenAnswer((
+          invocation,
+        ) async {
+          final messages =
+              invocation.positionalArguments[0] as IList<CoreMessage>;
           return messages;
         });
 
-        when(() => mockConversationManager.updateMessages(any())).thenAnswer((_) async {});
+        when(
+          () => mockConversationManager.updateMessages(any()),
+        ).thenAnswer((_) async {});
 
         controller = TestChatController(
           conversationManager: mockConversationManager,
@@ -359,9 +381,14 @@ void main() {
         fakeQueryEngine.reset();
         fakeQueryEngine.throwOnExecute = Exception('Query processing failed');
 
-        when(() => mockConversationManager.removeMessages(any())).thenAnswer((_) async {});
+        when(
+          () => mockConversationManager.removeMessages(any()),
+        ).thenAnswer((_) async {});
 
-        final result = await controller.submit('Hello', revertInputOnError: true);
+        final result = await controller.submit(
+          'Hello',
+          revertInputOnError: true,
+        );
 
         expect(result, isA<GenerationErrorState<GenerationResult>>());
 
@@ -396,9 +423,14 @@ void main() {
           testContextEngine: failingContextEngine,
         );
 
-        when(() => mockConversationManager.removeMessages(any())).thenAnswer((_) async {});
+        when(
+          () => mockConversationManager.removeMessages(any()),
+        ).thenAnswer((_) async {});
 
-        final result = await failingController.submit('Hello', revertInputOnError: true);
+        final result = await failingController.submit(
+          'Hello',
+          revertInputOnError: true,
+        );
 
         expect(result, isA<GenerationErrorState<GenerationResult>>());
 
@@ -415,7 +447,8 @@ void main() {
             config: any(named: 'config'),
           ),
         ).thenAnswer((_) {
-          final streamController = StreamController<GenerationState<GenerationResult>>();
+          final streamController =
+              StreamController<GenerationState<GenerationResult>>();
 
           Future.microtask(() {
             streamController.addError(Exception('Generation failed'));
@@ -424,9 +457,14 @@ void main() {
           return streamController.stream;
         });
 
-        when(() => mockConversationManager.removeMessages(any())).thenAnswer((_) async {});
+        when(
+          () => mockConversationManager.removeMessages(any()),
+        ).thenAnswer((_) async {});
 
-        final result = await controller.submit('Hello', revertInputOnError: true);
+        final result = await controller.submit(
+          'Hello',
+          revertInputOnError: true,
+        );
 
         expect(result, isA<GenerationErrorState<GenerationResult>>());
 
@@ -443,7 +481,8 @@ void main() {
             config: any(named: 'config'),
           ),
         ).thenAnswer((_) {
-          final streamController = StreamController<GenerationState<GenerationResult>>();
+          final streamController =
+              StreamController<GenerationState<GenerationResult>>();
 
           Future.microtask(() {
             streamController.close();
@@ -452,9 +491,14 @@ void main() {
           return streamController.stream;
         });
 
-        when(() => mockConversationManager.removeMessages(any())).thenAnswer((_) async {});
+        when(
+          () => mockConversationManager.removeMessages(any()),
+        ).thenAnswer((_) async {});
 
-        final result = await controller.submit('Hello', revertInputOnError: true);
+        final result = await controller.submit(
+          'Hello',
+          revertInputOnError: true,
+        );
 
         expect(result, isA<GenerationErrorState<GenerationResult>>());
 
@@ -462,70 +506,97 @@ void main() {
       });
 
       test('should handle post response engine failure gracefully', () async {
-        fakePostResponseEngine.throwOnExecute = Exception('Post response failed');
+        fakePostResponseEngine.throwOnExecute = Exception(
+          'Post response failed',
+        );
 
         final result = await controller.submit('Hello');
 
         expect(result, isA<GenerationErrorState<GenerationResult>>());
       });
 
-      test('should revert input on error when revertInputOnError is true', () async {
-        fakeQueryEngine.reset();
-        fakeQueryEngine.throwOnExecute = Exception('Query failed');
+      test(
+        'should revert input on error when revertInputOnError is true',
+        () async {
+          fakeQueryEngine.reset();
+          fakeQueryEngine.throwOnExecute = Exception('Query failed');
 
-        when(() => mockConversationManager.removeMessages(any())).thenAnswer((_) async {});
+          when(
+            () => mockConversationManager.removeMessages(any()),
+          ).thenAnswer((_) async {});
 
-        final result = await controller.submit('Hello', revertInputOnError: true);
+          final result = await controller.submit(
+            'Hello',
+            revertInputOnError: true,
+          );
 
-        expect(result, isA<GenerationErrorState<GenerationResult>>());
+          expect(result, isA<GenerationErrorState<GenerationResult>>());
 
-        verify(() => mockConversationManager.removeMessages(any())).called(1);
-      });
+          verify(() => mockConversationManager.removeMessages(any())).called(1);
+        },
+      );
 
-      test('should revert persisted message on error when revertInputOnError is true', () async {
-        when(
-          () => mockGenerationService.stream(
-            any(),
-            cancelToken: any(named: 'cancelToken'),
-            tools: any(named: 'tools'),
-            toolingConfig: any(named: 'toolingConfig'),
-            config: any(named: 'config'),
-          ),
-        ).thenThrow(Exception('Generation failed'));
+      test(
+        'should revert persisted message on error when revertInputOnError is true',
+        () async {
+          when(
+            () => mockGenerationService.stream(
+              any(),
+              cancelToken: any(named: 'cancelToken'),
+              tools: any(named: 'tools'),
+              toolingConfig: any(named: 'toolingConfig'),
+              config: any(named: 'config'),
+            ),
+          ).thenThrow(Exception('Generation failed'));
 
-        when(() => mockConversationManager.removeMessages(any())).thenAnswer((_) async {});
+          when(
+            () => mockConversationManager.removeMessages(any()),
+          ).thenAnswer((_) async {});
 
-        final result = await controller.submit('Hello', revertInputOnError: true);
+          final result = await controller.submit(
+            'Hello',
+            revertInputOnError: true,
+          );
 
-        expect(result, isA<GenerationErrorState<GenerationResult>>());
+          expect(result, isA<GenerationErrorState<GenerationResult>>());
 
-        verify(() => mockConversationManager.removeMessages(any())).called(1);
-      });
+          verify(() => mockConversationManager.removeMessages(any())).called(1);
+        },
+      );
 
-      test('should not revert input on error when revertInputOnError is false', () async {
-        fakeQueryEngine.reset();
-        fakeQueryEngine.throwOnExecute = Exception('Query failed');
+      test(
+        'should not revert input on error when revertInputOnError is false',
+        () async {
+          fakeQueryEngine.reset();
+          fakeQueryEngine.throwOnExecute = Exception('Query failed');
 
-        final result = await controller.submit('Hello', revertInputOnError: false);
+          final result = await controller.submit(
+            'Hello',
+            revertInputOnError: false,
+          );
 
-        expect(result, isA<GenerationErrorState<GenerationResult>>());
+          expect(result, isA<GenerationErrorState<GenerationResult>>());
 
-        verifyNever(() => mockConversationManager.removeMessages(any()));
-      });
+          verifyNever(() => mockConversationManager.removeMessages(any()));
+        },
+      );
 
-      test('should pass custom generative configs to generation service', () async {
-        await controller.submit('Hello');
+      test(
+        'should pass custom generative configs to generation service',
+        () async {
+          await controller.submit('Hello');
 
-        verify(
-          () => mockGenerationService.stream(
-            any(),
-            cancelToken: any(named: 'cancelToken'),
-            tools: [],
-            toolingConfig: any(named: 'toolingConfig'),
-            config: {'temperature': 0.7},
-          ),
-        ).called(1);
-      });
+          verify(
+            () => mockGenerationService.stream(
+              any(),
+              cancelToken: any(named: 'cancelToken'),
+              tools: [],
+              toolingConfig: any(named: 'toolingConfig'),
+              config: {'temperature': 0.7},
+            ),
+          ).called(1);
+        },
+      );
 
       test(
         'should handle stream error states and revert input when revertInputOnError is true',
@@ -539,11 +610,14 @@ void main() {
               config: any(named: 'config'),
             ),
           ).thenAnswer((_) {
-            final streamController = StreamController<GenerationState<GenerationResult>>();
+            final streamController =
+                StreamController<GenerationState<GenerationResult>>();
 
             Future.microtask(() {
               streamController.add(
-                GenerationState.error(KaiException.exception('Stream generation failed')),
+                GenerationState.error(
+                  KaiException.exception('Stream generation failed'),
+                ),
               );
               streamController.close();
             });
@@ -551,9 +625,14 @@ void main() {
             return streamController.stream;
           });
 
-          when(() => mockConversationManager.removeMessages(any())).thenAnswer((_) async {});
+          when(
+            () => mockConversationManager.removeMessages(any()),
+          ).thenAnswer((_) async {});
 
-          final result = await controller.submit('Hello', revertInputOnError: true);
+          final result = await controller.submit(
+            'Hello',
+            revertInputOnError: true,
+          );
 
           expect(result, isA<GenerationErrorState<GenerationResult>>());
 
@@ -573,11 +652,14 @@ void main() {
               config: any(named: 'config'),
             ),
           ).thenAnswer((_) {
-            final streamController = StreamController<GenerationState<GenerationResult>>();
+            final streamController =
+                StreamController<GenerationState<GenerationResult>>();
 
             Future.microtask(() {
               streamController.add(
-                GenerationState.error(KaiException.exception('Stream generation failed')),
+                GenerationState.error(
+                  KaiException.exception('Stream generation failed'),
+                ),
               );
               streamController.close();
             });
@@ -585,7 +667,10 @@ void main() {
             return streamController.stream;
           });
 
-          final result = await controller.submit('Hello', revertInputOnError: false);
+          final result = await controller.submit(
+            'Hello',
+            revertInputOnError: false,
+          );
 
           expect(result, isA<GenerationErrorState<GenerationResult>>());
 
@@ -595,7 +680,9 @@ void main() {
 
       test('should emit states through stream during submission', () async {
         final states = <GenerationState<GenerationResult>>[];
-        final subscription = controller.generationStateStream.listen(states.add);
+        final subscription = controller.generationStateStream.listen(
+          states.add,
+        );
 
         await controller.submit('Hello');
 
@@ -606,96 +693,112 @@ void main() {
         await subscription.cancel();
       });
 
-      test('should not have duplicate user input in prompt sent to stream', () async {
-        IList<CoreMessage>? capturedPrompts;
-
-        final localMockConversationManager = MockConversationManager();
-        final localMockGenerationService = MockGenerationService();
-        final localFakeQueryEngine = FakeQueryEngine();
-        final localFakePostResponseEngine = FakePostResponseEngine();
-
-        when(() => localMockConversationManager.getMessages()).thenAnswer(
-          (_) async => IList([
-            CoreMessage.user(content: 'Previous message'),
-            CoreMessage.ai(content: 'Previous response'),
-          ]),
-        );
+      test(
+        'should not have duplicate user input in prompt sent to stream',
+        () async {
+          IList<CoreMessage>? capturedPrompts;
+
+          final localMockConversationManager = MockConversationManager();
+          final localMockGenerationService = MockGenerationService();
+          final localFakeQueryEngine = FakeQueryEngine();
+          final localFakePostResponseEngine = FakePostResponseEngine();
+
+          when(() => localMockConversationManager.getMessages()).thenAnswer(
+            (_) async => IList([
+              CoreMessage.user(content: 'Previous message'),
+              CoreMessage.ai(content: 'Previous response'),
+            ]),
+          );
 
-        when(() => localMockConversationManager.addMessages(any())).thenAnswer((invocation) async {
-          final messages = invocation.positionalArguments[0] as IList<CoreMessage>;
-          return messages;
-        });
+          when(
+            () => localMockConversationManager.addMessages(any()),
+          ).thenAnswer((invocation) async {
+            final messages =
+                invocation.positionalArguments[0] as IList<CoreMessage>;
+            return messages;
+          });
 
-        when(() => localMockConversationManager.updateMessages(any())).thenAnswer((_) async {});
+          when(
+            () => localMockConversationManager.updateMessages(any()),
+          ).thenAnswer((_) async {});
 
-        localFakeQueryEngine.onExecute = (input) async => QueryContext(
-          originalQuery: input.rawInput,
-          processedQuery: input.rawInput,
-          session: const ConversationSession(id: 'test_session'),
-        );
+          localFakeQueryEngine.onExecute = (input) async => QueryContext(
+            originalQuery: input.rawInput,
+            processedQuery: input.rawInput,
+            session: const ConversationSession(id: 'test_session'),
+          );
 
-        when(
-          () => localMockGenerationService.stream(
-            captureAny(that: isA<IList<CoreMessage>>()),
-            cancelToken: any(named: 'cancelToken'),
-            tools: any(named: 'tools'),
-            toolingConfig: any(named: 'toolingConfig'),
-            config: any(named: 'config'),
-          ),
-        ).thenAnswer((invocation) {
-          capturedPrompts = invocation.positionalArguments[0] as IList<CoreMessage>;
+          when(
+            () => localMockGenerationService.stream(
+              captureAny(that: isA<IList<CoreMessage>>()),
+              cancelToken: any(named: 'cancelToken'),
+              tools: any(named: 'tools'),
+              toolingConfig: any(named: 'toolingConfig'),
+              config: any(named: 'config'),
+            ),
+          ).thenAnswer((invocation) {
+            capturedPrompts =
+                invocation.positionalArguments[0] as IList<CoreMessage>;
+
+            final streamController =
+                StreamController<GenerationState<GenerationResult>>();
+            final result = createTestGenerationResult(
+              messages: IList([
+                CoreMessage.ai(messageId: 'ai-1', content: 'Hi there'),
+              ]),
+            );
 
-          final streamController = StreamController<GenerationState<GenerationResult>>();
-          final result = createTestGenerationResult(
-            messages: IList([CoreMessage.ai(messageId: 'ai-1', content: 'Hi there')]),
-          );
+            Future.microtask(() {
+              streamController.add(GenerationState.complete(result));
+              streamController.close();
+            });
 
-          Future.microtask(() {
-            streamController.add(GenerationState.complete(result));
-            streamController.close();
+            return streamController.stream;
           });
 
-          return streamController.stream;
-        });
-
-        final localTestContextEngine = TestContextEngine(
-          generateFunction:
-              ({
-                required IList<CoreMessage> source,
-                required QueryContext inputQuery,
-                CoreMessage? providedUserMessage,
-              }) async => (
-                userMessage:
-                    providedUserMessage ?? CoreMessage.user(messageId: 'user-1', content: 'Hello'),
-                prompts: IList([
-                  CoreMessage.user(content: 'Previous message'),
-                  CoreMessage.ai(content: 'Previous response'),
-                  providedUserMessage ?? CoreMessage.user(messageId: 'user-1', content: 'Hello'),
-                ]),
-              ),
-        );
+          final localTestContextEngine = TestContextEngine(
+            generateFunction:
+                ({
+                  required IList<CoreMessage> source,
+                  required QueryContext inputQuery,
+                  CoreMessage? providedUserMessage,
+                }) async => (
+                  userMessage:
+                      providedUserMessage ??
+                      CoreMessage.user(messageId: 'user-1', content: 'Hello'),
+                  prompts: IList([
+                    CoreMessage.user(content: 'Previous message'),
+                    CoreMessage.ai(content: 'Previous response'),
+                    providedUserMessage ??
+                        CoreMessage.user(messageId: 'user-1', content: 'Hello'),
+                  ]),
+                ),
+          );
 
-        final localController = TestChatController(
-          conversationManager: localMockConversationManager,
-          generationService: localMockGenerationService,
-          queryEngine: localFakeQueryEngine,
-          postResponseEngine: localFakePostResponseEngine,
-          testContextEngine: localTestContextEngine,
-        );
+          final localController = TestChatController(
+            conversationManager: localMockConversationManager,
+            generationService: localMockGenerationService,
+            queryEngine: localFakeQueryEngine,
+            postResponseEngine: localFakePostResponseEngine,
+            testContextEngine: localTestContextEngine,
+          );
 
-        await localController.submit('Hello');
+          await localController.submit('Hello');
 
-        expect(capturedPrompts, isNotNull);
-        expect(capturedPrompts!.length, equals(3));
+          expect(capturedPrompts, isNotNull);
+          expect(capturedPrompts!.length, equals(3));
 
-        expect(capturedPrompts![0].content, equals('Previous message'));
-        expect(capturedPrompts![1].content, equals('Previous response'));
-        expect(capturedPrompts![2].content, equals('Hello'));
+          expect(capturedPrompts![0].content, equals('Previous message'));
+          expect(capturedPrompts![1].content, equals('Previous response'));
+          expect(capturedPrompts![2].content, equals('Hello'));
 
-        final helloMessages = capturedPrompts!.where((msg) => msg.content == 'Hello').toList();
-        expect(helloMessages.length, equals(1));
-        expect(helloMessages[0].messageId, isNotEmpty);
-      });
+          final helloMessages = capturedPrompts!
+              .where((msg) => msg.content == 'Hello')
+              .toList();
+          expect(helloMessages.length, equals(1));
+          expect(helloMessages[0].messageId, isNotEmpty);
+        },
+      );
     });
 
     group('streams', () {
@@ -726,9 +829,12 @@ void main() {
           config: any(named: 'config'),
         ),
       ).thenAnswer((_) {
-        final streamController = StreamController<GenerationState<GenerationResult>>();
+        final streamController =
+            StreamController<GenerationState<GenerationResult>>();
         final result = createTestGenerationResult(
-          messages: IList([CoreMessage.ai(messageId: 'ai-1', content: 'Response')]),
+          messages: IList([
+            CoreMessage.ai(messageId: 'ai-1', content: 'Response'),
+          ]),
         );
 
         Future.microtask(() {
@@ -851,16 +957,21 @@ void main() {
         controller.dispose();
       });
 
-      test('should use InMemoryMessageRepository when repository not provided', () async {
-        final controller = await KaiChatController.create(generationService: mockGenerationService);
+      test(
+        'should use InMemoryMessageRepository when repository not provided',
+        () async {
+          final controller = await KaiChatController.create(
+            generationService: mockGenerationService,
+          );
 
-        expect(controller, isNotNull);
+          expect(controller, isNotNull);
 
-        final messages = await controller.getAllMessages();
-        expect(messages, isA<IList<CoreMessage>>());
+          final messages = await controller.getAllMessages();
+          expect(messages, isA<IList<CoreMessage>>());
 
-        controller.dispose();
-      });
+          controller.dispose();
+        },
+      );
 
       test('should create controller with all parameters', () async {
         final queryEngine = FakeQueryEngine();
@@ -900,48 +1011,58 @@ void main() {
         controller.dispose();
       });
 
-      test('should create multiple controllers with different sessionIds', () async {
-        final controller1 = await KaiChatController.create(
-          sessionId: 'session-1',
-          generationService: mockGenerationService,
-        );
+      test(
+        'should create multiple controllers with different sessionIds',
+        () async {
+          final controller1 = await KaiChatController.create(
+            sessionId: 'session-1',
+            generationService: mockGenerationService,
+          );
 
-        final controller2 = await KaiChatController.create(
-          sessionId: 'session-2',
-          generationService: mockGenerationService,
-        );
+          final controller2 = await KaiChatController.create(
+            sessionId: 'session-2',
+            generationService: mockGenerationService,
+          );
 
-        expect(controller1, isNotNull);
-        expect(controller2, isNotNull);
-        expect(controller1, isNot(equals(controller2)));
+          expect(controller1, isNotNull);
+          expect(controller2, isNotNull);
+          expect(controller1, isNot(equals(controller2)));
 
-        controller1.dispose();
-        controller2.dispose();
-      });
+          controller1.dispose();
+          controller2.dispose();
+        },
+      );
     });
 
     group('submit', () {
-      test('should handle successful submission with default configuration', () async {
-        final controller = await KaiChatController.create(
-          generationService: mockGenerationService,
-          systemPrompt: 'You are helpful',
-        );
+      test(
+        'should handle successful submission with default configuration',
+        () async {
+          final controller = await KaiChatController.create(
+            generationService: mockGenerationService,
+            systemPrompt: 'You are helpful',
+          );
 
-        final result = await controller.submit('Hello, how are you?');
+          final result = await controller.submit('Hello, how are you?');
 
-        expect(result, isA<GenerationCompleteState<GenerationResult>>());
+          expect(result, isA<GenerationCompleteState<GenerationResult>>());
 
-        final messages = await controller.getAllMessages();
-        expect(messages.length, greaterThan(0));
+          final messages = await controller.getAllMessages();
+          expect(messages.length, greaterThan(0));
 
-        controller.dispose();
-      });
+          controller.dispose();
+        },
+      );
 
       test('should emit loading and complete states', () async {
-        final controller = await KaiChatController.create(generationService: mockGenerationService);
+        final controller = await KaiChatController.create(
+          generationService: mockGenerationService,
+        );
 
         final states = <GenerationState<GenerationResult>>[];
-        final subscription = controller.generationStateStream.listen(states.add);
+        final subscription = controller.generationStateStream.listen(
+          states.add,
+        );
 
         await controller.submit('Test message');
 
@@ -953,38 +1074,49 @@ void main() {
         controller.dispose();
       });
 
-      test('should handle error and revert input when revertInputOnError is true', () async {
-        when(
-          () => mockGenerationService.stream(
-            any(),
-            cancelToken: any(named: 'cancelToken'),
-            tools: any(named: 'tools'),
-            toolingConfig: any(named: 'toolingConfig'),
-            config: any(named: 'config'),
-          ),
-        ).thenAnswer((_) {
-          final streamController = StreamController<GenerationState<GenerationResult>>();
+      test(
+        'should handle error and revert input when revertInputOnError is true',
+        () async {
+          when(
+            () => mockGenerationService.stream(
+              any(),
+              cancelToken: any(named: 'cancelToken'),
+              tools: any(named: 'tools'),
+              toolingConfig: any(named: 'toolingConfig'),
+              config: any(named: 'config'),
+            ),
+          ).thenAnswer((_) {
+            final streamController =
+                StreamController<GenerationState<GenerationResult>>();
 
-          Future.microtask(() {
-            streamController.addError(Exception('Generation failed'));
-          });
+            Future.microtask(() {
+              streamController.addError(Exception('Generation failed'));
+            });
 
-          return streamController.stream;
-        });
+            return streamController.stream;
+          });
 
-        final controller = await KaiChatController.create(generationService: mockGenerationService);
+          final controller = await KaiChatController.create(
+            generationService: mockGenerationService,
+          );
 
-        final result = await controller.submit('Test message', revertInputOnError: true);
+          final result = await controller.submit(
+            'Test message',
+            revertInputOnError: true,
+          );
 
-        expect(result, isA<GenerationErrorState<GenerationResult>>());
+          expect(result, isA<GenerationErrorState<GenerationResult>>());
 
-        controller.dispose();
-      });
+          controller.dispose();
+        },
+      );
     });
 
     group('messagesStream', () {
       test('should emit messages through stream', () async {
-        final controller = await KaiChatController.create(generationService: mockGenerationService);
+        final controller = await KaiChatController.create(
+          generationService: mockGenerationService,
+        );
 
         final messagesList = <IList<CoreMessage>>[];
         final subscription = controller.messagesStream.listen(messagesList.add);
diff --git a/packages/kai_engine/test/src/message_repository_base_test.dart b/packages/kai_engine/test/src/message_repository_base_test.dart
index ebf4967..eb0f46d 100644
--- a/packages/kai_engine/test/src/message_repository_base_test.dart
+++ b/packages/kai_engine/test/src/message_repository_base_test.dart
@@ -42,7 +42,10 @@ void main() {
 
     group('constructor', () {
       test('should create instance with required callbacks', () {
-        final repo = CoreMessageRepository(onInitial: (s) async => [], onPut: (s, msgs) async {});
+        final repo = CoreMessageRepository(
+          onInitial: (s) async => [],
+          onPut: (s, msgs) async {},
+        );
 
         expect(repo, isA<CoreMessageRepository>());
       });
@@ -59,21 +62,24 @@ void main() {
     });
 
     group('getMessages', () {
-      test('should call onInitial callback with session and populate internal state', () async {
-        var capturedSession;
-        final repo = CoreMessageRepository(
-          onInitial: (s) {
-            capturedSession = s;
-            return Future.value([testMessage1]);
-          },
-          onPut: (s, msgs) async {},
-        );
-
-        final result = await repo.getMessages(session);
-
-        expect(capturedSession, equals(session));
-        expect(result, equals([testMessage1]));
-      });
+      test(
+        'should call onInitial callback with session and populate internal state',
+        () async {
+          var capturedSession;
+          final repo = CoreMessageRepository(
+            onInitial: (s) {
+              capturedSession = s;
+              return Future.value([testMessage1]);
+            },
+            onPut: (s, msgs) async {},
+          );
+
+          final result = await repo.getMessages(session);
+
+          expect(capturedSession, equals(session));
+          expect(result, equals([testMessage1]));
+        },
+      );
 
       test('should sort messages by timestamp on initial load', () async {
         final now = DateTime.now();
@@ -111,23 +117,26 @@ void main() {
         expect(result, equals([testMessage1, testMessage2]));
       });
 
-      test('should cache state and not reload on subsequent getMessages calls', () async {
-        var callCount = 0;
-        final repo = CoreMessageRepository(
-          onInitial: (s) async {
-            callCount++;
-            return callCount == 1 ? [testMessage1] : [testMessage2];
-          },
-          onPut: (s, msgs) async {},
-        );
-
-        final result1 = await repo.getMessages(session);
-        final result2 = await repo.getMessages(session);
-
-        expect(callCount, equals(1)); // Only called once
-        expect(result1, equals([testMessage1]));
-        expect(result2, equals([testMessage1])); // Returns cached state
-      });
+      test(
+        'should cache state and not reload on subsequent getMessages calls',
+        () async {
+          var callCount = 0;
+          final repo = CoreMessageRepository(
+            onInitial: (s) async {
+              callCount++;
+              return callCount == 1 ? [testMessage1] : [testMessage2];
+            },
+            onPut: (s, msgs) async {},
+          );
+
+          final result1 = await repo.getMessages(session);
+          final result2 = await repo.getMessages(session);
+
+          expect(callCount, equals(1)); // Only called once
+          expect(result1, equals([testMessage1]));
+          expect(result2, equals([testMessage1])); // Returns cached state
+        },
+      );
 
       test('should cache session for later use', () async {
         final repo = CoreMessageRepository(
@@ -145,45 +154,60 @@ void main() {
     });
 
     group('saveMessages', () {
-      test('should call onPut callback with session and messages for persistence', () async {
-        var capturedSession;
-        var capturedMessages;
+      test(
+        'should call onPut callback with session and messages for persistence',
+        () async {
+          var capturedSession;
+          var capturedMessages;
+          final repo = CoreMessageRepository(
+            onInitial: (s) async => [],
+            onPut: (s, msgs) {
+              capturedSession = s;
+              capturedMessages = msgs;
+              return Future.value(msgs);
+            },
+          );
+
+          await repo.saveMessages(session: session, messages: [testMessage1]);
+
+          expect(capturedSession, equals(session));
+          expect(capturedMessages, equals([testMessage1]));
+        },
+      );
+
+      test('should add messages to internal state', () async {
         final repo = CoreMessageRepository(
           onInitial: (s) async => [],
-          onPut: (s, msgs) {
-            capturedSession = s;
-            capturedMessages = msgs;
-            return Future.value(msgs);
-          },
+          onPut: (s, msgs) async {},
         );
 
-        await repo.saveMessages(session: session, messages: [testMessage1]);
-
-        expect(capturedSession, equals(session));
-        expect(capturedMessages, equals([testMessage1]));
-      });
-
-      test('should add messages to internal state', () async {
-        final repo = CoreMessageRepository(onInitial: (s) async => [], onPut: (s, msgs) async {});
-
         await repo.saveMessages(session: session, messages: [testMessage1]);
         final result = await repo.getMessages(session);
 
         expect(result, equals([testMessage1]));
       });
 
-      test('should accumulate messages in internal state across multiple saves', () async {
-        final repo = CoreMessageRepository(onInitial: (s) async => [], onPut: (s, msgs) async {});
+      test(
+        'should accumulate messages in internal state across multiple saves',
+        () async {
+          final repo = CoreMessageRepository(
+            onInitial: (s) async => [],
+            onPut: (s, msgs) async {},
+          );
 
-        await repo.saveMessages(session: session, messages: [testMessage1]);
-        await repo.saveMessages(session: session, messages: [testMessage2]);
-        final result = await repo.getMessages(session);
+          await repo.saveMessages(session: session, messages: [testMessage1]);
+          await repo.saveMessages(session: session, messages: [testMessage2]);
+          final result = await repo.getMessages(session);
 
-        expect(result, equals([testMessage1, testMessage2]));
-      });
+          expect(result, equals([testMessage1, testMessage2]));
+        },
+      );
 
       test('should cache session for later use', () async {
-        final repo = CoreMessageRepository(onInitial: (s) async => [], onPut: (s, msgs) async {});
+        final repo = CoreMessageRepository(
+          onInitial: (s) async => [],
+          onPut: (s, msgs) async {},
+        );
 
         await repo.saveMessages(session: session, messages: [testMessage1]);
 
@@ -195,24 +219,27 @@ void main() {
     });
 
     group('updateMessages', () {
-      test('should call onPut callback with cached session and messages for persistence', () async {
-        var capturedSession;
-        var capturedMessages;
-        final repo = CoreMessageRepository(
-          onInitial: (s) async => [testMessage1],
-          onPut: (s, msgs) {
-            capturedSession = s;
-            capturedMessages = msgs;
-            return Future.value(msgs);
-          },
-        );
-
-        await repo.getMessages(session);
-        await repo.updateMessages([testMessage1]);
-
-        expect(capturedSession, equals(session));
-        expect(capturedMessages, equals([testMessage1]));
-      });
+      test(
+        'should call onPut callback with cached session and messages for persistence',
+        () async {
+          var capturedSession;
+          var capturedMessages;
+          final repo = CoreMessageRepository(
+            onInitial: (s) async => [testMessage1],
+            onPut: (s, msgs) {
+              capturedSession = s;
+              capturedMessages = msgs;
+              return Future.value(msgs);
+            },
+          );
+
+          await repo.getMessages(session);
+          await repo.updateMessages([testMessage1]);
+
+          expect(capturedSession, equals(session));
+          expect(capturedMessages, equals([testMessage1]));
+        },
+      );
 
       test('should update messages in internal state', () async {
         final repo = CoreMessageRepository(
@@ -229,7 +256,10 @@ void main() {
       });
 
       test('should throw StateError when session not initialized', () async {
-        final repo = CoreMessageRepository(onInitial: (s) async => [], onPut: (s, msgs) async {});
+        final repo = CoreMessageRepository(
+          onInitial: (s) async => [],
+          onPut: (s, msgs) async {},
+        );
 
         expect(
           () => repo.updateMessages([testMessage1]),
@@ -243,16 +273,22 @@ void main() {
         );
       });
 
-      test('should use session from saveMessages if getMessages not called', () async {
-        final repo = CoreMessageRepository(onInitial: (s) async => [], onPut: (s, msgs) async {});
+      test(
+        'should use session from saveMessages if getMessages not called',
+        () async {
+          final repo = CoreMessageRepository(
+            onInitial: (s) async => [],
+            onPut: (s, msgs) async {},
+          );
 
-        await repo.saveMessages(session: session, messages: [testMessage1]);
-        final updated = testMessage1.copyWith(content: 'Updated');
-        await repo.updateMessages([updated]);
-        final result = await repo.getMessages(session);
+          await repo.saveMessages(session: session, messages: [testMessage1]);
+          final updated = testMessage1.copyWith(content: 'Updated');
+          await repo.updateMessages([updated]);
+          final result = await repo.getMessages(session);
 
-        expect(result.first.content, equals('Updated'));
-      });
+          expect(result.first.content, equals('Updated'));
+        },
+      );
 
       test('should only update existing messages, not add new ones', () async {
         final repo = CoreMessageRepository(
@@ -269,22 +305,25 @@ void main() {
     });
 
     group('removeMessages', () {
-      test('should call onRemove callback when provided for persistence', () async {
-        var capturedMessages;
-        final repo = CoreMessageRepository(
-          onInitial: (s) async => [testMessage1, testMessage2],
-          onPut: (s, msgs) async {},
-          onRemove: (msgs) {
-            capturedMessages = msgs;
-            return Future.value();
-          },
-        );
-
-        await repo.getMessages(session);
-        await repo.removeMessages([testMessage1]);
-
-        expect(capturedMessages, equals([testMessage1]));
-      });
+      test(
+        'should call onRemove callback when provided for persistence',
+        () async {
+          var capturedMessages;
+          final repo = CoreMessageRepository(
+            onInitial: (s) async => [testMessage1, testMessage2],
+            onPut: (s, msgs) async {},
+            onRemove: (msgs) {
+              capturedMessages = msgs;
+              return Future.value();
+            },
+          );
+
+          await repo.getMessages(session);
+          await repo.removeMessages([testMessage1]);
+
+          expect(capturedMessages, equals([testMessage1]));
+        },
+      );
 
       test('should remove messages from internal state', () async {
         final repo = CoreMessageRepository(
@@ -324,53 +363,64 @@ void main() {
     });
 
     group('integration scenarios', () {
-      test('should handle complete message lifecycle with persistence callbacks', () async {
-        final persistedMessages = <CoreMessage>[];
-
-        final repo = CoreMessageRepository(
-          onInitial: (s) async => persistedMessages.toList(),
-          onPut: (s, msgs) async {
-            // Simulate upsert in persistence layer
-            for (var msg in msgs) {
-              final index = persistedMessages.indexWhere((m) => m.messageId == msg.messageId);
-              if (index != -1) {
-                persistedMessages[index] = msg;
-              } else {
-                persistedMessages.add(msg);
+      test(
+        'should handle complete message lifecycle with persistence callbacks',
+        () async {
+          final persistedMessages = <CoreMessage>[];
+
+          final repo = CoreMessageRepository(
+            onInitial: (s) async => persistedMessages.toList(),
+            onPut: (s, msgs) async {
+              // Simulate upsert in persistence layer
+              for (var msg in msgs) {
+                final index = persistedMessages.indexWhere(
+                  (m) => m.messageId == msg.messageId,
+                );
+                if (index != -1) {
+                  persistedMessages[index] = msg;
+                } else {
+                  persistedMessages.add(msg);
+                }
               }
-            }
-          },
-          onRemove: (msgs) async {
-            persistedMessages.removeWhere((m) => msgs.contains(m));
-          },
-        );
-
-        // Initial load
-        var current = await repo.getMessages(session);
-        expect(current, isEmpty);
-
-        // Save messages
-        await repo.saveMessages(session: session, messages: [testMessage1, testMessage2]);
-        current = await repo.getMessages(session);
-        expect(current.length, equals(2));
-
-        // Update message
-        final updated = testMessage1.copyWith(content: 'Updated content');
-        await repo.updateMessages([updated]);
-        current = await repo.getMessages(session);
-        expect(current.any((m) => m.content == 'Updated content'), isTrue);
-        expect(current.length, equals(2)); // Still 2 messages, one updated
-
-        // Remove message
-        await repo.removeMessages([testMessage2]);
-        current = await repo.getMessages(session);
-        expect(current.length, equals(1));
-        expect(current.first.messageId, equals(testMessage1.messageId));
-        expect(current.first.content, equals('Updated content'));
-      });
+            },
+            onRemove: (msgs) async {
+              persistedMessages.removeWhere((m) => msgs.contains(m));
+            },
+          );
+
+          // Initial load
+          var current = await repo.getMessages(session);
+          expect(current, isEmpty);
+
+          // Save messages
+          await repo.saveMessages(
+            session: session,
+            messages: [testMessage1, testMessage2],
+          );
+          current = await repo.getMessages(session);
+          expect(current.length, equals(2));
+
+          // Update message
+          final updated = testMessage1.copyWith(content: 'Updated content');
+          await repo.updateMessages([updated]);
+          current = await repo.getMessages(session);
+          expect(current.any((m) => m.content == 'Updated content'), isTrue);
+          expect(current.length, equals(2)); // Still 2 messages, one updated
+
+          // Remove message
+          await repo.removeMessages([testMessage2]);
+          current = await repo.getMessages(session);
+          expect(current.length, equals(1));
+          expect(current.first.messageId, equals(testMessage1.messageId));
+          expect(current.first.content, equals('Updated content'));
+        },
+      );
 
       test('should handle append-only pattern without onRemove', () async {
-        final repo = CoreMessageRepository(onInitial: (s) async => [], onPut: (s, msgs) async {});
+        final repo = CoreMessageRepository(
+          onInitial: (s) async => [],
+          onPut: (s, msgs) async {},
+        );
 
         await repo.saveMessages(session: session, messages: [testMessage1]);
         await repo.saveMessages(session: session, messages: [testMessage2]);
@@ -385,25 +435,28 @@ void main() {
         expect(afterRemove.first.messageId, equals(testMessage2.messageId));
       });
 
-      test('should maintain state independently of persistence layer', () async {
-        var persistenceCallCount = 0;
-
-        final repo = CoreMessageRepository(
-          onInitial: (s) async {
-            persistenceCallCount++;
-            return [];
-          },
-          onPut: (s, msgs) async {},
-        );
-
-        // First load calls persistence
-        await repo.getMessages(session);
-        expect(persistenceCallCount, equals(1));
-
-        // Subsequent reads use internal state, don't call persistence
-        await repo.getMessages(session);
-        expect(persistenceCallCount, equals(1)); // Still 1, not 2
-      });
+      test(
+        'should maintain state independently of persistence layer',
+        () async {
+          var persistenceCallCount = 0;
+
+          final repo = CoreMessageRepository(
+            onInitial: (s) async {
+              persistenceCallCount++;
+              return [];
+            },
+            onPut: (s, msgs) async {},
+          );
+
+          // First load calls persistence
+          await repo.getMessages(session);
+          expect(persistenceCallCount, equals(1));
+
+          // Subsequent reads use internal state, don't call persistence
+          await repo.getMessages(session);
+          expect(persistenceCallCount, equals(1)); // Still 1, not 2
+        },
+      );
     });
   });
 }

From f4b0ff8a59f0d55cf8b4ea79dab9b3143823e287 Mon Sep 17 00:00:00 2001
From: pckimlong <pochkimlong88@gmail.com>
Date: Sat, 11 Apr 2026 16:48:21 +0700
Subject: [PATCH 10/10] optional format CI

---
 .github/workflows/ci.yml                      |  4 +--
 .../lib/src/adapters.dart                     | 28 ++++++++++++++-----
 2 files changed, 23 insertions(+), 9 deletions(-)

diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
index de9c14c..5537efe 100644
--- a/.github/workflows/ci.yml
+++ b/.github/workflows/ci.yml
@@ -28,10 +28,10 @@ jobs:
         run: echo "$HOME/.pub-cache/bin" >> $GITHUB_PATH
       - name: Install dependencies
         run: melos bootstrap
-      - name: Verify formatting
+      - name: Format code
         run: |
           cd packages/${{ matrix.package }}
-          dart format --output=none --set-exit-if-changed .
+          dart format .
       - name: Analyze project source
         run: |
           cd packages/${{ matrix.package }}
diff --git a/packages/kai_engine_firebase_ai/lib/src/adapters.dart b/packages/kai_engine_firebase_ai/lib/src/adapters.dart
index 3dccc74..75e03e4 100644
--- a/packages/kai_engine_firebase_ai/lib/src/adapters.dart
+++ b/packages/kai_engine_firebase_ai/lib/src/adapters.dart
@@ -5,7 +5,8 @@ import 'package:firebase_ai/src/content.dart'; // Access parsePart function
 import 'package:kai_engine/kai_engine.dart';
 
 /// Adapter for Firebase AI content to Core Message
-class FirebaseAiContentAdapter implements GenerativeMessageAdapterBase<Content> {
+class FirebaseAiContentAdapter
+    implements GenerativeMessageAdapterBase<Content> {
   const FirebaseAiContentAdapter();
 
   @override
@@ -32,7 +33,8 @@ class FirebaseAiContentAdapter implements GenerativeMessageAdapterBase<Content>
           final reconstructedParts = <Part>[];
           final originalJson = originalContentJson as Map<String, dynamic>;
           final originalParts =
-              (originalJson['parts'] as List?)?.cast<Map<String, dynamic>>() ?? [];
+              (originalJson['parts'] as List?)?.cast<Map<String, dynamic>>() ??
+              [];
 
           for (int i = 0; i < parsedContent.parts.length; i++) {
             final parsedPart = parsedContent.parts[i];
@@ -40,10 +42,13 @@ class FirebaseAiContentAdapter implements GenerativeMessageAdapterBase<Content>
             // If this was parsed as UnknownPart but we have original JSON, try to reconstruct manually
             // Also handle InlineDataPart with null willContinue field
             if ((parsedPart is UnknownPart ||
-                    (parsedPart is InlineDataPart && parsedPart.willContinue == null)) &&
+                    (parsedPart is InlineDataPart &&
+                        parsedPart.willContinue == null)) &&
                 i < originalParts.length) {
               final originalPartJson = originalParts[i];
-              final reconstructedPart = _reconstructPartFromJson(originalPartJson);
+              final reconstructedPart = _reconstructPartFromJson(
+                originalPartJson,
+              );
               reconstructedParts.add(reconstructedPart ?? parsedPart);
             } else {
               reconstructedParts.add(parsedPart);
@@ -107,7 +112,11 @@ class FirebaseAiContentAdapter implements GenerativeMessageAdapterBase<Content>
       extensions['_timestamp'] = timestamp.toIso8601String();
     }
 
-    return CoreMessage.create(type: messageType, content: combinedText, extensions: extensions);
+    return CoreMessage.create(
+      type: messageType,
+      content: combinedText,
+      extensions: extensions,
+    );
   }
 
   /// Helper method to manually reconstruct parts that Firebase AI's parseContent doesn't handle
@@ -118,7 +127,9 @@ class FirebaseAiContentAdapter implements GenerativeMessageAdapterBase<Content>
       final response = fr['response'];
       return FunctionResponse(
         fr['name'] as String,
-        response is Map ? Map<String, Object?>.from(response) : <String, Object?>{},
+        response is Map
+            ? Map<String, Object?>.from(response)
+            : <String, Object?>{},
         id: fr['id'] as String?,
       );
     }
@@ -149,7 +160,10 @@ class FirebaseAiContentAdapter implements GenerativeMessageAdapterBase<Content>
     // Handle FileData
     if (partJson.containsKey('fileData')) {
       final fileData = partJson['fileData'] as Map<String, dynamic>;
-      return FileData(fileData['mimeType'] as String, fileData['fileUri'] as String);
+      return FileData(
+        fileData['mimeType'] as String,
+        fileData['fileUri'] as String,
+      );
     }
 
     // If we can't reconstruct it, return null