Workflow Copilot: server update streaming with "cancel" ability (#4456)

2026-01-14 18:34:09 -07:00
parent 6b9ea59e67
commit 9cf1f87514
7 changed files with 462 additions and 170 deletions
--- a/skyvern-frontend/src/api/sse.ts
+++ b/skyvern-frontend/src/api/sse.ts
@@ -2,19 +2,65 @@ import { fetchEventSource } from "@microsoft/fetch-event-source";
 import type { CredentialGetter } from "@/api/AxiosClient";
 import { getRuntimeApiKey, runsApiBaseUrl } from "@/util/env";

-export type SseJsonPayload = Record<string, unknown>;
+export type SseMessageHandler<T> = (payload: T, eventName: string) => boolean;

-type SseClient = {
-  post: <T extends SseJsonPayload>(path: string, body: unknown) => Promise<T>;
+type SseStreamingOptions = {
+  signal?: AbortSignal;
 };

-export async function fetchJsonSse<T extends SseJsonPayload>(
+type SseClient = {
+  postStreaming: <T>(
+    path: string,
+    body: unknown,
+    onMessage: SseMessageHandler<T>,
+    options?: SseStreamingOptions,
+  ) => Promise<void>;
+};
+
+export async function fetchStreamingSse<T>(
  input: RequestInfo | URL,
  init: RequestInit,
-): Promise<T> {
+  onMessage: SseMessageHandler<T>,
+  options?: SseStreamingOptions,
+): Promise<void> {
  const controller = new AbortController();
+  const externalSignal = options?.signal;
+  let settled = false;
+  const resolveOnce = () => {
+    if (!settled) {
+      settled = true;
+      return true;
+    }
+    return false;
+  };
+  const onExternalAbort = () => {
+    controller.abort();
+  };
+  if (externalSignal) {
+    if (externalSignal.aborted) {
+      controller.abort();
+      return;
+    }
+    externalSignal.addEventListener("abort", onExternalAbort, { once: true });
+  }
  try {
-    const parsedPayload = await new Promise<T>((resolve, reject) => {
+    await new Promise<void>((resolve, reject) => {
+      const safeResolve = () => {
+        if (resolveOnce()) {
+          resolve();
+        }
+      };
+      const safeReject = (error: unknown) => {
+        if (controller.signal.aborted) {
+          safeResolve();
+          return;
+        }
+        if (!settled) {
+          settled = true;
+          reject(error);
+        }
+      };
+
      fetchEventSource(input instanceof URL ? input.toString() : input, {
        method: init.method,
        headers: init.headers as Record<string, string>,
@@ -26,25 +72,28 @@ export async function fetchJsonSse<T extends SseJsonPayload>(
          }
          try {
            const payload = JSON.parse(event.data) as T;
-            resolve(payload);
+            if (onMessage(payload, event.event)) {
+              safeResolve();
+            }
          } catch (error) {
-            reject(error);
+            safeReject(error);
          }
        },
        onerror: (error) => {
-          reject(error);
+          safeReject(error);
        },
        onopen: async (response) => {
          if (!response.ok) {
            const errorText = await response.text();
-            reject(new Error(errorText || "Failed to send request."));
+            safeReject(new Error(errorText || "Failed to send request."));
          }
        },
-      }).catch(reject);
+      }).catch(safeReject);
    });
-
-    return parsedPayload;
  } finally {
+    if (externalSignal) {
+      externalSignal.removeEventListener("abort", onExternalAbort);
+    }
    controller.abort();
  }
 }
@@ -73,14 +122,21 @@ export async function getSseClient(
  }

  return {
-    post: <T extends SseJsonPayload>(path: string, body: unknown) => {
-      return fetchJsonSse<T>(
+    postStreaming: <T>(
+      path: string,
+      body: unknown,
+      onMessage: SseMessageHandler<T>,
+      options?: SseStreamingOptions,
+    ) => {
+      return fetchStreamingSse<T>(
        `${runsApiBaseUrl.replace(/\/$/, "")}/${path.replace(/^\//, "")}`,
        {
          method: "POST",
          headers: requestHeaders,
          body: JSON.stringify(body),
        },
+        onMessage,
+        options,
      );
    },
  };
--- a/skyvern-frontend/src/routes/workflows/copilot/WorkflowCopilotButton.tsx
+++ b/skyvern-frontend/src/routes/workflows/copilot/WorkflowCopilotButton.tsx
--- a/skyvern-frontend/src/routes/workflows/copilot/WorkflowCopilotChat.tsx
+++ b/skyvern-frontend/src/routes/workflows/copilot/WorkflowCopilotChat.tsx
@@ -9,14 +9,26 @@ import { useWorkflowHasChangesStore } from "@/store/WorkflowHasChangesStore";
 import { WorkflowCreateYAMLRequest } from "@/routes/workflows/types/workflowYamlTypes";
 import { toast } from "@/components/ui/use-toast";
 import { getSseClient } from "@/api/sse";
+import {
+  WorkflowCopilotChatHistoryResponse,
+  WorkflowCopilotProcessingUpdate,
+  WorkflowCopilotStreamError,
+  WorkflowCopilotStreamResponse,
+  WorkflowCopilotChatSender,
+} from "./workflowCopilotTypes";

 interface ChatMessage {
  id: string;
-  sender: "ai" | "user";
+  sender: WorkflowCopilotChatSender;
  content: string;
  timestamp?: string;
 }

+type WorkflowCopilotSsePayload =
+  | WorkflowCopilotProcessingUpdate
+  | WorkflowCopilotStreamResponse
+  | WorkflowCopilotStreamError;
+
 const formatChatTimestamp = (value: string) => {
  let normalizedValue = value.replace(/\.(\d{3})\d*/, ".$1");
  if (!normalizedValue.endsWith("Z")) {
@@ -39,7 +51,9 @@ const MessageItem = memo(({ message }: { message: ChatMessage }) => {
        {message.sender === "ai" ? "AI" : "U"}
      </div>
      <div className="relative flex-1 rounded-lg bg-slate-800 p-3 pr-12">
-        <p className="text-sm text-slate-200">{message.content}</p>
+        <p className="whitespace-pre-wrap pr-3 text-sm text-slate-200">
+          {message.content}
+        </p>
        {message.timestamp ? (
          <span className="pointer-events-none absolute bottom-2 right-2 rounded bg-slate-900/70 px-1.5 py-0.5 text-[10px] text-slate-400">
            {formatChatTimestamp(message.timestamp)}
@@ -111,7 +125,10 @@ export function WorkflowCopilotChat({
  const [messages, setMessages] = useState<ChatMessage[]>([]);
  const [inputValue, setInputValue] = useState("");
  const [isLoading, setIsLoading] = useState(false);
+  const [processingStatus, setProcessingStatus] = useState<string>("");
  const [isLoadingHistory, setIsLoadingHistory] = useState(false);
+  const streamingAbortController = useRef<AbortController | null>(null);
+  const pendingMessageId = useRef<string | null>(null);
  const [workflowCopilotChatId, setWorkflowCopilotChatId] = useState<
    string | null
  >(null);
@@ -143,6 +160,7 @@ export function WorkflowCopilotChat({
  const credentialGetter = useCredentialGetter();
  const { workflowRunId, workflowPermanentId } = useParams();
  const messagesEndRef = useRef<HTMLDivElement>(null);
+  const textareaRef = useRef<HTMLTextAreaElement>(null);
  const { getSaveData } = useWorkflowHasChangesStore();
  const hasInitializedPosition = useRef(false);
  const hasScrolledOnLoad = useRef(false);
@@ -151,6 +169,18 @@ export function WorkflowCopilotChat({
    messagesEndRef.current?.scrollIntoView({ behavior });
  };

+  const adjustTextareaHeight = () => {
+    const textarea = textareaRef.current;
+    if (!textarea) return;
+
+    textarea.style.height = "auto";
+    textarea.style.height = `${Math.min(textarea.scrollHeight, 150)}px`;
+  };
+
+  useEffect(() => {
+    adjustTextareaHeight();
+  }, [inputValue]);
+
  const handleNewChat = () => {
    setMessages([]);
    setWorkflowCopilotChatId(null);
@@ -194,16 +224,12 @@ export function WorkflowCopilotChat({
      hasScrolledOnLoad.current = false;
      try {
        const client = await getClient(credentialGetter, "sans-api-v1");
-        const response = await client.get<{
-          workflow_copilot_chat_id: string | null;
-          chat_history: Array<{
-            sender: "ai" | "user";
-            content: string;
-            created_at: string;
-          }>;
-        }>("/workflow/copilot/chat-history", {
-          params: { workflow_permanent_id: workflowPermanentId },
-        });
+        const response = await client.get<WorkflowCopilotChatHistoryResponse>(
+          "/workflow/copilot/chat-history",
+          {
+            params: { workflow_permanent_id: workflowPermanentId },
+          },
+        );

        if (!isMounted) return;

@@ -233,6 +259,33 @@ export function WorkflowCopilotChat({
    };
  }, [credentialGetter, workflowPermanentId]);

+  useEffect(() => {
+    const handleKeyDown = (event: KeyboardEvent) => {
+      if (event.key !== "Escape" || !isOpen || !isLoading) {
+        return;
+      }
+      cancelSend();
+    };
+
+    window.addEventListener("keydown", handleKeyDown);
+    return () => {
+      window.removeEventListener("keydown", handleKeyDown);
+    };
+  }, [isLoading, isOpen]);
+
+  const cancelSend = async () => {
+    if (!streamingAbortController.current) return;
+
+    if (pendingMessageId.current) {
+      const messageId = pendingMessageId.current;
+      pendingMessageId.current = null;
+      setMessages((prev) => prev.filter((message) => message.id !== messageId));
+    }
+    setIsLoading(false);
+    setProcessingStatus("");
+    streamingAbortController.current?.abort();
+  };
+
  const handleSend = async () => {
    if (!inputValue.trim() || isLoading) return;
    if (!workflowPermanentId) {
@@ -251,10 +304,16 @@ export function WorkflowCopilotChat({
      content: inputValue,
    };

+    pendingMessageId.current = userMessageId;
    setMessages((prev) => [...prev, userMessage]);
    const messageContent = inputValue;
    setInputValue("");
    setIsLoading(true);
+    setProcessingStatus("Starting...");
+
+    const abortController = new AbortController();
+    streamingAbortController.current?.abort();
+    streamingAbortController.current = abortController;

    try {
      const saveData = getSaveData();
@@ -314,69 +373,97 @@ export function WorkflowCopilotChat({
        workflowYaml = convertToYAML(requestBody);
      }

-      const client = await getSseClient(credentialGetter);
-      const response = await client.post<{
-        workflow_copilot_chat_id?: string;
-        message?: string;
-        updated_workflow_yaml?: string | null;
-        request_time?: string;
-        response_time?: string;
-        error?: string;
-      }>("/workflow/copilot/chat-post", {
-        workflow_permanent_id: workflowPermanentId,
-        workflow_copilot_chat_id: workflowCopilotChatId,
-        workflow_run_id: workflowRunId,
-        message: messageContent,
-        workflow_yaml: workflowYaml,
-      });
+      const handleProcessingUpdate = (
+        payload: WorkflowCopilotProcessingUpdate,
+      ) => {
+        if (payload.status) {
+          setProcessingStatus(payload.status);
+        }

-      if (response.error) {
-        throw new Error(response.error);
-      }
+        const pendingId = pendingMessageId.current;
+        if (!pendingId || !payload.timestamp) {
+          return;
+        }

-      if (
-        !response.workflow_copilot_chat_id ||
-        !response.message ||
-        !response.request_time ||
-        !response.response_time
-      ) {
-        throw new Error("No response received.");
-      }
-
-      setWorkflowCopilotChatId(response.workflow_copilot_chat_id);
-
-      const aiMessage: ChatMessage = {
-        id: Date.now().toString(),
-        sender: "ai",
-        content: response.message || "I received your message.",
-        timestamp: response.response_time,
+        setMessages((prev) =>
+          prev.map((message) =>
+            message.id === pendingId
+              ? {
+                  ...message,
+                  timestamp: payload.timestamp,
+                }
+              : message,
+          ),
+        );
      };

-      setMessages((prev) => [
-        ...prev.map((message) =>
-          message.id === userMessageId
-            ? {
-                ...message,
-                timestamp: response.request_time,
-              }
-            : message,
-        ),
-        aiMessage,
-      ]);
+      const handleResponse = (response: WorkflowCopilotStreamResponse) => {
+        setWorkflowCopilotChatId(response.workflow_copilot_chat_id);

-      if (response.updated_workflow_yaml && onWorkflowUpdate) {
-        try {
-          onWorkflowUpdate(response.updated_workflow_yaml);
-        } catch (updateError) {
-          console.error("Failed to update workflow:", updateError);
-          toast({
-            title: "Update failed",
-            description: "Failed to apply workflow changes. Please try again.",
-            variant: "destructive",
-          });
+        const aiMessage: ChatMessage = {
+          id: Date.now().toString(),
+          sender: "ai",
+          content: response.message,
+          timestamp: response.response_time,
+        };
+
+        setMessages((prev) => [...prev, aiMessage]);
+
+        if (response.updated_workflow_yaml && onWorkflowUpdate) {
+          try {
+            onWorkflowUpdate(response.updated_workflow_yaml);
+          } catch (updateError) {
+            console.error("Failed to update workflow:", updateError);
+            toast({
+              title: "Update failed",
+              description:
+                "Failed to apply workflow changes. Please try again.",
+              variant: "destructive",
+            });
+          }
        }
-      }
+      };
+
+      const handleError = (payload: WorkflowCopilotStreamError) => {
+        const errorMessage: ChatMessage = {
+          id: Date.now().toString(),
+          sender: "ai",
+          content: payload.error,
+        };
+        setMessages((prev) => [...prev, errorMessage]);
+      };
+
+      const client = await getSseClient(credentialGetter);
+      await client.postStreaming<WorkflowCopilotSsePayload>(
+        "/workflow/copilot/chat-post",
+        {
+          workflow_permanent_id: workflowPermanentId,
+          workflow_copilot_chat_id: workflowCopilotChatId,
+          workflow_run_id: workflowRunId,
+          message: messageContent,
+          workflow_yaml: workflowYaml,
+        },
+        (payload) => {
+          switch (payload.type) {
+            case "processing_update":
+              handleProcessingUpdate(payload);
+              return false;
+            case "response":
+              handleResponse(payload);
+              return true;
+            case "error":
+              handleError(payload);
+              return true;
+            default:
+              return false;
+          }
+        },
+        { signal: abortController.signal },
+      );
    } catch (error) {
+      if (abortController.signal.aborted) {
+        return;
+      }
      console.error("Failed to send message:", error);
      const errorMessage: ChatMessage = {
        id: Date.now().toString(),
@@ -385,12 +472,18 @@ export function WorkflowCopilotChat({
      };
      setMessages((prev) => [...prev, errorMessage]);
    } finally {
+      if (streamingAbortController.current === abortController) {
+        streamingAbortController.current = null;
+      }
+      pendingMessageId.current = null;
      setIsLoading(false);
+      setProcessingStatus("");
    }
  };

-  const handleKeyPress = (e: React.KeyboardEvent<HTMLInputElement>) => {
-    if (e.key === "Enter") {
+  const handleKeyPress = (e: React.KeyboardEvent<HTMLTextAreaElement>) => {
+    if (e.key === "Enter" && !e.shiftKey) {
+      e.preventDefault();
      handleSend();
    }
  };
@@ -614,7 +707,7 @@ export function WorkflowCopilotChat({
              <div className="flex-1 rounded-lg bg-slate-800 p-3">
                <div className="flex items-center gap-2 text-sm text-slate-400">
                  <ReloadIcon className="h-4 w-4 animate-spin" />
-                  <span>Processing...</span>
+                  <span>{processingStatus || "Processing..."}</span>
                </div>
              </div>
            </div>
@@ -625,15 +718,21 @@ export function WorkflowCopilotChat({

      {/* Input */}
      <div className="border-t border-slate-700 p-3">
-        <div className="flex gap-2">
-          <input
-            type="text"
-            placeholder="Type your message..."
+        <div className="flex items-end gap-2">
+          <textarea
+            ref={textareaRef}
+            placeholder="Type your message... (Shift+Enter for new line)"
            value={inputValue}
            onChange={(e) => setInputValue(e.target.value)}
-            onKeyPress={handleKeyPress}
+            onKeyDown={handleKeyPress}
            disabled={isLoading}
-            className="flex-1 rounded-md border border-slate-600 bg-slate-800 px-3 py-2 text-sm text-slate-200 placeholder-slate-500 focus:border-blue-500 focus:outline-none disabled:cursor-not-allowed disabled:opacity-50"
+            rows={1}
+            className="flex-1 resize-none rounded-md border border-slate-600 bg-slate-800 px-3 py-2 text-sm text-slate-200 placeholder-slate-500 focus:border-blue-500 focus:outline-none disabled:cursor-not-allowed disabled:opacity-50"
+            style={{
+              minHeight: "38px",
+              maxHeight: "150px",
+              overflow: "auto",
+            }}
          />
          <button
            onClick={handleSend}
--- a/skyvern-frontend/src/routes/workflows/copilot/workflowCopilotTypes.ts
+++ b/skyvern-frontend/src/routes/workflows/copilot/workflowCopilotTypes.ts
@@ -0,0 +1,62 @@
+export type WorkflowCopilotChatSender = "user" | "ai";
+
+export interface WorkflowCopilotChat {
+  workflow_copilot_chat_id: string;
+  organization_id: string;
+  workflow_permanent_id: string;
+  created_at: string;
+  modified_at: string;
+}
+
+export interface WorkflowCopilotChatMessage {
+  workflow_copilot_chat_message_id: string;
+  workflow_copilot_chat_id: string;
+  sender: WorkflowCopilotChatSender;
+  content: string;
+  global_llm_context: string | null;
+  created_at: string;
+  modified_at: string;
+}
+
+export interface WorkflowCopilotChatRequest {
+  workflow_permanent_id: string;
+  workflow_copilot_chat_id?: string | null;
+  workflow_run_id?: string | null;
+  message: string;
+  workflow_yaml: string;
+}
+
+export interface WorkflowCopilotChatHistoryMessage {
+  sender: WorkflowCopilotChatSender;
+  content: string;
+  created_at: string;
+}
+
+export interface WorkflowCopilotChatHistoryResponse {
+  workflow_copilot_chat_id: string | null;
+  chat_history: WorkflowCopilotChatHistoryMessage[];
+}
+
+export type WorkflowCopilotStreamMessageType =
+  | "processing_update"
+  | "response"
+  | "error";
+
+export interface WorkflowCopilotProcessingUpdate {
+  type: "processing_update";
+  status: string;
+  timestamp: string;
+}
+
+export interface WorkflowCopilotStreamResponse {
+  type: "response";
+  workflow_copilot_chat_id: string;
+  message: string;
+  updated_workflow_yaml?: string | null;
+  response_time: string;
+}
+
+export interface WorkflowCopilotStreamError {
+  type: "error";
+  error: string;
+}
--- a/skyvern-frontend/src/routes/workflows/editor/Workspace.tsx
+++ b/skyvern-frontend/src/routes/workflows/editor/Workspace.tsx
@@ -103,8 +103,8 @@ import {
  upgradeWorkflowBlocksV1toV2,
 } from "./workflowEditorUtils";
 import { constructCacheKeyValue, getInitialParameters } from "./utils";
-import { WorkflowCopilotChat } from "./WorkflowCopilotChat";
-import { WorkflowCopilotButton } from "./WorkflowCopilotButton";
+import { WorkflowCopilotChat } from "../copilot/WorkflowCopilotChat";
+import { WorkflowCopilotButton } from "../copilot/WorkflowCopilotButton";
 import { parse as parseYAML } from "yaml";
 import {
  BlockYAML,