From a2dfa9e8dfde9f605fd23115d37f46bf1e229f1d Mon Sep 17 00:00:00 2001
From: keyonghan <54558023+keyonghan@users.noreply.github.com>
Date: Tue, 20 Feb 2024 11:12:20 -0800
Subject: [PATCH 1/6] Supports Firestore task update for
 /api/update-task-status (#3495)

Part of https://github.com/flutter/flutter/issues/142951
---
 app_dart/lib/src/model/firestore/task.dart    |  8 ++
 .../request_handlers/update_task_status.dart  | 24 ++++++
 .../update_task_status_test.dart              | 84 +++++++++++++++++++
 3 files changed, 116 insertions(+)

diff --git a/app_dart/lib/src/model/firestore/task.dart b/app_dart/lib/src/model/firestore/task.dart
index e92fcb13c..9a88e6c12 100644
--- a/app_dart/lib/src/model/firestore/task.dart
+++ b/app_dart/lib/src/model/firestore/task.dart
@@ -141,6 +141,14 @@ class Task extends Document {
     return value;
   }
 
+  void setEndTimestamp(int endTimestamp) {
+    fields!['endTimestamp'] = Value(integerValue: endTimestamp.toString());
+  }
+
+  void setTestFlaky(bool testFlaky) {
+    fields!['testFlaky'] = Value(booleanValue: testFlaky);
+  }
+
   /// Update [Task] fields based on a LUCI [Build].
   void updateFromBuild(Build build) {
     final List<String>? tags = build.tags;
diff --git a/app_dart/lib/src/request_handlers/update_task_status.dart b/app_dart/lib/src/request_handlers/update_task_status.dart
index e4e6ed4fd..afd41c855 100644
--- a/app_dart/lib/src/request_handlers/update_task_status.dart
+++ b/app_dart/lib/src/request_handlers/update_task_status.dart
@@ -5,14 +5,17 @@
 import 'dart:async';
 
 import 'package:gcloud/db.dart';
+import 'package:googleapis/firestore/v1.dart';
 import 'package:meta/meta.dart';
 
 import '../model/appengine/commit.dart';
 import '../model/appengine/task.dart';
+import '../model/firestore/task.dart' as firestore;
 import '../request_handling/api_request_handler.dart';
 import '../request_handling/body.dart';
 import '../request_handling/exceptions.dart';
 import '../service/datastore.dart';
+import '../service/firestore.dart';
 import '../service/logging.dart';
 
 /// Endpoint for task runners to update Cocoon with test run information.
@@ -60,9 +63,30 @@ class UpdateTaskStatus extends ApiRequestHandler<UpdateTaskStatusResponse> {
     task.isTestFlaky = isTestFlaky;
 
     await datastore.insert(<Task>[task]);
+
+    try {
+      await updateTaskDocument(task.status, task.endTimestamp!, task.isTestFlaky!);
+    } catch (error) {
+      log.warning('Failed to update task in Firestore: $error');
+    }
     return UpdateTaskStatusResponse(task);
   }
 
+  Future<void> updateTaskDocument(String status, int endTimestamp, bool isTestFlaky) async {
+    final FirestoreService firestoreService = await config.createFirestoreService();
+    final String sha = (requestData![gitShaParam] as String).trim();
+    final String? taskName = requestData![builderNameParam] as String?;
+    final String documentName = '$kDatabase/documents/tasks/${sha}_${taskName}_1';
+    log.info('getting firestore document: $documentName');
+    final firestore.Task firestoreTask =
+        await firestore.Task.fromFirestore(firestoreService: firestoreService, documentName: documentName);
+    firestoreTask.setStatus(status);
+    firestoreTask.setEndTimestamp(endTimestamp);
+    firestoreTask.setTestFlaky(isTestFlaky);
+    final List<Write> writes = documentsToWrites([firestoreTask], exists: true);
+    await firestoreService.batchWriteDocuments(BatchWriteRequest(writes: writes), kDatabase);
+  }
+
   /// Retrieve [Task] from [DatastoreService] when given [gitShaParam], [gitBranchParam], and [builderNameParam].
   ///
   /// This is used when the DeviceLab test runner is uploading results to Cocoon for runs on LUCI.
diff --git a/app_dart/test/request_handlers/update_task_status_test.dart b/app_dart/test/request_handlers/update_task_status_test.dart
index 9325d1f93..77a6bef70 100644
--- a/app_dart/test/request_handlers/update_task_status_test.dart
+++ b/app_dart/test/request_handlers/update_task_status_test.dart
@@ -4,10 +4,13 @@
 
 import 'package:cocoon_service/src/model/appengine/commit.dart';
 import 'package:cocoon_service/src/model/appengine/task.dart';
+import 'package:cocoon_service/src/model/firestore/task.dart' as firestore;
 import 'package:cocoon_service/src/request_handlers/update_task_status.dart';
 import 'package:cocoon_service/src/service/config.dart';
 import 'package:cocoon_service/src/service/datastore.dart';
 import 'package:gcloud/db.dart';
+import 'package:googleapis/firestore/v1.dart';
+import 'package:mockito/mockito.dart';
 import 'package:test/test.dart';
 
 import '../src/bigquery/fake_tabledata_resource.dart';
@@ -15,11 +18,14 @@ import '../src/datastore/fake_config.dart';
 import '../src/datastore/fake_datastore.dart';
 import '../src/request_handling/api_request_handler_tester.dart';
 import '../src/request_handling/fake_authentication.dart';
+import '../src/utilities/entity_generators.dart';
+import '../src/utilities/mocks.dart';
 
 void main() {
   group('UpdateTaskStatus', () {
     late FakeConfig config;
     late ApiRequestHandlerTester tester;
+    late MockFirestoreService mockFirestoreService;
     late UpdateTaskStatus handler;
     final FakeTabledataResource tabledataResourceApi = FakeTabledataResource();
     late Commit commit;
@@ -27,11 +33,13 @@ void main() {
     const int taskId = 4506830800027648;
 
     setUp(() {
+      mockFirestoreService = MockFirestoreService();
       final FakeDatastoreDB datastoreDB = FakeDatastoreDB();
       config = FakeConfig(
         dbValue: datastoreDB,
         tabledataResource: tabledataResourceApi,
         maxTaskRetriesValue: 2,
+        firestoreService: mockFirestoreService,
       );
       tester = ApiRequestHandlerTester();
       handler = UpdateTaskStatus(
@@ -48,6 +56,24 @@ void main() {
     });
 
     test('TestFlaky is false when not injected', () async {
+      final firestore.Task firestoreTask = generateFirestoreTask(1);
+      when(
+        mockFirestoreService.getDocument(
+          captureAny,
+        ),
+      ).thenAnswer((Invocation invocation) {
+        return Future<Document>.value(
+          firestoreTask,
+        );
+      });
+      when(
+        mockFirestoreService.batchWriteDocuments(
+          captureAny,
+          captureAny,
+        ),
+      ).thenAnswer((Invocation invocation) {
+        return Future<BatchWriteResponse>.value(BatchWriteResponse());
+      });
       final Task task = Task(
         key: commit.key.append(Task, id: taskId),
         name: 'integration_ui_ios',
@@ -69,9 +95,28 @@ void main() {
       await tester.post(handler);
 
       expect(task.isTestFlaky, false);
+      expect(firestoreTask.testFlaky, false);
     });
 
     test('TestFlaky is true when injected', () async {
+      final firestore.Task firestoreTask = generateFirestoreTask(1);
+      when(
+        mockFirestoreService.getDocument(
+          captureAny,
+        ),
+      ).thenAnswer((Invocation invocation) {
+        return Future<Document>.value(
+          firestoreTask,
+        );
+      });
+      when(
+        mockFirestoreService.batchWriteDocuments(
+          captureAny,
+          captureAny,
+        ),
+      ).thenAnswer((Invocation invocation) {
+        return Future<BatchWriteResponse>.value(BatchWriteResponse());
+      });
       final Task task = Task(
         key: commit.key.append(Task, id: taskId),
         name: 'integration_ui_ios',
@@ -94,9 +139,28 @@ void main() {
       await tester.post(handler);
 
       expect(task.isTestFlaky, true);
+      expect(firestoreTask.testFlaky, true);
     });
 
     test('task name requests can update tasks', () async {
+      final firestore.Task firestoreTask = generateFirestoreTask(1);
+      when(
+        mockFirestoreService.getDocument(
+          captureAny,
+        ),
+      ).thenAnswer((Invocation invocation) {
+        return Future<Document>.value(
+          firestoreTask,
+        );
+      });
+      when(
+        mockFirestoreService.batchWriteDocuments(
+          captureAny,
+          captureAny,
+        ),
+      ).thenAnswer((Invocation invocation) {
+        return Future<BatchWriteResponse>.value(BatchWriteResponse());
+      });
       final Task task = Task(
         key: commit.key.append(Task, id: taskId),
         name: 'integration_ui_ios',
@@ -119,6 +183,8 @@ void main() {
 
       expect(task.status, 'Failed');
       expect(task.attempts, 1);
+      expect(firestoreTask.status, 'Failed');
+      expect(firestoreTask.attempts, 1);
     });
 
     test('task name requests when task does not exists returns exception', () async {
@@ -132,6 +198,24 @@ void main() {
     });
 
     test('task name request updates when input has whitespace', () async {
+      final firestore.Task firestoreTask = generateFirestoreTask(1);
+      when(
+        mockFirestoreService.getDocument(
+          captureAny,
+        ),
+      ).thenAnswer((Invocation invocation) {
+        return Future<Document>.value(
+          firestoreTask,
+        );
+      });
+      when(
+        mockFirestoreService.batchWriteDocuments(
+          captureAny,
+          captureAny,
+        ),
+      ).thenAnswer((Invocation invocation) {
+        return Future<BatchWriteResponse>.value(BatchWriteResponse());
+      });
       config.db.values[commit.key] = commit;
       final Task cocoonTask = Task(
         key: commit.key.append(Task, id: taskId),

From 83a44f0d441786b5832d47db3d3a6224065d78e5 Mon Sep 17 00:00:00 2001
From: keyonghan <54558023+keyonghan@users.noreply.github.com>
Date: Tue, 20 Feb 2024 11:20:17 -0800
Subject: [PATCH 2/6] Support multiple records for retries on a same task
 (#3494)

This PR adds support to save multiple entries for a task retry. Old logic overrides task statuses on a single task, which fails to track statuses.

This PR:
1) makes it customizable to collect `attempts` info
2) adds `current_attempt` tag to make it easy to track `attempts` info.
3) handles `attempt` increase when retry is triggered.

Part of https://github.com/flutter/flutter/issues/142951.
---
 app_dart/lib/src/model/firestore/task.dart    | 19 ++++-
 .../postsubmit_luci_subscription.dart         | 25 ++++--
 .../presubmit_luci_subscription.dart          |  2 +-
 .../src/service/github_checks_service.dart    |  3 +-
 .../lib/src/service/luci_build_service.dart   | 26 ++++--
 app_dart/test/model/firestore/task_test.dart  | 15 ++++
 .../postsubmit_luci_subscription_test.dart    | 81 ++++++++++++++++++-
 .../test/service/luci_build_service_test.dart | 46 ++++++++++-
 .../test/src/utilities/entity_generators.dart |  3 +-
 app_dart/test/src/utilities/mocks.mocks.dart  | 17 ++--
 cloud_build/verify_provenance.sh              |  2 +-
 11 files changed, 211 insertions(+), 28 deletions(-)

diff --git a/app_dart/lib/src/model/firestore/task.dart b/app_dart/lib/src/model/firestore/task.dart
index 9a88e6c12..527584535 100644
--- a/app_dart/lib/src/model/firestore/task.dart
+++ b/app_dart/lib/src/model/firestore/task.dart
@@ -94,7 +94,10 @@ class Task extends Document {
   ///
   /// This is a human-readable name, typically a test name (e.g.
   /// "hello_world__memory").
-  String? get taskName => fields!['endTimestamp']!.stringValue!;
+  String? get taskName => fields!['name']!.stringValue!;
+
+  /// The sha of the task commit.
+  String? get commitSha => fields!['commitSha']!.stringValue!;
 
   /// The number of attempts that have been made to run this task successfully.
   ///
@@ -166,6 +169,20 @@ class Task extends Document {
     _setStatusFromLuciStatus(build);
   }
 
+  void resetAsRetry({int attempt = 1}) {
+    name = '$kDatabase/documents/tasks/${commitSha}_${taskName}_$attempt';
+    fields = <String, Value>{
+      'createTimestamp': Value(integerValue: DateTime.now().millisecondsSinceEpoch.toString()),
+      'endTimestamp': Value(integerValue: '0'),
+      'bringup': Value(booleanValue: bringup),
+      'name': Value(stringValue: taskName),
+      'startTimestamp': Value(integerValue: '0'),
+      'status': Value(stringValue: Task.statusNew),
+      'testFlaky': Value(booleanValue: false),
+      'commitSha': Value(stringValue: commitSha),
+    };
+  }
+
   /// Get a [Task] status from a LUCI [Build] status/result.
   String _setStatusFromLuciStatus(Build build) {
     // Updates can come out of order. Ensure completed statuses are kept.
diff --git a/app_dart/lib/src/request_handlers/postsubmit_luci_subscription.dart b/app_dart/lib/src/request_handlers/postsubmit_luci_subscription.dart
index 90aa1da6b..2a5bb783e 100644
--- a/app_dart/lib/src/request_handlers/postsubmit_luci_subscription.dart
+++ b/app_dart/lib/src/request_handlers/postsubmit_luci_subscription.dart
@@ -9,7 +9,7 @@ import 'package:meta/meta.dart';
 
 import '../model/appengine/commit.dart';
 import '../model/appengine/task.dart';
-import '../model/firestore/task.dart' as f;
+import '../model/firestore/task.dart' as firestore;
 import '../model/luci/push_message.dart';
 import '../request_handling/body.dart';
 import '../request_handling/exceptions.dart';
@@ -45,6 +45,7 @@ class PostsubmitLuciSubscription extends SubscriptionHandler {
   @override
   Future<Body> post() async {
     final DatastoreService datastore = datastoreProvider(config.db);
+    final FirestoreService firestoreService = await config.createFirestoreService();
 
     final BuildPushMessage buildPushMessage = BuildPushMessage.fromPushMessage(message);
     log.fine('userData=${buildPushMessage.userData}');
@@ -83,12 +84,14 @@ class PostsubmitLuciSubscription extends SubscriptionHandler {
     }
     log.fine('Found $task');
 
+    firestore.Task taskDocument = firestore.Task();
+
     if (_shouldUpdateTask(build, task)) {
       final String oldTaskStatus = task.status;
       task.updateFromBuild(build);
       await datastore.insert(<Task>[task]);
       try {
-        await updateFirestore(build, rawCommitKey, task.name!);
+        taskDocument = await updateFirestore(build, rawCommitKey, task.name!, firestoreService);
       } catch (error) {
         log.warning('Failed to update task in Firestore: $error');
       }
@@ -114,6 +117,8 @@ class PostsubmitLuciSubscription extends SubscriptionHandler {
         target: target,
         task: task,
         datastore: datastore,
+        taskDocument: taskDocument,
+        firestoreService: firestoreService,
       );
       log.info('Retried: $retried');
     }
@@ -144,17 +149,23 @@ class PostsubmitLuciSubscription extends SubscriptionHandler {
   }
 
   /// Queries the task document and updates based on the latest build data.
-  Future<void> updateFirestore(Build build, String commitKeyId, String taskName) async {
-    final FirestoreService firestoreService = await config.createFirestoreService();
+  Future<firestore.Task> updateFirestore(
+    Build build,
+    String commitKeyId,
+    String taskName,
+    FirestoreService firestoreService,
+  ) async {
+    final int currentAttempt = githubChecksService.currentAttempt(build);
     final String sha = commitKeyId.split('/').last;
-    final String documentName = '$kDatabase/documents/tasks/${sha}_${taskName}_1';
+    final String documentName = '$kDatabase/documents/tasks/${sha}_${taskName}_$currentAttempt';
     log.info('getting firestore document: $documentName');
-    final f.Task firestoreTask =
-        await f.Task.fromFirestore(firestoreService: firestoreService, documentName: documentName);
+    final firestore.Task firestoreTask =
+        await firestore.Task.fromFirestore(firestoreService: firestoreService, documentName: documentName);
     log.info('updating firestoreTask based on build');
     firestoreTask.updateFromBuild(build);
     log.info('finished updating firestoreTask based on builds');
     final List<Write> writes = documentsToWrites([firestoreTask], exists: true);
     await firestoreService.batchWriteDocuments(BatchWriteRequest(writes: writes), kDatabase);
+    return firestoreTask;
   }
 }
diff --git a/app_dart/lib/src/request_handlers/presubmit_luci_subscription.dart b/app_dart/lib/src/request_handlers/presubmit_luci_subscription.dart
index 63d484106..8049eb282 100644
--- a/app_dart/lib/src/request_handlers/presubmit_luci_subscription.dart
+++ b/app_dart/lib/src/request_handlers/presubmit_luci_subscription.dart
@@ -72,7 +72,7 @@ class PresubmitLuciSubscription extends SubscriptionHandler {
       );
       bool rescheduled = false;
       if (githubChecksService.taskFailed(buildPushMessage)) {
-        final int currentAttempt = githubChecksService.currentAttempt(buildPushMessage);
+        final int currentAttempt = githubChecksService.currentAttempt(build);
         final int maxAttempt = await _getMaxAttempt(buildPushMessage, slug, builderName);
         if (currentAttempt < maxAttempt) {
           rescheduled = true;
diff --git a/app_dart/lib/src/service/github_checks_service.dart b/app_dart/lib/src/service/github_checks_service.dart
index 9a962e24f..48e208505 100644
--- a/app_dart/lib/src/service/github_checks_service.dart
+++ b/app_dart/lib/src/service/github_checks_service.dart
@@ -144,8 +144,7 @@ class GithubChecksService {
   /// Returns current reschedule attempt.
   ///
   /// It returns 1 if this is the first run, and +1 with each reschedule.
-  int currentAttempt(push_message.BuildPushMessage buildPushMessage) {
-    final push_message.Build build = buildPushMessage.build!;
+  int currentAttempt(push_message.Build build) {
     if (build.tagsByName('current_attempt').isEmpty) {
       return 1;
     } else {
diff --git a/app_dart/lib/src/service/luci_build_service.dart b/app_dart/lib/src/service/luci_build_service.dart
index 7537d5967..4327ac519 100644
--- a/app_dart/lib/src/service/luci_build_service.dart
+++ b/app_dart/lib/src/service/luci_build_service.dart
@@ -7,26 +7,23 @@ import 'dart:convert';
 import 'dart:math';
 import 'dart:typed_data';
 
+import 'package:cocoon_service/cocoon_service.dart';
 import 'package:github/github.dart' as github;
 import 'package:github/hooks.dart';
+import 'package:googleapis/firestore/v1.dart' hide Status;
 import 'package:googleapis/pubsub/v1.dart';
 
 import '../foundation/github_checks_util.dart';
-import '../foundation/utils.dart';
 import '../model/appengine/commit.dart';
 import '../model/appengine/task.dart';
+import '../model/firestore/task.dart' as f;
 import '../model/ci_yaml/target.dart';
 import '../model/github/checks.dart' as cocoon_checks;
 import '../model/luci/buildbucket.dart';
 import '../model/luci/push_message.dart' as push_message;
-import '../request_handling/pubsub.dart';
 import '../service/datastore.dart';
 import '../service/logging.dart';
-import 'buildbucket.dart';
-import 'cache_service.dart';
-import 'config.dart';
 import 'exceptions.dart';
-import 'gerrit_service.dart';
 import 'github_service.dart';
 
 const Set<String> taskFailStatusSet = <String>{
@@ -619,6 +616,8 @@ class LuciBuildService {
     tags['user_agent'] = <String>['flutter-cocoon'];
     // Tag `scheduler_job_id` is needed when calling buildbucket search build API.
     tags['scheduler_job_id'] = <String>['flutter/${target.value.name}'];
+    // Default attempt is the initial attempt, which is 1.
+    tags['current_attempt'] = tags['current_attempt'] ?? <String>['1'];
     final Map<String, Object> processedProperties = target.getProperties();
     processedProperties.addAll(properties ?? <String, Object>{});
     processedProperties['git_branch'] = commit.branch!;
@@ -682,8 +681,10 @@ class LuciBuildService {
     required Target target,
     required Task task,
     required DatastoreService datastore,
+    FirestoreService? firestoreService,
     Map<String, List<String>>? tags,
     bool ignoreChecks = false,
+    f.Task? taskDocument,
   }) async {
     if (ignoreChecks == false && await _shouldRerunBuilder(task, commit, datastore) == false) {
       return false;
@@ -692,6 +693,19 @@ class LuciBuildService {
     tags ??= <String, List<String>>{};
     tags['trigger_type'] = <String>['retry'];
 
+    // TODO(keyonghan): remove check when [ResetProdTask] supports firestore update.
+    if (taskDocument != null) {
+      try {
+        final int newAttempt = int.parse(taskDocument.name!.split('_').last) + 1;
+        tags['current_attempt'] = <String>[newAttempt.toString()];
+        taskDocument.resetAsRetry(attempt: newAttempt);
+        final List<Write> writes = documentsToWrites([taskDocument]);
+        await firestoreService!.batchWriteDocuments(BatchWriteRequest(writes: writes), kDatabase);
+      } catch (error) {
+        log.warning('Failed to insert retried task in Firestore: $error');
+      }
+    }
+
     final BatchRequest request = BatchRequest(
       requests: <Request>[
         Request(
diff --git a/app_dart/test/model/firestore/task_test.dart b/app_dart/test/model/firestore/task_test.dart
index 836ad2ae8..997630880 100644
--- a/app_dart/test/model/firestore/task_test.dart
+++ b/app_dart/test/model/firestore/task_test.dart
@@ -146,4 +146,19 @@ void main() {
       expect(resultedTask.fields, firestoreTask.fields);
     });
   });
+
+  group('resert as retry', () {
+    test('success', () {
+      final Task task = generateFirestoreTask(
+        1,
+        status: Task.statusFailed,
+        testFlaky: true,
+      );
+      task.resetAsRetry(attempt: 2);
+
+      expect(int.parse(task.name!.split('_').last), 2);
+      expect(task.status, Task.statusNew);
+      expect(task.testFlaky, false);
+    });
+  });
 }
diff --git a/app_dart/test/request_handlers/postsubmit_luci_subscription_test.dart b/app_dart/test/request_handlers/postsubmit_luci_subscription_test.dart
index f87eb5a77..119f28149 100644
--- a/app_dart/test/request_handlers/postsubmit_luci_subscription_test.dart
+++ b/app_dart/test/request_handlers/postsubmit_luci_subscription_test.dart
@@ -5,7 +5,7 @@
 import 'package:cocoon_service/cocoon_service.dart';
 import 'package:cocoon_service/src/model/appengine/commit.dart';
 import 'package:cocoon_service/src/model/appengine/task.dart';
-import 'package:cocoon_service/src/model/firestore/task.dart' as f;
+import 'package:cocoon_service/src/model/firestore/task.dart' as firestore;
 import 'package:cocoon_service/src/request_handling/exceptions.dart';
 import 'package:cocoon_service/src/service/datastore.dart';
 import 'package:gcloud/db.dart';
@@ -92,7 +92,7 @@ void main() {
   });
 
   test('updates task based on message', () async {
-    final f.Task firestoreTask = generateFirestoreTask(1);
+    final firestore.Task firestoreTask = generateFirestoreTask(1, attempts: 2);
     when(
       mockFirestoreService.getDocument(
         captureAny,
@@ -110,6 +110,7 @@ void main() {
     ).thenAnswer((Invocation invocation) {
       return Future<BatchWriteResponse>.value(BatchWriteResponse());
     });
+    when(mockGithubChecksService.currentAttempt(any)).thenAnswer((_) => 2);
     final Commit commit = generateCommit(1, sha: '87f88734747805589f2131753620d61b22922822');
     final Task task = generateTask(
       4507531199512576,
@@ -301,6 +302,25 @@ void main() {
   });
 
   test('fallback to build parameters if task_key is not present', () async {
+    when(mockGithubChecksService.currentAttempt(any)).thenAnswer((_) => 2);
+    final firestore.Task firestoreTask = generateFirestoreTask(1, attempts: 2);
+    when(
+      mockFirestoreService.getDocument(
+        captureAny,
+      ),
+    ).thenAnswer((Invocation invocation) {
+      return Future<Document>.value(
+        firestoreTask,
+      );
+    });
+    when(
+      mockFirestoreService.batchWriteDocuments(
+        captureAny,
+        captureAny,
+      ),
+    ).thenAnswer((Invocation invocation) {
+      return Future<BatchWriteResponse>.value(BatchWriteResponse());
+    });
     final Commit commit = generateCommit(1, sha: '87f88734747805589f2131753620d61b22922822');
     final Task task = generateTask(
       4507531199512576,
@@ -325,6 +345,25 @@ void main() {
   test('non-bringup target updates check run', () async {
     scheduler.ciYaml = nonBringupPackagesConfig;
     when(mockGithubChecksService.updateCheckStatus(any, any, any)).thenAnswer((_) async => true);
+    when(mockGithubChecksService.currentAttempt(any)).thenAnswer((_) => 2);
+    final firestore.Task firestoreTask = generateFirestoreTask(1, attempts: 2);
+    when(
+      mockFirestoreService.getDocument(
+        captureAny,
+      ),
+    ).thenAnswer((Invocation invocation) {
+      return Future<Document>.value(
+        firestoreTask,
+      );
+    });
+    when(
+      mockFirestoreService.batchWriteDocuments(
+        captureAny,
+        captureAny,
+      ),
+    ).thenAnswer((Invocation invocation) {
+      return Future<BatchWriteResponse>.value(BatchWriteResponse());
+    });
     final Commit commit = generateCommit(1, sha: '87f88734747805589f2131753620d61b22922822', repo: 'packages');
     final Task task = generateTask(
       4507531199512576,
@@ -349,6 +388,25 @@ void main() {
   test('bringup target does not update check run', () async {
     scheduler.ciYaml = bringupPackagesConfig;
     when(mockGithubChecksService.updateCheckStatus(any, any, any)).thenAnswer((_) async => true);
+    when(mockGithubChecksService.currentAttempt(any)).thenAnswer((_) => 2);
+    final firestore.Task firestoreTask = generateFirestoreTask(1, attempts: 2);
+    when(
+      mockFirestoreService.getDocument(
+        captureAny,
+      ),
+    ).thenAnswer((Invocation invocation) {
+      return Future<Document>.value(
+        firestoreTask,
+      );
+    });
+    when(
+      mockFirestoreService.batchWriteDocuments(
+        captureAny,
+        captureAny,
+      ),
+    ).thenAnswer((Invocation invocation) {
+      return Future<BatchWriteResponse>.value(BatchWriteResponse());
+    });
     final Commit commit = generateCommit(1, sha: '87f88734747805589f2131753620d61b22922822');
     final Task task = generateTask(
       4507531199512576,
@@ -373,6 +431,25 @@ void main() {
   test('unsupported repo target does not update check run', () async {
     scheduler.ciYaml = unsupportedPostsubmitCheckrunConfig;
     when(mockGithubChecksService.updateCheckStatus(any, any, any)).thenAnswer((_) async => true);
+    when(mockGithubChecksService.currentAttempt(any)).thenAnswer((_) => 2);
+    final firestore.Task firestoreTask = generateFirestoreTask(1, attempts: 2);
+    when(
+      mockFirestoreService.getDocument(
+        captureAny,
+      ),
+    ).thenAnswer((Invocation invocation) {
+      return Future<Document>.value(
+        firestoreTask,
+      );
+    });
+    when(
+      mockFirestoreService.batchWriteDocuments(
+        captureAny,
+        captureAny,
+      ),
+    ).thenAnswer((Invocation invocation) {
+      return Future<BatchWriteResponse>.value(BatchWriteResponse());
+    });
     final Commit commit = generateCommit(1, sha: '87f88734747805589f2131753620d61b22922822');
     final Task task = generateTask(
       4507531199512576,
diff --git a/app_dart/test/service/luci_build_service_test.dart b/app_dart/test/service/luci_build_service_test.dart
index bb687bc28..283ef94ee 100644
--- a/app_dart/test/service/luci_build_service_test.dart
+++ b/app_dart/test/service/luci_build_service_test.dart
@@ -8,13 +8,15 @@ import 'dart:core';
 import 'package:cocoon_service/cocoon_service.dart';
 import 'package:cocoon_service/src/model/appengine/commit.dart';
 import 'package:cocoon_service/src/model/appengine/task.dart';
+import 'package:cocoon_service/src/model/firestore/task.dart' as f;
 import 'package:cocoon_service/src/model/ci_yaml/target.dart';
+import 'package:cocoon_service/src/model/github/checks.dart' as cocoon_checks;
 import 'package:cocoon_service/src/model/luci/buildbucket.dart';
 import 'package:cocoon_service/src/model/luci/push_message.dart' as push_message;
 import 'package:cocoon_service/src/service/exceptions.dart';
 import 'package:cocoon_service/src/service/datastore.dart';
 import 'package:github/github.dart';
-import 'package:cocoon_service/src/model/github/checks.dart' as cocoon_checks;
+import 'package:googleapis/firestore/v1.dart' hide Status;
 import 'package:mockito/mockito.dart';
 import 'package:test/test.dart';
 
@@ -911,11 +913,13 @@ void main() {
     late Commit totCommit;
     late DatastoreService datastore;
     late MockGithubChecksUtil mockGithubChecksUtil;
+    late MockFirestoreService mockFirestoreService;
     setUp(() {
       cache = CacheService(inMemory: true);
       config = FakeConfig();
       mockBuildBucketClient = MockBuildBucketClient();
       mockGithubChecksUtil = MockGithubChecksUtil();
+      mockFirestoreService = MockFirestoreService();
       when(mockGithubChecksUtil.createCheckRun(any, any, any, any, output: anyNamed('output')))
           .thenAnswer((realInvocation) async => generateCheckRun(1));
       pubsub = FakePubSub();
@@ -1064,5 +1068,45 @@ void main() {
       );
       expect(rerunFlag, isFalse);
     });
+
+    test('insert retried task document to firestore', () async {
+      final f.Task firestoreTask = generateFirestoreTask(1, attempts: 1);
+      when(
+        mockFirestoreService.batchWriteDocuments(
+          captureAny,
+          captureAny,
+        ),
+      ).thenAnswer((Invocation invocation) {
+        return Future<BatchWriteResponse>.value(BatchWriteResponse());
+      });
+      totCommit = generateCommit(1);
+      config.db.values[totCommit.key] = totCommit;
+      config.maxLuciTaskRetriesValue = 1;
+      final Task task = generateTask(
+        1,
+        status: Task.statusInfraFailure,
+        parent: totCommit,
+        buildNumber: 1,
+      );
+      final Target target = generateTarget(1);
+      expect(firestoreTask.attempts, 1);
+      final bool rerunFlag = await service.checkRerunBuilder(
+        commit: totCommit,
+        task: task,
+        target: target,
+        datastore: datastore,
+        firestoreService: mockFirestoreService,
+        taskDocument: firestoreTask,
+      );
+      expect(rerunFlag, isTrue);
+
+      expect(firestoreTask.attempts, 2);
+      final List<dynamic> captured = verify(mockFirestoreService.batchWriteDocuments(captureAny, captureAny)).captured;
+      expect(captured.length, 2);
+      final BatchWriteRequest batchWriteRequest = captured[0] as BatchWriteRequest;
+      expect(batchWriteRequest.writes!.length, 1);
+      final Document insertedTaskDocument = batchWriteRequest.writes![0].update!;
+      expect(insertedTaskDocument, firestoreTask);
+    });
   });
 }
diff --git a/app_dart/test/src/utilities/entity_generators.dart b/app_dart/test/src/utilities/entity_generators.dart
index a28229840..957d83fe1 100644
--- a/app_dart/test/src/utilities/entity_generators.dart
+++ b/app_dart/test/src/utilities/entity_generators.dart
@@ -112,7 +112,7 @@ f.Task generateFirestoreTask(
   final String taskName = name ?? 'task$i';
   final String sha = commitSha ?? 'testSha';
   final f.Task task = f.Task()
-    ..name = '${sha}_${taskName}_1'
+    ..name = '${sha}_${taskName}_$attempts'
     ..fields = <String, Value>{
       'createTimestamp': Value(integerValue: (created?.millisecondsSinceEpoch ?? 0).toString()),
       'startTimestamp': Value(integerValue: (started?.millisecondsSinceEpoch ?? 0).toString()),
@@ -121,6 +121,7 @@ f.Task generateFirestoreTask(
       'testFlaky': Value(booleanValue: testFlaky),
       'status': Value(stringValue: status),
       'name': Value(stringValue: taskName),
+      'commitSha': Value(stringValue: sha),
     };
   if (buildNumber != null) {
     task.fields!['buildNumber'] = Value(integerValue: buildNumber.toString());
diff --git a/app_dart/test/src/utilities/mocks.mocks.dart b/app_dart/test/src/utilities/mocks.mocks.dart
index 2258efeae..3948a8d1e 100644
--- a/app_dart/test/src/utilities/mocks.mocks.dart
+++ b/app_dart/test/src/utilities/mocks.mocks.dart
@@ -19,6 +19,7 @@ import 'package:cocoon_service/src/model/appengine/key_helper.dart' as _i12;
 import 'package:cocoon_service/src/model/appengine/stage.dart' as _i34;
 import 'package:cocoon_service/src/model/appengine/task.dart' as _i33;
 import 'package:cocoon_service/src/model/ci_yaml/target.dart' as _i38;
+import 'package:cocoon_service/src/model/firestore/task.dart' as _i40;
 import 'package:cocoon_service/src/model/github/checks.dart' as _i39;
 import 'package:cocoon_service/src/model/luci/buildbucket.dart' as _i8;
 import 'package:cocoon_service/src/model/luci/push_message.dart' as _i37;
@@ -41,11 +42,11 @@ import 'package:http/http.dart' as _i2;
 import 'package:mockito/mockito.dart' as _i1;
 import 'package:mockito/src/dummies.dart' as _i28;
 import 'package:neat_cache/neat_cache.dart' as _i25;
-import 'package:process/src/interface/process_manager.dart' as _i40;
+import 'package:process/src/interface/process_manager.dart' as _i41;
 import 'package:retry/retry.dart' as _i27;
 
 import '../../service/cache_service_test.dart' as _i35;
-import 'mocks.dart' as _i41;
+import 'mocks.dart' as _i42;
 
 // ignore_for_file: type=lint
 // ignore_for_file: avoid_redundant_argument_values
@@ -3395,10 +3396,10 @@ class MockGithubChecksService extends _i1.Mock implements _i15.GithubChecksServi
       ) as bool);
 
   @override
-  int currentAttempt(_i37.BuildPushMessage? buildPushMessage) => (super.noSuchMethod(
+  int currentAttempt(_i37.Build? build) => (super.noSuchMethod(
         Invocation.method(
           #currentAttempt,
-          [buildPushMessage],
+          [build],
         ),
         returnValue: 0,
       ) as int);
@@ -6600,8 +6601,10 @@ class MockLuciBuildService extends _i1.Mock implements _i15.LuciBuildService {
     required _i38.Target? target,
     required _i33.Task? task,
     required _i9.DatastoreService? datastore,
+    _i15.FirestoreService? firestoreService,
     Map<String, List<String>>? tags,
     bool? ignoreChecks = false,
+    _i40.Task? taskDocument,
   }) =>
       (super.noSuchMethod(
         Invocation.method(
@@ -6612,8 +6615,10 @@ class MockLuciBuildService extends _i1.Mock implements _i15.LuciBuildService {
             #target: target,
             #task: task,
             #datastore: datastore,
+            #firestoreService: firestoreService,
             #tags: tags,
             #ignoreChecks: ignoreChecks,
+            #taskDocument: taskDocument,
           },
         ),
         returnValue: _i20.Future<bool>.value(false),
@@ -6623,7 +6628,7 @@ class MockLuciBuildService extends _i1.Mock implements _i15.LuciBuildService {
 /// A class which mocks [ProcessManager].
 ///
 /// See the documentation for Mockito's code generation for more information.
-class MockProcessManager extends _i1.Mock implements _i40.ProcessManager {
+class MockProcessManager extends _i1.Mock implements _i41.ProcessManager {
   MockProcessManager() {
     _i1.throwOnMissingStub(this);
   }
@@ -9354,7 +9359,7 @@ class MockGitHub extends _i1.Mock implements _i13.GitHub {
             #preview: preview,
           },
         ),
-        returnValue: _i41.postJsonShim<S, T>(
+        returnValue: _i42.postJsonShim<S, T>(
           path,
           statusCode: statusCode,
           fail: fail,
diff --git a/cloud_build/verify_provenance.sh b/cloud_build/verify_provenance.sh
index 9b11808bf..80d8a6158 100755
--- a/cloud_build/verify_provenance.sh
+++ b/cloud_build/verify_provenance.sh
@@ -10,7 +10,7 @@ set -e
 # legitimate, then the script will exit with a non-zero exit code.
 PROVENANCE_PATH=$1
 BUILDER_ID=https://cloudbuild.googleapis.com/GoogleHostedWorker
-SOURCE_URI=github.com/flutter/cocoon
+SOURCE_URI=github.com/keyonghan/cocoon
 
 # Download the jq binary in order to obtain the artifact registry url from the
 # docker image provenance.

From 5d1ed59deb63b2a47d73799a357dac4fe544ad64 Mon Sep 17 00:00:00 2001
From: keyonghan <54558023+keyonghan@users.noreply.github.com>
Date: Tue, 20 Feb 2024 13:04:05 -0800
Subject: [PATCH 3/6] Restore deploy repo url (#3504)

Mistakenly merged in https://github.com/flutter/cocoon/pull/3494. This PR restores the URL.
---
 cloud_build/verify_provenance.sh | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/cloud_build/verify_provenance.sh b/cloud_build/verify_provenance.sh
index 80d8a6158..9b11808bf 100755
--- a/cloud_build/verify_provenance.sh
+++ b/cloud_build/verify_provenance.sh
@@ -10,7 +10,7 @@ set -e
 # legitimate, then the script will exit with a non-zero exit code.
 PROVENANCE_PATH=$1
 BUILDER_ID=https://cloudbuild.googleapis.com/GoogleHostedWorker
-SOURCE_URI=github.com/keyonghan/cocoon
+SOURCE_URI=github.com/flutter/cocoon
 
 # Download the jq binary in order to obtain the artifact registry url from the
 # docker image provenance.

From eff005a8b315061c7af6225933118b1e66670ada Mon Sep 17 00:00:00 2001
From: Matan Lurey <matanlurey@users.noreply.github.com>
Date: Tue, 20 Feb 2024 15:33:19 -0800
Subject: [PATCH 4/6] Treat `flutter-gold` failures in `flutter/engine` as
 failures. (#3487)

Closes https://github.com/flutter/flutter/issues/143352.

In practice this reverts the functionality in https://github.com/flutter/flutter/issues/123979 (https://github.com/flutter/cocoon/pull/2585). Sorry for the back-and-forth, we tried this model but ended up with (currently) [~170 untriaged failures](https://flutter-engine-gold.skia.org/), so we want to go back to treating these as failures and improve our tools/process if needed instead.

/cc @jonahwilliams and @Piinks
---
 auto_submit/lib/validations/ci_successful.dart       | 11 ++++-------
 auto_submit/test/validations/ci_successful_test.dart | 12 ++++++------
 2 files changed, 10 insertions(+), 13 deletions(-)

diff --git a/auto_submit/lib/validations/ci_successful.dart b/auto_submit/lib/validations/ci_successful.dart
index 27f8c8774..8cccdd597 100644
--- a/auto_submit/lib/validations/ci_successful.dart
+++ b/auto_submit/lib/validations/ci_successful.dart
@@ -145,12 +145,6 @@ class CiSuccessful extends Validation {
       // How can name be null but presumed to not be null below when added to failure?
       final String? name = status.context;
 
-      // If the account author is a roller account do not block merge on flutter-gold check.
-      if (isToEngineRoller(author, slug) && name == 'flutter-gold') {
-        log.info('Skipping status check for flutter-gold for ${slug.fullName}/$prNumber, pr author: $author.');
-        continue;
-      }
-
       if (status.state != STATUS_SUCCESS) {
         if (notInAuthorsControl.contains(name) && labelNames.contains(overrideTreeStatusLabel)) {
           continue;
@@ -159,7 +153,10 @@ class CiSuccessful extends Validation {
         if (status.state == STATUS_FAILURE && !notInAuthorsControl.contains(name)) {
           failures.add(FailureDetail(name!, status.targetUrl!));
         }
-        if (status.state == STATUS_PENDING && isStale(status.createdAt!) && supportStale(author, slug)) {
+        if (status.state == STATUS_PENDING &&
+            status.createdAt != null &&
+            isStale(status.createdAt!) &&
+            supportStale(author, slug)) {
           staleStatuses.add(status);
         }
       }
diff --git a/auto_submit/test/validations/ci_successful_test.dart b/auto_submit/test/validations/ci_successful_test.dart
index e325118b2..1935677b3 100644
--- a/auto_submit/test/validations/ci_successful_test.dart
+++ b/auto_submit/test/validations/ci_successful_test.dart
@@ -318,7 +318,7 @@ void main() {
       expect(failures.length, 2);
     });
 
-    test('Validate flutter-gold is not checked for engine auto roller pull requests.', () {
+    test('Validate flutter-gold is checked for engine auto roller pull requests.', () {
       final List<ContextNode> contextNodeList = getContextNodeListFromJson(repositoryStatusesWithGoldMock);
       const bool allSuccess = true;
       final Author author = Author(login: 'skia-flutter-autoroll');
@@ -331,13 +331,13 @@ void main() {
       convertContextNodeStatuses(contextNodeList);
       expect(
         ciSuccessful.validateStatuses(slug, prNumber, author, labelNames, contextNodeList, failures, allSuccess),
-        isTrue,
+        isFalse,
       );
       expect(failures, isEmpty);
       expect(failures.length, 0);
     });
 
-    test('Validate flutter-gold is not checked even if failing for engine auto roller pull requests.', () {
+    test('Validate flutter-gold is checked even if failing for engine auto roller pull requests.', () {
       final List<ContextNode> contextNodeList = getContextNodeListFromJson(repositoryStatusesWithFailedGoldMock);
       const bool allSuccess = true;
       final Author author = Author(login: 'skia-flutter-autoroll');
@@ -350,10 +350,10 @@ void main() {
       convertContextNodeStatuses(contextNodeList);
       expect(
         ciSuccessful.validateStatuses(slug, prNumber, author, labelNames, contextNodeList, failures, allSuccess),
-        isTrue,
+        isFalse,
       );
-      expect(failures, isEmpty);
-      expect(failures.length, 0);
+      expect(failures, isNotEmpty);
+      expect(failures.length, 1);
     });
 
     test('Validate flutter-gold is checked for non engine auto roller pull requests.', () {

From 0e74e8c2b5e3d72023cf96a7df704f53ad9a73a6 Mon Sep 17 00:00:00 2001
From: keyonghan <54558023+keyonghan@users.noreply.github.com>
Date: Wed, 21 Feb 2024 09:54:01 -0800
Subject: [PATCH 5/6] Update Firestore task when mulitple reruns exist (#3506)

This is a follow up of https://github.com/flutter/cocoon/pull/3495.

This PR:
1) updates the latest task status correctly when multiple records exist
2) uses const variables to avoid potential typos

Part of https://github.com/flutter/flutter/issues/142951
---
 app_dart/lib/src/model/firestore/task.dart    | 72 +++++++++--------
 .../request_handlers/update_task_status.dart  |  7 +-
 app_dart/lib/src/service/firestore.dart       | 77 +++++++++++++++----
 .../update_task_status_test.dart              | 43 ++++++-----
 app_dart/test/service/firestore_test.dart     | 34 ++++----
 app_dart/test/src/utilities/mocks.mocks.dart  | 61 ++++++++-------
 6 files changed, 180 insertions(+), 114 deletions(-)

diff --git a/app_dart/lib/src/model/firestore/task.dart b/app_dart/lib/src/model/firestore/task.dart
index 527584535..773e8d70a 100644
--- a/app_dart/lib/src/model/firestore/task.dart
+++ b/app_dart/lib/src/model/firestore/task.dart
@@ -2,6 +2,7 @@
 // Use of this source code is governed by a BSD-style license that can be
 // found in the LICENSE file.
 
+import 'package:cocoon_service/cocoon_service.dart';
 import 'package:googleapis/firestore/v1.dart' hide Status;
 
 import '../../request_handling/exceptions.dart';
@@ -77,27 +78,27 @@ class Task extends Document {
   ///
   /// This is _not_ when the task first started running, as tasks start out in
   /// the 'New' state until they've been picked up by an [Agent].
-  int? get createTimestamp => int.parse(fields!['createTimestamp']!.integerValue!);
+  int? get createTimestamp => int.parse(fields![kTaskCreateTimestampField]!.integerValue!);
 
   /// The timestamp (in milliseconds since the Epoch) that this task started
   /// running.
   ///
   /// Tasks may be run more than once. If this task has been run more than
   /// once, this timestamp represents when the task was most recently started.
-  int? get startTimestamp => int.parse(fields!['startTimestamp']!.integerValue!);
+  int? get startTimestamp => int.parse(fields![kTaskStartTimestampField]!.integerValue!);
 
   /// The timestamp (in milliseconds since the Epoch) that this task last
   /// finished running.
-  int? get endTimestamp => int.parse(fields!['endTimestamp']!.integerValue!);
+  int? get endTimestamp => int.parse(fields![kTaskEndTimestampField]!.integerValue!);
 
   /// The name of the task.
   ///
   /// This is a human-readable name, typically a test name (e.g.
   /// "hello_world__memory").
-  String? get taskName => fields!['name']!.stringValue!;
+  String? get taskName => fields![kTaskNameField]!.stringValue!;
 
   /// The sha of the task commit.
-  String? get commitSha => fields!['commitSha']!.stringValue!;
+  String? get commitSha => fields![kTaskCommitShaField]!.stringValue!;
 
   /// The number of attempts that have been made to run this task successfully.
   ///
@@ -112,7 +113,7 @@ class Task extends Document {
   ///  * <https://github.com/flutter/flutter/blob/master/.ci.yaml>
   ///
   /// A flaky (`bringup: true`) task will not block the tree.
-  bool? get bringup => fields!['bringup']!.booleanValue!;
+  bool? get bringup => fields![kTaskBringupField]!.booleanValue!;
 
   /// Whether the test execution of this task shows flake.
   ///
@@ -120,16 +121,17 @@ class Task extends Document {
   ///
   /// See also:
   ///  * <https://github.com/flutter/flutter/blob/master/dev/devicelab/lib/framework/runner.dart>
-  bool? get testFlaky => fields!['testFlaky']!.booleanValue!;
+  bool? get testFlaky => fields![kTaskTestFlakyField]!.booleanValue!;
 
   /// The build number of luci build: https://chromium.googlesource.com/infra/luci/luci-go/+/master/buildbucket/proto/build.proto#146
-  int? get buildNumber => fields!.containsKey('buildNumber') ? int.parse(fields!['buildNumber']!.integerValue!) : null;
+  int? get buildNumber =>
+      fields!.containsKey(kTaskBuildNumberField) ? int.parse(fields![kTaskBuildNumberField]!.integerValue!) : null;
 
   /// The status of the task.
   ///
   /// Legal values and their meanings are defined in [legalStatusValues].
   String get status {
-    final String taskStatus = fields!['status']!.stringValue!;
+    final String taskStatus = fields![kTaskStatusField]!.stringValue!;
     if (!legalStatusValues.contains(taskStatus)) {
       throw ArgumentError('Invalid state: "$taskStatus"');
     }
@@ -140,16 +142,16 @@ class Task extends Document {
     if (!legalStatusValues.contains(value)) {
       throw ArgumentError('Invalid state: "$value"');
     }
-    fields!['status'] = Value(stringValue: value);
+    fields![kTaskStatusField] = Value(stringValue: value);
     return value;
   }
 
   void setEndTimestamp(int endTimestamp) {
-    fields!['endTimestamp'] = Value(integerValue: endTimestamp.toString());
+    fields![kTaskEndTimestampField] = Value(integerValue: endTimestamp.toString());
   }
 
   void setTestFlaky(bool testFlaky) {
-    fields!['testFlaky'] = Value(booleanValue: testFlaky);
+    fields![kTaskTestFlakyField] = Value(booleanValue: testFlaky);
   }
 
   /// Update [Task] fields based on a LUCI [Build].
@@ -161,25 +163,31 @@ class Task extends Document {
       log.warning('Tags: $tags');
       throw const BadRequestException('build_address does not contain build number');
     }
-    fields!['buildNumber'] = Value(integerValue: buildAddress.split('/').last);
-    fields!['createTimestamp'] = Value(integerValue: (build.createdTimestamp?.millisecondsSinceEpoch ?? 0).toString());
-    fields!['startTimestamp'] = Value(integerValue: (build.startedTimestamp?.millisecondsSinceEpoch ?? 0).toString());
-    fields!['endTimestamp'] = Value(integerValue: (build.completedTimestamp?.millisecondsSinceEpoch ?? 0).toString());
+    fields![kTaskBuildNumberField] = Value(integerValue: buildAddress.split('/').last);
+    fields![kTaskCreateTimestampField] = Value(
+      integerValue: (build.createdTimestamp?.millisecondsSinceEpoch ?? kTaskDefaultTimestampValue).toString(),
+    );
+    fields![kTaskStartTimestampField] = Value(
+      integerValue: (build.startedTimestamp?.millisecondsSinceEpoch ?? kTaskDefaultTimestampValue).toString(),
+    );
+    fields![kTaskEndTimestampField] = Value(
+      integerValue: (build.completedTimestamp?.millisecondsSinceEpoch ?? kTaskDefaultTimestampValue).toString(),
+    );
 
     _setStatusFromLuciStatus(build);
   }
 
   void resetAsRetry({int attempt = 1}) {
-    name = '$kDatabase/documents/tasks/${commitSha}_${taskName}_$attempt';
+    name = '$kDatabase/documents/$kTaskCollectionId/${commitSha}_${taskName}_$attempt';
     fields = <String, Value>{
-      'createTimestamp': Value(integerValue: DateTime.now().millisecondsSinceEpoch.toString()),
-      'endTimestamp': Value(integerValue: '0'),
-      'bringup': Value(booleanValue: bringup),
-      'name': Value(stringValue: taskName),
-      'startTimestamp': Value(integerValue: '0'),
-      'status': Value(stringValue: Task.statusNew),
-      'testFlaky': Value(booleanValue: false),
-      'commitSha': Value(stringValue: commitSha),
+      kTaskCreateTimestampField: Value(integerValue: DateTime.now().millisecondsSinceEpoch.toString()),
+      kTaskEndTimestampField: Value(integerValue: kTaskDefaultTimestampValue.toString()),
+      kTaskBringupField: Value(booleanValue: bringup),
+      kTaskNameField: Value(stringValue: taskName),
+      kTaskStartTimestampField: Value(integerValue: kTaskDefaultTimestampValue.toString()),
+      kTaskStatusField: Value(stringValue: Task.statusNew),
+      kTaskTestFlakyField: Value(booleanValue: false),
+      kTaskCommitShaField: Value(stringValue: commitSha),
     };
   }
 
@@ -226,13 +234,13 @@ class Task extends Document {
   String toString() {
     final StringBuffer buf = StringBuffer()
       ..write('$runtimeType(')
-      ..write(', createTimestamp: $createTimestamp')
-      ..write(', startTimestamp: $startTimestamp')
-      ..write(', endTimestamp: $endTimestamp')
-      ..write(', name: $name')
-      ..write(', bringup: $bringup')
-      ..write(', testRunFlaky: $testFlaky')
-      ..write(', status: $status')
+      ..write(', $kTaskCreateTimestampField: $createTimestamp')
+      ..write(', $kTaskStartTimestampField: $startTimestamp')
+      ..write(', $kTaskEndTimestampField: $endTimestamp')
+      ..write(', $kTaskNameField: $name')
+      ..write(', $kTaskBringupField: $bringup')
+      ..write(', $kTaskTestFlakyField: $testFlaky')
+      ..write(', $kTaskStatusField: $status')
       ..write(')');
     return buf.toString();
   }
diff --git a/app_dart/lib/src/request_handlers/update_task_status.dart b/app_dart/lib/src/request_handlers/update_task_status.dart
index afd41c855..34ebe4d0c 100644
--- a/app_dart/lib/src/request_handlers/update_task_status.dart
+++ b/app_dart/lib/src/request_handlers/update_task_status.dart
@@ -78,8 +78,11 @@ class UpdateTaskStatus extends ApiRequestHandler<UpdateTaskStatusResponse> {
     final String? taskName = requestData![builderNameParam] as String?;
     final String documentName = '$kDatabase/documents/tasks/${sha}_${taskName}_1';
     log.info('getting firestore document: $documentName');
-    final firestore.Task firestoreTask =
-        await firestore.Task.fromFirestore(firestoreService: firestoreService, documentName: documentName);
+    final List<firestore.Task> initialTasks = await firestoreService.queryCommitTasks(sha);
+    // Targets the latest task. This assumes only one task is running at any time.
+    final firestore.Task firestoreTask = initialTasks.where((firestore.Task task) => task.taskName == taskName).reduce(
+          (firestore.Task current, firestore.Task next) => current.name!.compareTo(next.name!) > 0 ? current : next,
+        );
     firestoreTask.setStatus(status);
     firestoreTask.setEndTimestamp(endTimestamp);
     firestoreTask.setTestFlaky(isTestFlaky);
diff --git a/app_dart/lib/src/service/firestore.dart b/app_dart/lib/src/service/firestore.dart
index 8494f62d3..50418fb96 100644
--- a/app_dart/lib/src/service/firestore.dart
+++ b/app_dart/lib/src/service/firestore.dart
@@ -4,16 +4,41 @@
 
 import 'dart:async';
 
+import 'package:cocoon_service/cocoon_service.dart';
 import 'package:googleapis/firestore/v1.dart';
 import 'package:http/http.dart';
 
 import '../model/appengine/commit.dart';
 import '../model/appengine/task.dart';
+import '../model/firestore/task.dart' as firestore;
 import '../model/ci_yaml/target.dart';
 import 'access_client_provider.dart';
 import 'config.dart';
 
 const String kDatabase = 'projects/${Config.flutterGcpProjectId}/databases/${Config.flutterGcpFirestoreDatabase}';
+const String kDocumentParent = '$kDatabase/documents';
+const String kFieldFilterOpEqual = 'EQUAL';
+
+const String kTaskCollectionId = 'tasks';
+const int kTaskDefaultTimestampValue = 0;
+const String kTaskBringupField = 'bringup';
+const String kTaskBuildNumberField = 'buildNumber';
+const String kTaskCommitShaField = 'commitSha';
+const String kTaskCreateTimestampField = 'createTimestamp';
+const String kTaskEndTimestampField = 'endTimestamp';
+const String kTaskNameField = 'name';
+const String kTaskStartTimestampField = 'startTimestamp';
+const String kTaskStatusField = 'status';
+const String kTaskTestFlakyField = 'testFlaky';
+
+const String kCommitCollectionId = 'commits';
+const String kCommitAvatarField = 'avatar';
+const String kCommitBranchField = 'branch';
+const String kCommitCreateTimestampField = 'createTimestamp';
+const String kCommitAuthorField = 'author';
+const String kCommitMessageField = 'message';
+const String kCommitRepositoryPathField = 'repositoryPath';
+const String kCommitShaField = 'sha';
 
 class FirestoreService {
   const FirestoreService(this.accessClientProvider);
@@ -65,22 +90,40 @@ class FirestoreService {
         CommitRequest(transaction: beginTransactionResponse.transaction, writes: writes);
     return databasesDocumentsResource.commit(commitRequest, kDatabase);
   }
+
+  Future<List<firestore.Task>> queryCommitTasks(String commitSha) async {
+    final ProjectsDatabasesDocumentsResource databasesDocumentsResource = await documentResource();
+    final List<CollectionSelector> from = <CollectionSelector>[CollectionSelector(collectionId: kTaskCollectionId)];
+    final Filter filter = Filter(
+      fieldFilter: FieldFilter(
+        field: FieldReference(fieldPath: kTaskCommitShaField),
+        op: kFieldFilterOpEqual,
+        value: Value(stringValue: commitSha),
+      ),
+    );
+    final RunQueryRequest runQueryRequest =
+        RunQueryRequest(structuredQuery: StructuredQuery(from: from, where: filter));
+    final List<RunQueryResponseElement> runQueryResponseElements =
+        await databasesDocumentsResource.runQuery(runQueryRequest, kDocumentParent);
+    final List<Document> documents = runQueryResponseElements.map((e) => e.document!).toList();
+    return documents.map((Document document) => firestore.Task.fromDocument(taskDocument: document)).toList();
+  }
 }
 
 /// Generates task documents based on targets.
 List<Document> targetsToTaskDocuments(Commit commit, List<Target> targets) {
   final Iterable<Document> iterableDocuments = targets.map(
     (Target target) => Document(
-      name: '$kDatabase/documents/tasks/${commit.sha}_${target.value.name}_1',
+      name: '$kDatabase/documents/$kTaskCollectionId/${commit.sha}_${target.value.name}_1',
       fields: <String, Value>{
-        'createTimestamp': Value(integerValue: commit.timestamp!.toString()),
-        'endTimestamp': Value(integerValue: '0'),
-        'bringup': Value(booleanValue: target.value.bringup),
-        'name': Value(stringValue: target.value.name.toString()),
-        'startTimestamp': Value(integerValue: '0'),
-        'status': Value(stringValue: Task.statusNew),
-        'testFlaky': Value(booleanValue: false),
-        'commitSha': Value(stringValue: commit.sha),
+        kTaskCreateTimestampField: Value(integerValue: commit.timestamp!.toString()),
+        kTaskEndTimestampField: Value(integerValue: kTaskDefaultTimestampValue.toString()),
+        kTaskBringupField: Value(booleanValue: target.value.bringup),
+        kTaskNameField: Value(stringValue: target.value.name.toString()),
+        kTaskStartTimestampField: Value(integerValue: kTaskDefaultTimestampValue.toString()),
+        kTaskStatusField: Value(stringValue: Task.statusNew),
+        kTaskTestFlakyField: Value(booleanValue: false),
+        kTaskCommitShaField: Value(stringValue: commit.sha),
       },
     ),
   );
@@ -90,15 +133,15 @@ List<Document> targetsToTaskDocuments(Commit commit, List<Target> targets) {
 /// Generates commit document based on datastore commit data model.
 Document commitToCommitDocument(Commit commit) {
   return Document(
-    name: '$kDatabase/documents/commits/${commit.sha}',
+    name: '$kDatabase/documents/$kCommitCollectionId/${commit.sha}',
     fields: <String, Value>{
-      'avatar': Value(stringValue: commit.authorAvatarUrl),
-      'branch': Value(stringValue: commit.branch),
-      'createTimestamp': Value(integerValue: commit.timestamp.toString()),
-      'author': Value(stringValue: commit.author),
-      'message': Value(stringValue: commit.message),
-      'repositoryPath': Value(stringValue: commit.repository),
-      'sha': Value(stringValue: commit.sha),
+      kCommitAvatarField: Value(stringValue: commit.authorAvatarUrl),
+      kCommitBranchField: Value(stringValue: commit.branch),
+      kCommitCreateTimestampField: Value(integerValue: commit.timestamp.toString()),
+      kCommitAuthorField: Value(stringValue: commit.author),
+      kCommitMessageField: Value(stringValue: commit.message),
+      kCommitRepositoryPathField: Value(stringValue: commit.repository),
+      kCommitShaField: Value(stringValue: commit.sha),
     },
   );
 }
diff --git a/app_dart/test/request_handlers/update_task_status_test.dart b/app_dart/test/request_handlers/update_task_status_test.dart
index 77a6bef70..b495ef36f 100644
--- a/app_dart/test/request_handlers/update_task_status_test.dart
+++ b/app_dart/test/request_handlers/update_task_status_test.dart
@@ -56,14 +56,15 @@ void main() {
     });
 
     test('TestFlaky is false when not injected', () async {
-      final firestore.Task firestoreTask = generateFirestoreTask(1);
+      final firestore.Task firestoreTask1 = generateFirestoreTask(1, name: 'linux_integration_ui_ios', attempts: 1);
+      final firestore.Task firestoreTask2 = generateFirestoreTask(2, name: 'linux_integration_ui_ios', attempts: 2);
       when(
-        mockFirestoreService.getDocument(
+        mockFirestoreService.queryCommitTasks(
           captureAny,
         ),
       ).thenAnswer((Invocation invocation) {
-        return Future<Document>.value(
-          firestoreTask,
+        return Future<List<firestore.Task>>.value(
+          <firestore.Task>[firestoreTask1, firestoreTask2],
         );
       });
       when(
@@ -95,18 +96,19 @@ void main() {
       await tester.post(handler);
 
       expect(task.isTestFlaky, false);
-      expect(firestoreTask.testFlaky, false);
+      expect(firestoreTask2.testFlaky, false);
     });
 
     test('TestFlaky is true when injected', () async {
-      final firestore.Task firestoreTask = generateFirestoreTask(1);
+      final firestore.Task firestoreTask1 = generateFirestoreTask(1, name: 'linux_integration_ui_ios', attempts: 1);
+      final firestore.Task firestoreTask2 = generateFirestoreTask(2, name: 'linux_integration_ui_ios', attempts: 2);
       when(
-        mockFirestoreService.getDocument(
+        mockFirestoreService.queryCommitTasks(
           captureAny,
         ),
       ).thenAnswer((Invocation invocation) {
-        return Future<Document>.value(
-          firestoreTask,
+        return Future<List<firestore.Task>>.value(
+          <firestore.Task>[firestoreTask1, firestoreTask2],
         );
       });
       when(
@@ -139,18 +141,18 @@ void main() {
       await tester.post(handler);
 
       expect(task.isTestFlaky, true);
-      expect(firestoreTask.testFlaky, true);
     });
 
     test('task name requests can update tasks', () async {
-      final firestore.Task firestoreTask = generateFirestoreTask(1);
+      final firestore.Task firestoreTask1 = generateFirestoreTask(1, name: 'linux_integration_ui_ios', attempts: 1);
+      final firestore.Task firestoreTask2 = generateFirestoreTask(2, name: 'linux_integration_ui_ios', attempts: 2);
       when(
-        mockFirestoreService.getDocument(
+        mockFirestoreService.queryCommitTasks(
           captureAny,
         ),
       ).thenAnswer((Invocation invocation) {
-        return Future<Document>.value(
-          firestoreTask,
+        return Future<List<firestore.Task>>.value(
+          <firestore.Task>[firestoreTask1, firestoreTask2],
         );
       });
       when(
@@ -183,8 +185,8 @@ void main() {
 
       expect(task.status, 'Failed');
       expect(task.attempts, 1);
-      expect(firestoreTask.status, 'Failed');
-      expect(firestoreTask.attempts, 1);
+      expect(firestoreTask2.status, 'Failed');
+      expect(firestoreTask2.attempts, 2);
     });
 
     test('task name requests when task does not exists returns exception', () async {
@@ -198,14 +200,15 @@ void main() {
     });
 
     test('task name request updates when input has whitespace', () async {
-      final firestore.Task firestoreTask = generateFirestoreTask(1);
+      final firestore.Task firestoreTask1 = generateFirestoreTask(1, name: 'linux_integration_ui_ios', attempts: 1);
+      final firestore.Task firestoreTask2 = generateFirestoreTask(2, name: 'linux_integration_ui_ios', attempts: 2);
       when(
-        mockFirestoreService.getDocument(
+        mockFirestoreService.queryCommitTasks(
           captureAny,
         ),
       ).thenAnswer((Invocation invocation) {
-        return Future<Document>.value(
-          firestoreTask,
+        return Future<List<firestore.Task>>.value(
+          <firestore.Task>[firestoreTask1, firestoreTask2],
         );
       });
       when(
diff --git a/app_dart/test/service/firestore_test.dart b/app_dart/test/service/firestore_test.dart
index 4d642671a..96bc50198 100644
--- a/app_dart/test/service/firestore_test.dart
+++ b/app_dart/test/service/firestore_test.dart
@@ -21,28 +21,28 @@ void main() {
     ];
     final List<Document> taskDocuments = targetsToTaskDocuments(commit, targets);
     expect(taskDocuments.length, 2);
-    expect(taskDocuments[0].name, '$kDatabase/documents/tasks/${commit.sha}_${targets[0].value.name}_1');
-    expect(taskDocuments[0].fields!['createTimestamp']!.integerValue, commit.timestamp.toString());
-    expect(taskDocuments[0].fields!['endTimestamp']!.integerValue, '0');
-    expect(taskDocuments[0].fields!['bringup']!.booleanValue, false);
-    expect(taskDocuments[0].fields!['name']!.stringValue, targets[0].value.name);
-    expect(taskDocuments[0].fields!['startTimestamp']!.integerValue, '0');
-    expect(taskDocuments[0].fields!['status']!.stringValue, Task.statusNew);
-    expect(taskDocuments[0].fields!['testFlaky']!.booleanValue, false);
-    expect(taskDocuments[0].fields!['commitSha']!.stringValue, commit.sha);
+    expect(taskDocuments[0].name, '$kDatabase/documents/$kTaskCollectionId/${commit.sha}_${targets[0].value.name}_1');
+    expect(taskDocuments[0].fields![kTaskCreateTimestampField]!.integerValue, commit.timestamp.toString());
+    expect(taskDocuments[0].fields![kTaskEndTimestampField]!.integerValue, '0');
+    expect(taskDocuments[0].fields![kTaskBringupField]!.booleanValue, false);
+    expect(taskDocuments[0].fields![kTaskNameField]!.stringValue, targets[0].value.name);
+    expect(taskDocuments[0].fields![kTaskStartTimestampField]!.integerValue, '0');
+    expect(taskDocuments[0].fields![kTaskStatusField]!.stringValue, Task.statusNew);
+    expect(taskDocuments[0].fields![kTaskTestFlakyField]!.booleanValue, false);
+    expect(taskDocuments[0].fields![kTaskCommitShaField]!.stringValue, commit.sha);
   });
 
   test('creates commit document correctly from commit data model', () async {
     final Commit commit = generateCommit(1);
     final Document commitDocument = commitToCommitDocument(commit);
-    expect(commitDocument.name, '$kDatabase/documents/commits/${commit.sha}');
-    expect(commitDocument.fields!['avatar']!.stringValue, commit.authorAvatarUrl);
-    expect(commitDocument.fields!['branch']!.stringValue, commit.branch);
-    expect(commitDocument.fields!['createTimestamp']!.integerValue, commit.timestamp.toString());
-    expect(commitDocument.fields!['author']!.stringValue, commit.author);
-    expect(commitDocument.fields!['message']!.stringValue, commit.message);
-    expect(commitDocument.fields!['repositoryPath']!.stringValue, commit.repository);
-    expect(commitDocument.fields!['sha']!.stringValue, commit.sha);
+    expect(commitDocument.name, '$kDatabase/documents/$kCommitCollectionId/${commit.sha}');
+    expect(commitDocument.fields![kCommitAvatarField]!.stringValue, commit.authorAvatarUrl);
+    expect(commitDocument.fields![kCommitBranchField]!.stringValue, commit.branch);
+    expect(commitDocument.fields![kCommitCreateTimestampField]!.integerValue, commit.timestamp.toString());
+    expect(commitDocument.fields![kCommitAuthorField]!.stringValue, commit.author);
+    expect(commitDocument.fields![kCommitMessageField]!.stringValue, commit.message);
+    expect(commitDocument.fields![kCommitRepositoryPathField]!.stringValue, commit.repository);
+    expect(commitDocument.fields![kCommitShaField]!.stringValue, commit.sha);
   });
 
   test('creates writes correctly from documents', () async {
diff --git a/app_dart/test/src/utilities/mocks.mocks.dart b/app_dart/test/src/utilities/mocks.mocks.dart
index 3948a8d1e..a2e8283d1 100644
--- a/app_dart/test/src/utilities/mocks.mocks.dart
+++ b/app_dart/test/src/utilities/mocks.mocks.dart
@@ -18,11 +18,11 @@ import 'package:cocoon_service/src/model/appengine/github_gold_status_update.dar
 import 'package:cocoon_service/src/model/appengine/key_helper.dart' as _i12;
 import 'package:cocoon_service/src/model/appengine/stage.dart' as _i34;
 import 'package:cocoon_service/src/model/appengine/task.dart' as _i33;
-import 'package:cocoon_service/src/model/ci_yaml/target.dart' as _i38;
-import 'package:cocoon_service/src/model/firestore/task.dart' as _i40;
-import 'package:cocoon_service/src/model/github/checks.dart' as _i39;
+import 'package:cocoon_service/src/model/ci_yaml/target.dart' as _i39;
+import 'package:cocoon_service/src/model/firestore/task.dart' as _i37;
+import 'package:cocoon_service/src/model/github/checks.dart' as _i40;
 import 'package:cocoon_service/src/model/luci/buildbucket.dart' as _i8;
-import 'package:cocoon_service/src/model/luci/push_message.dart' as _i37;
+import 'package:cocoon_service/src/model/luci/push_message.dart' as _i38;
 import 'package:cocoon_service/src/service/access_client_provider.dart' as _i5;
 import 'package:cocoon_service/src/service/access_token_provider.dart' as _i26;
 import 'package:cocoon_service/src/service/bigquery.dart' as _i16;
@@ -2591,6 +2591,15 @@ class MockFirestoreService extends _i1.Mock implements _i15.FirestoreService {
           ),
         )),
       ) as _i20.Future<_i21.CommitResponse>);
+
+  @override
+  _i20.Future<List<_i37.Task>> queryCommitTasks(String? commitSha) => (super.noSuchMethod(
+        Invocation.method(
+          #queryCommitTasks,
+          [commitSha],
+        ),
+        returnValue: _i20.Future<List<_i37.Task>>.value(<_i37.Task>[]),
+      ) as _i20.Future<List<_i37.Task>>);
 }
 
 /// A class which mocks [IssuesService].
@@ -3368,7 +3377,7 @@ class MockGithubChecksService extends _i1.Mock implements _i15.GithubChecksServi
 
   @override
   _i20.Future<bool> updateCheckStatus(
-    _i37.BuildPushMessage? buildPushMessage,
+    _i38.BuildPushMessage? buildPushMessage,
     _i15.LuciBuildService? luciBuildService,
     _i13.RepositorySlug? slug, {
     bool? rescheduled = false,
@@ -3387,7 +3396,7 @@ class MockGithubChecksService extends _i1.Mock implements _i15.GithubChecksServi
       ) as _i20.Future<bool>);
 
   @override
-  bool taskFailed(_i37.BuildPushMessage? buildPushMessage) => (super.noSuchMethod(
+  bool taskFailed(_i38.BuildPushMessage? buildPushMessage) => (super.noSuchMethod(
         Invocation.method(
           #taskFailed,
           [buildPushMessage],
@@ -3396,7 +3405,7 @@ class MockGithubChecksService extends _i1.Mock implements _i15.GithubChecksServi
       ) as bool);
 
   @override
-  int currentAttempt(_i37.Build? build) => (super.noSuchMethod(
+  int currentAttempt(_i38.Build? build) => (super.noSuchMethod(
         Invocation.method(
           #currentAttempt,
           [build],
@@ -3420,7 +3429,7 @@ class MockGithubChecksService extends _i1.Mock implements _i15.GithubChecksServi
       ) as String);
 
   @override
-  _i13.CheckRunConclusion conclusionForResult(_i37.Result? result) => (super.noSuchMethod(
+  _i13.CheckRunConclusion conclusionForResult(_i38.Result? result) => (super.noSuchMethod(
         Invocation.method(
           #conclusionForResult,
           [result],
@@ -3435,7 +3444,7 @@ class MockGithubChecksService extends _i1.Mock implements _i15.GithubChecksServi
       ) as _i13.CheckRunConclusion);
 
   @override
-  _i13.CheckRunStatus statusForResult(_i37.Status? status) => (super.noSuchMethod(
+  _i13.CheckRunStatus statusForResult(_i38.Status? status) => (super.noSuchMethod(
         Invocation.method(
           #statusForResult,
           [status],
@@ -6379,8 +6388,8 @@ class MockLuciBuildService extends _i1.Mock implements _i15.LuciBuildService {
       ) as _i20.Future<Iterable<_i8.Build>>);
 
   @override
-  _i20.Future<List<_i38.Target>> scheduleTryBuilds({
-    required List<_i38.Target>? targets,
+  _i20.Future<List<_i39.Target>> scheduleTryBuilds({
+    required List<_i39.Target>? targets,
     required _i13.PullRequest? pullRequest,
     _i30.CheckSuiteEvent? checkSuiteEvent,
   }) =>
@@ -6394,8 +6403,8 @@ class MockLuciBuildService extends _i1.Mock implements _i15.LuciBuildService {
             #checkSuiteEvent: checkSuiteEvent,
           },
         ),
-        returnValue: _i20.Future<List<_i38.Target>>.value(<_i38.Target>[]),
-      ) as _i20.Future<List<_i38.Target>>);
+        returnValue: _i20.Future<List<_i39.Target>>.value(<_i39.Target>[]),
+      ) as _i20.Future<List<_i39.Target>>);
 
   @override
   _i20.Future<void> cancelBuilds(
@@ -6417,7 +6426,7 @@ class MockLuciBuildService extends _i1.Mock implements _i15.LuciBuildService {
   @override
   _i20.Future<List<_i8.Build?>> failedBuilds(
     _i13.PullRequest? pullRequest,
-    List<_i38.Target>? targets,
+    List<_i39.Target>? targets,
   ) =>
       (super.noSuchMethod(
         Invocation.method(
@@ -6433,7 +6442,7 @@ class MockLuciBuildService extends _i1.Mock implements _i15.LuciBuildService {
   @override
   _i20.Future<_i8.Build> rescheduleBuild({
     required String? builderName,
-    required _i37.BuildPushMessage? buildPushMessage,
+    required _i38.BuildPushMessage? buildPushMessage,
     required int? rescheduleAttempt,
   }) =>
       (super.noSuchMethod(
@@ -6461,7 +6470,7 @@ class MockLuciBuildService extends _i1.Mock implements _i15.LuciBuildService {
       ) as _i20.Future<_i8.Build>);
 
   @override
-  _i20.Future<_i8.Build> reschedulePresubmitBuildUsingCheckRunEvent(_i39.CheckRunEvent? checkRunEvent) =>
+  _i20.Future<_i8.Build> reschedulePresubmitBuildUsingCheckRunEvent(_i40.CheckRunEvent? checkRunEvent) =>
       (super.noSuchMethod(
         Invocation.method(
           #reschedulePresubmitBuildUsingCheckRunEvent,
@@ -6491,10 +6500,10 @@ class MockLuciBuildService extends _i1.Mock implements _i15.LuciBuildService {
 
   @override
   _i20.Future<_i8.Build> reschedulePostsubmitBuildUsingCheckRunEvent(
-    _i39.CheckRunEvent? checkRunEvent, {
+    _i40.CheckRunEvent? checkRunEvent, {
     required _i32.Commit? commit,
     required _i33.Task? task,
-    required _i38.Target? target,
+    required _i39.Target? target,
   }) =>
       (super.noSuchMethod(
         Invocation.method(
@@ -6559,9 +6568,9 @@ class MockLuciBuildService extends _i1.Mock implements _i15.LuciBuildService {
       ) as _i20.Future<Set<String>>);
 
   @override
-  _i20.Future<List<_i15.Tuple<_i38.Target, _i33.Task, int>>> schedulePostsubmitBuilds({
+  _i20.Future<List<_i15.Tuple<_i39.Target, _i33.Task, int>>> schedulePostsubmitBuilds({
     required _i32.Commit? commit,
-    required List<_i15.Tuple<_i38.Target, _i33.Task, int>>? toBeScheduled,
+    required List<_i15.Tuple<_i39.Target, _i33.Task, int>>? toBeScheduled,
   }) =>
       (super.noSuchMethod(
         Invocation.method(
@@ -6572,14 +6581,14 @@ class MockLuciBuildService extends _i1.Mock implements _i15.LuciBuildService {
             #toBeScheduled: toBeScheduled,
           },
         ),
-        returnValue: _i20.Future<List<_i15.Tuple<_i38.Target, _i33.Task, int>>>.value(
-            <_i15.Tuple<_i38.Target, _i33.Task, int>>[]),
-      ) as _i20.Future<List<_i15.Tuple<_i38.Target, _i33.Task, int>>>);
+        returnValue: _i20.Future<List<_i15.Tuple<_i39.Target, _i33.Task, int>>>.value(
+            <_i15.Tuple<_i39.Target, _i33.Task, int>>[]),
+      ) as _i20.Future<List<_i15.Tuple<_i39.Target, _i33.Task, int>>>);
 
   @override
   _i20.Future<void> createPostsubmitCheckRun(
     _i32.Commit? commit,
-    _i38.Target? target,
+    _i39.Target? target,
     Map<String, dynamic>? rawUserData,
   ) =>
       (super.noSuchMethod(
@@ -6598,13 +6607,13 @@ class MockLuciBuildService extends _i1.Mock implements _i15.LuciBuildService {
   @override
   _i20.Future<bool> checkRerunBuilder({
     required _i32.Commit? commit,
-    required _i38.Target? target,
+    required _i39.Target? target,
     required _i33.Task? task,
     required _i9.DatastoreService? datastore,
     _i15.FirestoreService? firestoreService,
     Map<String, List<String>>? tags,
     bool? ignoreChecks = false,
-    _i40.Task? taskDocument,
+    _i37.Task? taskDocument,
   }) =>
       (super.noSuchMethod(
         Invocation.method(

From 1c56d9e25fbe13e2f51cd1bb42bf1a0b326c2c8a Mon Sep 17 00:00:00 2001
From: Matan Lurey <matanlurey@users.noreply.github.com>
Date: Wed, 21 Feb 2024 10:14:53 -0800
Subject: [PATCH 6/6] Update status checks for Gold Status -> Github (#3508)

These checks are out of date:

- It's `web_engine` not `web engine`
- We're missing Impeller, Scenario App goldens for the native engine entirely

/cc @mdebbar as well.

I also filed https://github.com/flutter/flutter/issues/143863 since it's likely we'll drift again in the future.
---
 .../push_gold_status_to_github.dart           | 21 ++++++++++++++-----
 .../push_gold_status_to_github_test.dart      |  4 ++--
 2 files changed, 18 insertions(+), 7 deletions(-)

diff --git a/app_dart/lib/src/request_handlers/push_gold_status_to_github.dart b/app_dart/lib/src/request_handlers/push_gold_status_to_github.dart
index d3fdcd9a6..52d6a54db 100644
--- a/app_dart/lib/src/request_handlers/push_gold_status_to_github.dart
+++ b/app_dart/lib/src/request_handlers/push_gold_status_to_github.dart
@@ -119,11 +119,22 @@ class PushGoldStatusToGithub extends ApiRequestHandler<Body> {
       for (Map<String, dynamic> checkRun in checkRuns) {
         log.fine('Check run: $checkRun');
         final String name = checkRun['name'].toLowerCase() as String;
-        // Framework shards run framework goldens
-        // Web shards run web version of framework goldens
-        // Misc shard runs API docs goldens
-        if (name.contains('framework') || name.contains('web engine') || name.contains('misc')) {
-          runsGoldenFileTests = true;
+        if (slug == Config.engineSlug) {
+          if (const <String>[
+            'linux_android_emulator',
+            'linux_host_engine',
+            'mac_host_engine',
+            'linux_web_engine',
+          ].any((String shardSubString) => name.contains(shardSubString))) {
+            runsGoldenFileTests = true;
+          }
+        } else if (slug == Config.flutterSlug) {
+          if (const <String>[
+            'framework',
+            'misc',
+          ].any((String shardSubString) => name.contains(shardSubString))) {
+            runsGoldenFileTests = true;
+          }
         }
         if (checkRun['conclusion'] == null || checkRun['conclusion'].toUpperCase() != 'SUCCESS') {
           incompleteChecks.add(name);
diff --git a/app_dart/test/request_handlers/push_gold_status_to_github_test.dart b/app_dart/test/request_handlers/push_gold_status_to_github_test.dart
index 80eb75882..261cf28da 100644
--- a/app_dart/test/request_handlers/push_gold_status_to_github_test.dart
+++ b/app_dart/test/request_handlers/push_gold_status_to_github_test.dart
@@ -767,7 +767,7 @@ void main() {
             <String, String>{'name': 'framework', 'status': 'in_progress', 'conclusion': 'neutral'},
           ];
           engineCheckRuns = <dynamic>[
-            <String, String>{'name': 'web engine', 'status': 'in_progress', 'conclusion': 'neutral'},
+            <String, String>{'name': 'Linux linux_web_engine', 'status': 'in_progress', 'conclusion': 'neutral'},
           ];
 
           final Body body = await tester.get<Body>(handler);
@@ -1482,7 +1482,7 @@ void main() {
           <String, String>{'name': 'framework', 'status': 'completed', 'conclusion': 'success'},
         ];
         engineCheckRuns = <dynamic>[
-          <String, String>{'name': 'web engine', 'status': 'completed', 'conclusion': 'success'},
+          <String, String>{'name': 'Linux linux_web_engine', 'status': 'completed', 'conclusion': 'success'},
         ];
 
         // Requests sent to Gold.