Fail safe error handling

commanddash/lib/agent/agent_handler.dart

@@ -50,39 +50,49 @@ class AgentHandler {
     );
   }
 
-  void runTask(TaskAssist taskAssist) async {
+  Future<void> runTask(TaskAssist taskAssist) async {
     DashRepository? dashRepository;
     if (githubAccessToken != null) {
       dashRepository = DashRepository.fromKeys(githubAccessToken!, taskAssist);
     }
     try {
       for (Map<String, dynamic> stepJson in steps) {
-        final step =
-            Step.fromJson(stepJson, inputs, outputs, agentName, agentVersion);
-        final results =
-            await step.run(taskAssist, generationRepository, dashRepository) ??
-                [];
-        if (step.outputIds != null) {
-          if (results.isEmpty) {
-            taskAssist.sendErrorMessage(
-                message:
-                    'No output received from the step where output was expected.',
-                data: {});
-          }
-          for (int i = 0; i < results.length; i++) {
-            final output = results[i];
-            if (output is ContinueToNextStepOutput && !output.value) {
-              break;
-            } else {
-              outputs[step.outputIds![i]] = output;
+        try {
+          final step =
+              Step.fromJson(stepJson, inputs, outputs, agentName, agentVersion);
+          final results = await step.run(
+                  taskAssist, generationRepository, dashRepository) ??
+              [];
+          if (step.outputIds != null) {
+            if (results.isEmpty) {
+              taskAssist.sendErrorMessage(
+                  message:
+                      'No output received from the step where output was expected.',
+                  data: {});
+            }
+            for (int i = 0; i < results.length; i++) {
+              final output = results[i];
+              if (output is ContinueToNextStepOutput && !output.value) {
+                break;
+              } else {
+                outputs[step.outputIds![i]] = output;
+              }
             }
           }
+        } catch (e, stackTrace) {
+          taskAssist.sendErrorMessage(
+              message:
+                  'Error processing step ${stepJson['type']}: ${e.toString()}',
+              data: {},
+              stackTrace: stackTrace);
         }
       }
-      taskAssist.sendResultMessage(message: "TASK_COMPLETE", data: outputs);
+      taskAssist.sendResultMessage(message: "TASK_COMPLETE", data: {});
     } catch (e, stackTrace) {
       taskAssist.sendErrorMessage(
-          message: e.toString(), data: {}, stackTrace: stackTrace);
+          message: 'Error processing request: ${e.toString()}',
+          data: {},
+          stackTrace: stackTrace);
     }
   }
 }

commanddash/lib/agent/step_model.dart

@@ -26,7 +26,6 @@ abstract class Step {
 
   factory Step.fromJson(Map<String, dynamic> json, Map<String, Input> inputs,
       Map<String, Output> outputs, String agentName, String agentVersion) {
-    // TODO: handle parsing error
     switch (json['type']) {
       case 'search_in_workspace':
         return SearchInWorkspaceStep.fromJson(
@@ -88,7 +87,10 @@ abstract class Step {
   Future<List<Output>?> run(
       TaskAssist taskAssist, GenerationRepository generationRepository,
       [DashRepository? dashRepository]) async {
-    await taskAssist.processStep(kind: 'loader_update', args: loader.toJson());
+    await taskAssist.processStep(
+        kind: 'loader_update',
+        args: loader.toJson(),
+        timeoutKind: TimeoutKind.sync);
     return null;
   }
 }

commanddash/lib/server/task_assist.dart

@@ -1,20 +1,55 @@
+import 'dart:async';
+
 import 'package:commanddash/server/messages.dart';
 import 'package:commanddash/server/operation_message.dart';
 import 'package:commanddash/server/server.dart';
 import 'package:rxdart/rxdart.dart';
 
+enum TimeoutKind {
+  /// 6s. Use for almost synchronous tasks like fetching local data.
+  sync,
+
+  /// 60s. Use for finite API call related tasks, like refreshing accessToken.
+  async,
+
+  /// 6 mins. Use for long running tasks like LLM calls (if any at client level).
+  stretched
+}
+
+Duration durationFromTimeoutKind(TimeoutKind kind) {
+  switch (kind) {
+    case TimeoutKind.sync:
+      return Duration(seconds: 6);
+    case TimeoutKind.async:
+      return Duration(seconds: 60);
+    case TimeoutKind.stretched:
+      return Duration(minutes: 6);
+  }
+}
+
 class TaskAssist {
   TaskAssist(this._server, this._taskId);
   final Server _server;
   final int _taskId;
 
-  Future<Map<String, dynamic>> processStep(
-      {required String kind, required Map<String, dynamic> args}) async {
+  /// [timeoutKind] should be suitably added depending upon what is expected from the operation
+  Future<Map<String, dynamic>> processStep({
+    required String kind,
+    required Map<String, dynamic> args,
+    required TimeoutKind timeoutKind,
+  }) async {
     _server.sendMessage(StepMessage(_taskId, kind: kind, args: args));
     final dataResponse = await _server.messagesStream
         .whereType<StepResponseMessage>()
         .where((event) => event.id == _taskId)
-        .first;
+        .first
+        .timeout(durationFromTimeoutKind(timeoutKind), onTimeout: () {
+      throw TimeoutException('Step timed out fetching $kind');
+    });
+
+    if (dataResponse.data['result'] == 'error') {
+      throw Exception(dataResponse.data['message']);
+    }
     return dataResponse.data;
   }
 
@@ -38,16 +73,26 @@ class TaskAssist {
   }
 
   /// To perform task independent operationas such as asking client to refresh and provide new access token.
+  ///
+  /// [timeoutKind] should be suitably added depending upon what is expected from the operation
   Future<Map<String, dynamic>> processOperation({
     required String kind,
     required Map<String, dynamic> args,
+    TimeoutKind timeoutKind = TimeoutKind.async,
   }) async {
     _server.sendMessage(OperationMessage(kind: kind, args: args));
 
     final dataResponse = await _server.messagesStream
         .whereType<OperationResponseMessage>()
         .where((event) => event.kind == kind)
-        .first;
+        .first
+        .timeout(durationFromTimeoutKind(timeoutKind), onTimeout: () {
+      throw TimeoutException('Operation timed out fetching $kind');
+    });
+
+    if (dataResponse.data['result'] == 'error') {
+      throw Exception(dataResponse.data['message']);
+    }
     return dataResponse.data;
   }
 }

commanddash/lib/server/task_handler.dart

@@ -15,22 +15,43 @@ class TaskHandler {
         .listen((TaskStartMessage message) async {
       final taskAssist = TaskAssist(_server, message.id);
       switch (message.taskKind) {
+        case 'random_task_global_error':
+          throw Exception('Some unhandled exception not tracked to a task.');
         case 'random_task_with_step':
-          randomFunctionWithStep(taskAssist);
+          try {
+            await randomFunctionWithStep(taskAssist);
+          } catch (e, stackTrace) {
+            taskAssist.sendErrorMessage(
+                message: 'Error processing request: ${e.toString()}',
+                data: {},
+                stackTrace: stackTrace);
+          }
           break;
         case 'random_task_with_side_operation':
-          randomFunctionWithSideOperation(taskAssist);
+          try {
+            await randomFunctionWithSideOperation(taskAssist);
+          } catch (e, stackTrace) {
+            taskAssist.sendErrorMessage(
+                message: 'Error processing request: ${e.toString()}',
+                data: {},
+                stackTrace: stackTrace);
+          }
+
           break;
         case 'get-agents':
           final client = getClient(
               message.data['auth']['github_access_token'],
-              () async => taskAssist
-                  .processOperation(kind: 'refresh_access_token', args: {}));
+              () async => taskAssist.processOperation(
+                    kind: 'refresh_access_token',
+                    args: {},
+                  ));
           final repo = DashRepository(client);
           try {
             final agents = await repo.getAgents();
             taskAssist.sendResultMessage(
-                message: "Agent get successful", data: {"agents": agents});
+              message: "Agent get successful",
+              data: {"agents": agents},
+            );
           } catch (e, stackTrace) {
             taskAssist.sendErrorMessage(
                 message: "Failed getting agents.",
@@ -41,25 +62,26 @@ class TaskHandler {
         case 'refresh_token_test':
           final client = getClient(
               message.data['auth']['github_access_token'],
-              () async => taskAssist
-                  .processOperation(kind: 'refresh_access_token', args: {}));
+              () async => taskAssist.processOperation(
+                    kind: 'refresh_access_token',
+                    args: {},
+                  ));
           DashRepository(client);
 
           ///Other repositories using the backend client
           ///Pass this to the agent.
           break;
         case 'agent-execute':
-          final handler = AgentHandler.fromJson(message.data);
-          handler.runTask(taskAssist);
+          try {
+            final handler = AgentHandler.fromJson(message.data);
+            await handler.runTask(taskAssist);
+          } catch (e, stackTrace) {
+            taskAssist.sendErrorMessage(
+                message: 'Error processing request: ${e.toString()}',
+                data: {},
+                stackTrace: stackTrace);
+          }
           break;
-        // case 'find_closest_files':
-        //   EmbeddingGenerator().findClosesResults(
-        //     taskAssist,
-        //     message.data['query'],
-        //     message.data['workspacePath'],
-        //     GeminiRepository(message.data['apiKey']),
-        //   );
-        // break;
         default:
           taskAssist.sendErrorMessage(message: 'INVALID_TASK_KIND', data: {});
       }
@@ -69,7 +91,8 @@ class TaskHandler {
 
 /// Function for Integration Test of the step communication
 Future<void> randomFunctionWithStep(TaskAssist taskAssist) async {
-  final data = await taskAssist.processStep(kind: 'step_data_kind', args: {});
+  final data = await taskAssist.processStep(
+      kind: 'step_data_kind', args: {}, timeoutKind: TimeoutKind.sync);
   if (data['value'] == 'unique_value') {
     taskAssist.sendResultMessage(message: 'TASK_COMPLETED', data: {});
   } else {
@@ -82,6 +105,9 @@ Future<void> randomFunctionWithStep(TaskAssist taskAssist) async {
 
 /// Function for Integration Test of the side operation communication
 Future<void> randomFunctionWithSideOperation(TaskAssist taskAssist) async {
-  await taskAssist.processOperation(kind: 'operation_data_kind', args: {});
-  taskAssist.sendResultMessage(message: 'TASK_COMPLETED', data: {});
+  await taskAssist.processOperation(
+      kind: 'operation_data_kind', args: {}, timeoutKind: TimeoutKind.sync);
+  taskAssist.sendLogMessage(message: 'response received', data: {});
+  taskAssist
+      .sendResultMessage(message: 'TASK_COMPLETED', data: {'success': true});
 }

commanddash/lib/steps/append_to_chat/append_to_chat_step.dart

@@ -30,8 +30,10 @@ class AppendToChatStep extends Step {
       TaskAssist taskAssist, GenerationRepository generationRepository,
       [DashRepository? dashRepository]) async {
     await super.run(taskAssist, generationRepository);
-    final response = await taskAssist
-        .processStep(kind: 'append_to_chat', args: {'message': message});
+    final response = await taskAssist.processStep(
+        kind: 'append_to_chat',
+        args: {'message': message},
+        timeoutKind: TimeoutKind.sync);
     if (response['error'] != null) {
       throw Exception(response['error']['message']);
     }

commanddash/lib/steps/find_closest_files/search_in_workspace_step.dart

@@ -43,7 +43,8 @@ class SearchInWorkspaceStep extends Step {
       [DashRepository? dashRepository]) async {
     await super.run(taskAssist, generationRepository);
     final dartFiles = EmbeddingGenerator.getDartFiles(workspacePath);
-    final codeCacheHash = await taskAssist.processStep(kind: 'cache', args: {});
+    final codeCacheHash = await taskAssist.processStep(
+        kind: 'cache', args: {}, timeoutKind: TimeoutKind.sync);
     final filesToUpdate =
         EmbeddingGenerator.getFilesToUpdate(dartFiles, codeCacheHash);
     final embeddedFiles = await EmbeddingGenerator.updateEmbeddings(

commanddash/lib/steps/replace_in_file/replace_in_file_step.dart

@@ -28,10 +28,12 @@ class ReplaceInFileStep extends Step {
       TaskAssist taskAssist, GenerationRepository generationRepository,
       [DashRepository? dashRepository]) async {
     await super.run(taskAssist, generationRepository);
-    final response =
-        await taskAssist.processStep(kind: 'replace_in_file', args: {
-      'file': file.getReplaceFileJson(newContent),
-    });
+    final response = await taskAssist.processStep(
+        kind: 'replace_in_file',
+        args: {
+          'file': file.getReplaceFileJson(newContent),
+        },
+        timeoutKind: TimeoutKind.stretched);
     final userChoice = response['value'] as bool;
     if (userChoice == false && continueIfDeclined == false) {
       return [ContinueToNextStepOutput(false)];

commanddash/test/server/task_assist_test.dart

@@ -1,13 +1,16 @@
 import 'dart:async';
 
+import 'package:async/async.dart';
 import 'package:commanddash/server/messages.dart';
 import 'package:commanddash/server/operation_message.dart';
 import 'package:commanddash/server/server.dart';
 import 'package:commanddash/server/task_assist.dart';
+import 'package:commanddash/server/task_handler.dart';
 import 'package:mockito/annotations.dart';
 import 'package:mockito/mockito.dart';
 import 'package:test/test.dart';
 
+import '../test_utils.dart';
 import 'task_assist_test.mocks.dart';
 
 @GenerateMocks([Server])
@@ -67,4 +70,50 @@ void main() {
       },
     );
   }, timeout: Timeout(Duration(milliseconds: 100)));
+
+  group('Task Assist E2E Tests', () {
+    late Server server;
+    late TaskHandler handler;
+    late StreamController<IncomingMessage> messageStreamController;
+    late TestOutWrapper outwrapper;
+    setUp(() async {
+      await EnvReader.load();
+      server = Server();
+      messageStreamController = StreamController.broadcast();
+      outwrapper = TestOutWrapper();
+      server.replaceMessageStreamController(messageStreamController);
+      server.stdout = outwrapper;
+      handler = TaskHandler(server);
+    });
+    test('Flutter agent', () async {
+      handler.initProcessing();
+
+      messageStreamController.add(
+        IncomingMessage.fromJson({
+          "method": "task_start",
+          "params": {
+            "kind": "random_task_with_side_operation",
+            "data": <String, dynamic>{}
+          },
+          "id": 1
+        }),
+      );
+
+      final queue =
+          StreamQueue<OutgoingMessage>(outwrapper.outputStream.stream);
+      var result = await queue.next;
+      expect(result, isA<OperationMessage>());
+      expect(result.id, -1);
+      expect((result as OperationMessage).kind, 'operation_data_kind');
+
+      messageStreamController.add(IncomingMessage.fromJson({
+        "method": "operation_response",
+        "kind": "operation_data_kind",
+        "data": {"result": "success", "value": "unique_value"}
+      }));
+      result = await queue.next;
+      expect(result, isA<ResultMessage>());
+      expect(result.id, 1);
+    });
+  });
 }