fix: address KDS PR review comments (round 2)\n\n- Remove Category.kinesis from amplify_categories (KDS is independent of plugin system)\n- Remove KinesisServiceException (SDK errors are silently retried, not exposed)\n- Add KinesisPartitionKeyInvalidException and partition key validation (1-256 chars)\n- Add in-memory cache size tracking to avoid DB query on every record() call\n- Add async lock on record() to prevent concurrent cache size races\n- Add maxIterations safety bound on flush loop\n- Wrap storage query failures in KinesisStorageException in flush\n- Scope sender error catch in _sendStreamBatch to only the putRecords call\n- Make KinesisDataStreamsInterval, KinesisDataStreamsNone, RecordStorage final\n- Use kKinesisMaxBatchBytes constant in getRecordsBatch defaults\n- Fix Timer.periodic arrow syntax style nit

Author: ekjotmultani

packages/amplify_core/lib/src/amplify_class_impl.dart

@@ -46,8 +46,6 @@ class AmplifyClassImpl extends AmplifyClass {
           plugin.cast(),
           authProviderRepo: authProviderRepo,
         );
-      case Category.kinesis:
-        throw UnimplementedError('Kinesis is not an Amplify category plugin');
     }
   }
 }

packages/amplify_core/lib/src/category/amplify_categories.dart

@@ -50,10 +50,7 @@ enum Category {
   storage,
 
   /// Push Notifications
-  pushNotifications,
-
-  /// Kinesis (Data Streams & Firehose)
-  kinesis;
+  pushNotifications;
 
   String get name => switch (this) {
     Category.analytics => 'Analytics',
@@ -63,7 +60,6 @@ enum Category {
     Category.hub => 'Hub',
     Category.storage => 'Storage',
     Category.pushNotifications => 'PushNotifications',
-    Category.kinesis => 'kinesis',
   };
 }
 

packages/kinesis/aws_kinesis_datastreams/lib/src/exception/amplify_kinesis_exception.dart

@@ -37,21 +37,6 @@ final class KinesisStorageException extends AmplifyKinesisException {
         );
 }
 
-/// {@template aws_kinesis_datastreams.kinesis_service_exception}
-/// Thrown when a Kinesis SDK/API error occurs. Inspect [sdkException] for details.
-/// {@endtemplate}
-final class KinesisServiceException extends AmplifyKinesisException {
-  /// {@macro aws_kinesis_datastreams.kinesis_service_exception}
-  KinesisServiceException(String message, {this.sdkException, super.cause})
-      : super(
-          message: message,
-          recoverySuggestion: 'Check sdkException for details.',
-        );
-
-  /// The underlying SDK exception, if any.
-  final Object? sdkException;
-}
-
 /// {@template aws_kinesis_datastreams.kinesis_limit_exceeded_exception}
 /// Thrown when the local cache is full.
 /// {@endtemplate}
@@ -107,6 +92,22 @@ final class KinesisRecordTooLargeException extends AmplifyKinesisException {
         );
 }
 
+/// {@template aws_kinesis_datastreams.kinesis_partition_key_invalid_exception}
+/// Thrown when a partition key is empty or exceeds the Kinesis limit (256 characters).
+/// {@endtemplate}
+final class KinesisPartitionKeyInvalidException extends AmplifyKinesisException {
+  /// {@macro aws_kinesis_datastreams.kinesis_partition_key_invalid_exception}
+  KinesisPartitionKeyInvalidException({
+    required int keyLength,
+  }) : super(
+          message:
+              'Partition key length ($keyLength) is invalid. '
+              'Kinesis requires partition keys to be between 1 and 256 characters.',
+          recoverySuggestion:
+              'Use a partition key between 1 and 256 characters.',
+        );
+}
+
 /// {@template aws_kinesis_datastreams.client_closed_exception}
 /// Thrown when an operation is attempted on a closed client.
 /// {@endtemplate}

packages/kinesis/aws_kinesis_datastreams/lib/src/flush_strategy/flush_strategy.dart

