Implement STREAM_AICONFIG_CHUNK action

python/src/aiconfig/editor/client/src/LocalEditor.tsx

@@ -98,8 +98,8 @@ export default function Editor() {
           output_chunk: (data) => {
             onStream({ type: "output_chunk", data: data as Output });
           },
-          aiconfig: (data) => {
-            onStream({ type: "aiconfig", data: data as AIConfig });
+          aiconfig_chunk: (data) => {
+            onStream({ type: "aiconfig_chunk", data: data as AIConfig });
           },
           stop_streaming: (_data) => {
             onStream({ type: "stop_streaming", data: null });

python/src/aiconfig/editor/client/src/components/AIConfigEditor.tsx

@@ -69,7 +69,7 @@ export type RunPromptStreamEvent =
       data: Output;
     }
   | {
-      type: "aiconfig";
+      type: "aiconfig_chunk";
       data: AIConfig;
     }
   | {
@@ -636,18 +636,11 @@ export default function EditorContainer({
                 id: promptId,
                 output: event.data,
               });
-            } else if (event.type === "aiconfig") {
-              // Next PR: Change this to aiconfig_stream to make it more obvious
-              // and make STREAM_AICONFIG it's own event so we don't need to pass
-              // the `isRunning` state to set. See Ryan's comments about this in
+            } else if (event.type === "aiconfig_chunk") {
               dispatch({
-                type: "CONSOLIDATE_AICONFIG",
-                action: {
-                  ...action,
-                  // Keep the prompt running state until the end of streaming
-                  isRunning: true,
-                },
+                type: "STREAM_AICONFIG_CHUNK",
                 config: event.data,
+                cancellationToken,
               });
             } else if (event.type === "stop_streaming") {
               // Pass this event at the end of streaming to signal

python/src/aiconfig/editor/client/src/components/aiconfigReducer.ts

@@ -15,6 +15,7 @@ export type MutateAIConfigAction =
   | RunPromptAction
   | SetDescriptionAction
   | SetNameAction
+  | StreamAIConfigChunkAction
   | StreamOutputChunkAction
   | StopStreamingAction
   | UpdatePromptInputAction
@@ -24,9 +25,15 @@ export type MutateAIConfigAction =
   | UpdatePromptParametersAction
   | UpdateGlobalParametersAction;
 
+// Actions that appear when called via ConsolidateAIConfigAction
+export type ConsolidateAIConfigSubAction =
+  | AddPromptAction
+  | RunPromptAction
+  | UpdatePromptInputAction;
+
 export type ConsolidateAIConfigAction = {
   type: "CONSOLIDATE_AICONFIG";
-  action: MutateAIConfigAction;
+  action: ConsolidateAIConfigSubAction;
   config: AIConfig;
 };
 
@@ -72,6 +79,12 @@ export type SetNameAction = {
   name: string;
 };
 
+export type StreamAIConfigChunkAction = {
+  type: "STREAM_AICONFIG_CHUNK";
+  config: AIConfig;
+  cancellationToken?: string;
+};
+
 export type StreamOutputChunkAction = {
   type: "STREAM_OUTPUT_CHUNK";
   id: string;
@@ -159,7 +172,7 @@ function reduceInsertPromptAtIndex(
 
 function reduceConsolidateAIConfig(
   state: ClientAIConfig,
-  action: MutateAIConfigAction,
+  action: ConsolidateAIConfigSubAction,
   responseConfig: AIConfig
 ): ClientAIConfig {
   // Make sure prompt structure is properly updated. Client input and metadata takes precedence
@@ -186,20 +199,16 @@ function reduceConsolidateAIConfig(
         consolidatePrompt
       );
     }
+    // Next PR: Split "RUN_PROMPT" into two actions:
+    // 1) "RUN_PROMPT_START"
+    // 2) "RUN_PROMPT_SUCCESS"
+    // 3) (Already exists) "RUN_PROMPT_ERROR"
     case "RUN_PROMPT": {
-      // Note: If we are calling "RUN_PROMPT" directly as a dispatched event
-      // type, we automatically set the state there to `isRunning` for that
-      // prompt. That logic does not happen here, it happens in
-      // `aiconfigReducer`.
-      // If we are calling "RUN_PROMPT" indirectly via the action of a
-      // "CONSOLIDATE_AICONFIG" dispatch, we end up here. We need to check
-      // if we actually want to set the prompt state to `isRunning`
-      const isRunning = action.isRunning ?? false;
       const stateWithUpdatedRunningPromptId = {
         ...state,
         _ui: {
           ...state._ui,
-          runningPromptId: isRunning ? action.id : undefined,
+          runningPromptId: undefined,
         },
       };
       return reduceReplacePrompt(
@@ -216,7 +225,7 @@ function reduceConsolidateAIConfig(
             ...prompt,
             _ui: {
               ...prompt._ui,
-              isRunning,
+              isRunning: false,
             },
             outputs,
           };
@@ -340,6 +349,24 @@ export default function aiconfigReducer(
         name: action.name,
       };
     }
+    case "STREAM_AICONFIG_CHUNK": {
+      const replaceOutput = (statePrompt: ClientPrompt) => {
+        const responsePrompt = action.config.prompts.find(
+          (resPrompt) => resPrompt.name === statePrompt.name
+        );
+        return {
+          // Note: Don't need to set `isRunning` or `cancellationToken`
+          // because we already call RUN_PROMPT earlier in `onRunPrompt`
+          ...statePrompt,
+          outputs: responsePrompt?.outputs,
+        } as ClientPrompt;
+      };
+      return reduceReplacePrompt(
+        dirtyState,
+        dirtyState._ui.runningPromptId as string,
+        replaceOutput
+      );
+    }
     case "STREAM_OUTPUT_CHUNK": {
       return reduceReplacePrompt(dirtyState, action.id, (prompt) => ({
         ...prompt,

python/src/aiconfig/editor/server/server.py

@@ -348,7 +348,7 @@ def kill_thread(thread_id: int | None):
 
             aiconfig_json = aiconfig.model_dump(exclude=EXCLUDE_OPTIONS) if aiconfig is not None else None
             yield "["
-            yield json.dumps({"aiconfig": aiconfig_json})
+            yield json.dumps({"aiconfig_chunk": aiconfig_json})
             yield "]"
 
         yield "["