@@ -16,7 +16,7 @@ sealed class KinesisDataStreamsFlushStrategy {
 /// {@template aws_kinesis_datastreams.interval_flush_strategy}
 /// A flush strategy that triggers automatic flushes at a fixed interval.
 /// {@endtemplate}
-class KinesisDataStreamsInterval extends KinesisDataStreamsFlushStrategy {
+final class KinesisDataStreamsInterval extends KinesisDataStreamsFlushStrategy {
   /// {@macro aws_kinesis_datastreams.interval_flush_strategy}
   const KinesisDataStreamsInterval({
     this.interval = const Duration(seconds: 30),
@@ -29,7 +29,7 @@ class KinesisDataStreamsInterval extends KinesisDataStreamsFlushStrategy {
 /// {@template aws_kinesis_datastreams.none_flush_strategy}
 /// A flush strategy that disables automatic flushing.
 /// {@endtemplate}
-class KinesisDataStreamsNone extends KinesisDataStreamsFlushStrategy {
+final class KinesisDataStreamsNone extends KinesisDataStreamsFlushStrategy {
   /// {@macro aws_kinesis_datastreams.none_flush_strategy}
   const KinesisDataStreamsNone();
 }

packages/kinesis/aws_kinesis_datastreams/lib/src/impl/auto_flush_scheduler.dart

@@ -31,7 +31,7 @@ class AutoFlushScheduler {
 
     switch (_strategy) {
       case KinesisDataStreamsInterval(:final interval):
-        _timer = Timer.periodic(interval, (_) { _handleTimerTick(); });
+        _timer = Timer.periodic(interval, (_) => _handleTimerTick());
       case KinesisDataStreamsNone():
         break;
     }

packages/kinesis/aws_kinesis_datastreams/lib/src/impl/record_client.dart

@@ -3,6 +3,8 @@
 
 import 'dart:async';
 
+import 'package:aws_kinesis_datastreams/aws_kinesis_datastreams.dart' show AmplifyKinesisClient;
+import 'package:aws_kinesis_datastreams/src/amplify_kinesis_client.dart' show AmplifyKinesisClient;
 import 'package:aws_kinesis_datastreams/src/db/kinesis_record_database.dart';
 import 'package:aws_kinesis_datastreams/src/exception/amplify_kinesis_exception.dart';
 import 'package:aws_kinesis_datastreams/src/impl/auto_flush_scheduler.dart';
@@ -16,6 +18,8 @@ import 'package:aws_kinesis_datastreams/src/model/flush_data.dart';
 /// {@template aws_kinesis_datastreams.record_client}
 /// Orchestrates record operations, managing the flow between storage,
 /// scheduling, and sending.
+///
+/// Not `final` to allow mocking in tests via [AmplifyKinesisClient.withRecordClient].
 /// {@endtemplate}
 class RecordClient {
   /// {@macro aws_kinesis_datastreams.record_client}
@@ -41,38 +45,79 @@ class RecordClient {
   bool _closed = false;
   bool _flushing = false;
 
+  /// In-memory cache size tracker to avoid a DB query on every record() call.
+  /// Recalculated from DB after deletes (flush/clearCache).
+  int _cachedSize = 0;
+  bool _cacheSizeInitialized = false;
+
+  /// Simple async lock to prevent concurrent record() calls from
+  /// racing on cache size checks.
+  Completer<void>? _recordLock;
+
   /// Maximum batch size in bytes (10 MiB Kinesis PutRecords limit).
-  static const int maxBatchSizeBytes = 10 * 1024 * 1024;
+  static const int maxBatchSizeBytes = kKinesisMaxBatchBytes;
 
   /// Whether the client is currently enabled.
   bool get isEnabled => _enabled;
 
   /// Whether the client has been closed.
   bool get isClosed => _closed;
 
+  /// Ensures the in-memory cache size is initialized from the database.
+  Future<void> _ensureCacheSizeInitialized() async {
+    if (!_cacheSizeInitialized) {
+      _cachedSize = await _storage.getCurrentCacheSize();
+      _cacheSizeInitialized = true;
+    }
+  }
+
   /// Records data to the local cache.
   ///
   /// Throws [ClientClosedException] if the client has been closed.
+  /// Throws [KinesisPartitionKeyInvalidException] if the partition key is
+  /// empty or exceeds 256 characters.
   /// Throws [KinesisRecordTooLargeException] if the record exceeds the
   /// per-record size limit (10 MiB, partition key + data blob).
   /// Throws [KinesisLimitExceededException] if the cache is full.
   Future<void> record(KinesisRecord record) async {
     if (_closed) throw ClientClosedException();
     if (!_enabled) return;
 
+    // Validate partition key length (Kinesis requires 1-256 characters).
+    if (record.partitionKey.isEmpty ||
+        record.partitionKey.length > kKinesisMaxPartitionKeyLength) {
+      throw KinesisPartitionKeyInvalidException(
+        keyLength: record.partitionKey.length,
+      );
+    }
+
     if (record.dataSize > kKinesisMaxRecordBytes) {
       throw KinesisRecordTooLargeException(
         recordBytes: record.dataSize,
         maxBytes: kKinesisMaxRecordBytes,
       );
     }
 
-    final currentSize = await _storage.getCurrentCacheSize();
-    if (currentSize + record.dataSize > _storage.maxCacheBytes) {
-      throw KinesisLimitExceededException();
+    // Acquire async lock to prevent concurrent cache size races.
+    while (_recordLock != null) {
+      await _recordLock!.future;
     }
+    _recordLock = Completer<void>();
+
+    try {
+      await _ensureCacheSizeInitialized();
 
-    await _storage.saveRecord(record);
+      if (_cachedSize + record.dataSize > _storage.maxCacheBytes) {
+        throw KinesisLimitExceededException();
+      }
+
+      await _storage.saveRecord(record);
+      _cachedSize += record.dataSize;
+    } finally {
+      final lock = _recordLock!;
+      _recordLock = null;
+      lock.complete();
+    }
   }
 
   /// Flushes all cached records to Kinesis.
@@ -86,14 +131,36 @@ class RecordClient {
     var totalFlushed = 0;
 
     try {
-      while (true) {
-        final batch = await _storage.getRecordsBatch(
-          maxCount: _maxRecords,
-          maxBytes: maxBatchSizeBytes,
-        );
+      // Safety bound: limit iterations to prevent infinite loops if records
+      // keep failing but never exceed retries within a single flush cycle.
+      var iterations = 0;
+      const maxIterations = 100;
+
+      var consecutiveNoProgress = 0;
+      // Allow enough no-progress iterations for records to exhaust their
+      // retries before considering the batch stuck.
+      final maxConsecutiveNoProgress = _maxRetries + 2;
+
+      while (iterations < maxIterations) {
+        iterations++;
+
+        List<StoredRecord> batch;
+        try {
+          batch = await _storage.getRecordsBatch(
+            maxCount: _maxRecords,
+            maxBytes: maxBatchSizeBytes,
+          );
+        } on Exception catch (e) {
+          throw KinesisStorageException(
+            'Failed to retrieve records batch',
+            cause: e,
+          );
+        }
 
         if (batch.isEmpty) break;
 
+        final countBefore = await _storage.getRecordCount();
+
         final recordsByStream = <String, List<StoredRecord>>{};
         for (final record in batch) {
           recordsByStream.putIfAbsent(record.streamName, () => []).add(record);
@@ -105,58 +172,102 @@ class RecordClient {
         }
 
         await _storage.deleteRecordsExceedingRetries(_maxRetries);
+
+        // Track whether the batch is making progress. If the record count
+        // hasn't decreased for several consecutive iterations, the batch
+        // is stuck (e.g. all records are retryable but haven't exceeded
+        // max retries yet) — break to avoid spinning.
+        final countAfter = await _storage.getRecordCount();
+        if (countAfter < countBefore) {
+          consecutiveNoProgress = 0;
+        } else {
+          consecutiveNoProgress++;
+          if (consecutiveNoProgress >= maxConsecutiveNoProgress) break;
+        }
       }
+
+      // Recalculate in-memory cache size from DB after deletes.
+      _cachedSize = await _storage.getCurrentCacheSize();
     } finally {
       _flushing = false;
     }
 
     return FlushData(recordsFlushed: totalFlushed);
   }
 
-  Future<int> _sendStreamBatch(String streamName, List<StoredRecord> records) async {
+  Future<int> _sendStreamBatch(
+    String streamName,
+    List<StoredRecord> records,
+  ) async {
     final senderRecords = records
-        .map((r) => KinesisSenderRecord(data: r.data, partitionKey: r.partitionKey))
+        .map(
+          (r) => KinesisSenderRecord(
+            data: r.data,
+            partitionKey: r.partitionKey,
+          ),
+        )
         .toList();
 
+    PutRecordsResult result;
     try {
-      final result = await _sender.putRecords(streamName: streamName, records: senderRecords);
-
-      await _storage.deleteRecords(result.successfulRecordIndices.map((i) => records[i].id));
-      await _storage.incrementRetryCount(result.retryableRecordIndices.map((i) => records[i].id));
-      await _storage.deleteRecords(result.failedRecordIndices.map((i) => records[i].id));
-
-      return result.successfulRecordIndices.length;
+      result = await _sender.putRecords(
+        streamName: streamName,
+        records: senderRecords,
+      );
     } on Exception {
+      // Sender/SDK errors — increment retry count and continue.
+      // Non-SDK exceptions (e.g. storage errors) are not caught here
+      // because they originate from _storage calls below, not _sender.
       await _storage.incrementRetryCount(records.map((r) => r.id));
       return 0;
     }
+
+    // Storage operations after a successful send propagate errors to caller.
+    await _storage.deleteRecords(
+      result.successfulRecordIndices.map((i) => records[i].id),
+    );
+    await _storage.incrementRetryCount(
+      result.retryableRecordIndices.map((i) => records[i].id),
+    );
+    await _storage.deleteRecords(
+      result.failedRecordIndices.map((i) => records[i].id),
+    );
+
+    return result.successfulRecordIndices.length;
   }
 
   /// Clears all cached records.
   Future<ClearCacheData> clearCache() async {
     final count = await _storage.getRecordCount();
     await _storage.clear();
+    // Reset in-memory cache size after clearing.
+    _cachedSize = 0;
     return ClearCacheData(recordsCleared: count);
   }
 
+  /// Enables the client to accept and flush records.
   void enable() {
     _enabled = true;
     _scheduler.enable();
   }
 
+  /// Disables the client from accepting and flushing records.
   void disable() {
     _enabled = false;
     _scheduler.disable();
   }
 
+  /// Enables automatic flush operations.
   void enableAutoFlush() {
     _scheduler.enable();
   }
 
+  /// Disables automatic flush operations.
   void disableAutoFlush() {
     _scheduler.disable();
   }
 
+  /// Closes the client and releases all resources.
   Future<void> close() async {
     _closed = true;
     await _scheduler.close();

packages/kinesis/aws_kinesis_datastreams/lib/src/impl/record_storage.dart

@@ -5,12 +5,14 @@ import 'dart:async';
 
 import 'package:aws_kinesis_datastreams/src/db/kinesis_record_database.dart';
 import 'package:aws_kinesis_datastreams/src/impl/kinesis_record.dart';
+import 'package:aws_kinesis_datastreams/src/kinesis_data_streams_options.dart'
+    show kKinesisMaxBatchBytes, kKinesisMaxRecordsPerBatch;
 import 'package:drift/drift.dart';
 
 /// {@template aws_kinesis_datastreams.record_storage}
 /// Manages SQLite database operations for record persistence.
 /// {@endtemplate}
-class RecordStorage {
+final class RecordStorage {
   /// {@macro aws_kinesis_datastreams.record_storage}
   RecordStorage({
     required KinesisRecordDatabase database,
@@ -42,8 +44,8 @@ class RecordStorage {
   /// Returns records up to [maxCount] records and [maxBytes] total size.
   /// Uses window functions to efficiently limit at the database level.
   Future<List<StoredRecord>> getRecordsBatch({
-    int maxCount = 500,
-    int maxBytes = 10 * 1024 * 1024, // 10 MiB for Kinesis PutRecords
+    int maxCount = kKinesisMaxRecordsPerBatch,
+    int maxBytes = kKinesisMaxBatchBytes,
   }) async {
     // Use window functions to compute row number and running size,
     // then filter to get records within both limits.

packages/kinesis/aws_kinesis_datastreams/lib/src/kinesis_data_streams_options.dart

@@ -21,6 +21,9 @@ const int kKinesisMaxBatchBytes = 10 * 1024 * 1024;
 /// partition key and data blob."
 const int kKinesisMaxRecordBytes = 10 * 1024 * 1024;
 
+/// Maximum length of a Kinesis partition key (256 characters).
+const int kKinesisMaxPartitionKeyLength = 256;
+
 /// {@template aws_kinesis_datastreams.amplify_kinesis_client_options}
 /// Configuration options for [AmplifyKinesisClient].
 ///