Frontend UX improvements and WebSocket fix

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "uipath-dev"
-version = "0.0.33"
+version = "0.0.34"
 description = "UiPath Developer Console"
 readme = { file = "README.md", content-type = "text/markdown" }
 requires-python = ">=3.11"

src/uipath/dev/server/frontend/src/api/websocket.ts

@@ -8,6 +8,8 @@ export class WsClient {
   private handlers: Set<MessageHandler> = new Set();
   private reconnectTimer: ReturnType<typeof setTimeout> | null = null;
   private shouldReconnect = true;
+  private pendingMessages: string[] = [];
+  private activeSubscriptions: Set<string> = new Set();
 
   constructor(url?: string) {
     const protocol = window.location.protocol === "https:" ? "wss:" : "ws:";
@@ -21,6 +23,15 @@ export class WsClient {
 
     this.ws.onopen = () => {
       console.log("[ws] connected");
+      // Re-subscribe to active subscriptions after reconnect
+      for (const runId of this.activeSubscriptions) {
+        this.sendRaw(JSON.stringify({ type: "subscribe", payload: { run_id: runId } }));
+      }
+      // Flush any messages queued while connecting
+      for (const msg of this.pendingMessages) {
+        this.sendRaw(msg);
+      }
+      this.pendingMessages = [];
     };
 
     this.ws.onmessage = (event) => {
@@ -56,17 +67,28 @@ export class WsClient {
     return () => this.handlers.delete(handler);
   }
 
+  private sendRaw(data: string): void {
+    if (this.ws?.readyState === WebSocket.OPEN) {
+      this.ws.send(data);
+    }
+  }
+
   send(type: ClientCommandType, payload: Record<string, unknown>): void {
+    const data = JSON.stringify({ type, payload });
     if (this.ws?.readyState === WebSocket.OPEN) {
-      this.ws.send(JSON.stringify({ type, payload }));
+      this.ws.send(data);
+    } else {
+      this.pendingMessages.push(data);
     }
   }
 
   subscribe(runId: string): void {
+    this.activeSubscriptions.add(runId);
     this.send("subscribe", { run_id: runId });
   }
 
   unsubscribe(runId: string): void {
+    this.activeSubscriptions.delete(runId);
     this.send("unsubscribe", { run_id: runId });
   }
 

src/uipath/dev/server/frontend/src/components/chat/ChatMessage.tsx

@@ -9,6 +9,8 @@ interface ChatMsg {
 
 interface Props {
   message: ChatMsg;
+  onToolCallClick?: (name: string, occurrenceIndex: number) => void;
+  toolCallIndices?: number[];
 }
 
 const ROLE_CONFIG: Record<string, { label: string; color: string }> = {
@@ -17,7 +19,7 @@ const ROLE_CONFIG: Record<string, { label: string; color: string }> = {
   assistant: { label: "AI", color: "var(--success)" },
 };
 
-export default function ChatMessage({ message }: Props) {
+export default function ChatMessage({ message, onToolCallClick, toolCallIndices }: Props) {
   const isUser = message.role === "user";
   const hasTool = message.tool_calls && message.tool_calls.length > 0;
   const roleKey = isUser ? "user" : hasTool ? "tool" : "assistant";
@@ -61,15 +63,16 @@ export default function ChatMessage({ message }: Props) {
       {/* Tool calls */}
       {message.tool_calls && message.tool_calls.length > 0 && (
         <div className="flex flex-wrap gap-1 mt-1 pl-2.5">
-          {message.tool_calls.map((tc) => (
+          {message.tool_calls.map((tc, i) => (
             <span
-              key={tc.name}
-              className="inline-flex items-center gap-1 text-[10px] font-mono px-1.5 py-0.5 rounded"
+              key={`${tc.name}-${i}`}
+              className="inline-flex items-center gap-1 text-[10px] font-mono px-1.5 py-0.5 rounded cursor-pointer hover:brightness-125"
               style={{
                 background: "var(--bg-primary)",
                 border: "1px solid var(--border)",
                 color: tc.has_result ? "var(--success)" : "var(--text-muted)",
               }}
+              onClick={() => onToolCallClick?.(tc.name, toolCallIndices?.[i] ?? 0)}
             >
               {tc.has_result ? "\u2713" : "\u2022"} {tc.name}
             </span>

src/uipath/dev/server/frontend/src/components/chat/ChatPanel.tsx

@@ -1,4 +1,4 @@
-import { useEffect, useRef } from "react";
+import { useEffect, useMemo, useRef } from "react";
 import type { WsClient } from "../../api/websocket";
 import { useRunStore } from "../../store/useRunStore";
 import ChatMessage from "./ChatMessage";
@@ -22,6 +22,25 @@ export default function ChatPanel({ messages, runId, runStatus, ws }: Props) {
   const scrollRef = useRef<HTMLDivElement>(null);
   const stickToBottom = useRef(true);
   const addLocalChatMessage = useRunStore((s) => s.addLocalChatMessage);
+  const setFocusedSpan = useRunStore((s) => s.setFocusedSpan);
+
+  // Precompute per-tool-call occurrence indices across all messages
+  const toolCallIndicesMap = useMemo(() => {
+    const map = new Map<string, number[]>();
+    const counts = new Map<string, number>();
+    for (const msg of messages) {
+      if (msg.tool_calls) {
+        const indices: number[] = [];
+        for (const tc of msg.tool_calls) {
+          const count = counts.get(tc.name) ?? 0;
+          indices.push(count);
+          counts.set(tc.name, count + 1);
+        }
+        map.set(msg.message_id, indices);
+      }
+    }
+    return map;
+  }, [messages]);
 
   // Track whether user has scrolled away from bottom
   const handleScroll = () => {
@@ -64,7 +83,12 @@ export default function ChatPanel({ messages, runId, runStatus, ws }: Props) {
           </p>
         )}
         {messages.map((msg) => (
-          <ChatMessage key={msg.message_id} message={msg} />
+          <ChatMessage
+            key={msg.message_id}
+            message={msg}
+            toolCallIndices={toolCallIndicesMap.get(msg.message_id)}
+            onToolCallClick={(name, idx) => setFocusedSpan({ name, index: idx })}
+          />
         ))}
       </div>
       <ChatInput

src/uipath/dev/server/frontend/src/components/runs/RunDetailsPanel.tsx

@@ -6,6 +6,7 @@ import GraphPanel from "../graph/GraphPanel";
 import TraceTree from "../traces/TraceTree";
 import LogPanel from "../logs/LogPanel";
 import ChatPanel from "../chat/ChatPanel";
+import JsonHighlight from "../shared/JsonHighlight";
 
 type Tab = "traces" | "output";
 
@@ -135,7 +136,7 @@ export default function RunDetailsPanel({ run, ws, activeTab, onTabChange }: Pro
   return (
     <div className="flex flex-col h-full">
       {/* Tab bar */}
-      <div className="flex items-center gap-1 px-2 py-1.5 border-b border-[var(--border)] bg-[var(--bg-primary)]">
+      <div className="flex items-center gap-1 px-2 py-2.5 border-b border-[var(--border)] bg-[var(--bg-primary)]">
         {tabs.map((tab) => (
           <button
             key={tab.id}
@@ -329,12 +330,11 @@ function IOView({ run }: { run: RunSummary }) {
     <div className="p-4 overflow-y-auto h-full space-y-4">
       {/* Input */}
       <DataSection title="Input" color="var(--success)" copyText={JSON.stringify(run.input_data, null, 2)}>
-        <pre
+        <JsonHighlight
+          json={JSON.stringify(run.input_data, null, 2)}
           className="p-3 rounded-lg text-xs font-mono whitespace-pre-wrap break-words"
-          style={{ background: "var(--bg-secondary)", color: "var(--text-primary)", border: "1px solid var(--border)" }}
-        >
-          {JSON.stringify(run.input_data, null, 2)}
-        </pre>
+          style={{ background: "var(--bg-secondary)", border: "1px solid var(--border)" }}
+        />
       </DataSection>
 
       {/* Output */}
@@ -344,14 +344,13 @@ function IOView({ run }: { run: RunSummary }) {
           color="var(--accent)"
           copyText={typeof run.output_data === "string" ? run.output_data : JSON.stringify(run.output_data, null, 2)}
         >
-          <pre
-            className="p-3 rounded-lg text-xs font-mono whitespace-pre-wrap break-words"
-            style={{ background: "var(--bg-secondary)", color: "var(--text-primary)", border: "1px solid var(--border)" }}
-          >
-            {typeof run.output_data === "string"
+          <JsonHighlight
+            json={typeof run.output_data === "string"
               ? run.output_data
               : JSON.stringify(run.output_data, null, 2)}
-          </pre>
+            className="p-3 rounded-lg text-xs font-mono whitespace-pre-wrap break-words"
+            style={{ background: "var(--bg-secondary)", border: "1px solid var(--border)" }}
+          />
         </DataSection>
       )}
 

src/uipath/dev/server/frontend/src/components/shared/JsonHighlight.tsx

@@ -0,0 +1,85 @@
+import { useMemo } from "react";
+
+interface Token {
+  type: "key" | "string" | "number" | "boolean" | "null" | "punctuation";
+  text: string;
+}
+
+const TOKEN_COLORS: Record<Token["type"], string> = {
+  key: "var(--info)",
+  string: "var(--success)",
+  number: "var(--warning)",
+  boolean: "var(--accent)",
+  null: "var(--accent)",
+  punctuation: "var(--text-muted)",
+};
+
+// Tokenize a pre-formatted JSON string into colored spans
+function tokenize(json: string): Token[] {
+  const tokens: Token[] = [];
+  // Match JSON tokens: strings, numbers, booleans, null, punctuation
+  const re = /("(?:[^"\\]|\\.)*")\s*:|("(?:[^"\\]|\\.)*")|(-?\d+(?:\.\d+)?(?:[eE][+-]?\d+)?)\b|(true|false)\b|(null)\b|([{}[\]:,])/g;
+  let lastIndex = 0;
+  let m: RegExpExecArray | null;
+
+  while ((m = re.exec(json)) !== null) {
+    // Whitespace between tokens
+    if (m.index > lastIndex) {
+      tokens.push({ type: "punctuation", text: json.slice(lastIndex, m.index) });
+    }
+
+    if (m[1] !== undefined) {
+      // Key (string followed by colon)
+      tokens.push({ type: "key", text: m[1] });
+      // Find the colon after the key
+      const colonIdx = json.indexOf(":", m.index + m[1].length);
+      if (colonIdx !== -1) {
+        // Add any whitespace between key and colon
+        if (colonIdx > m.index + m[1].length) {
+          tokens.push({ type: "punctuation", text: json.slice(m.index + m[1].length, colonIdx) });
+        }
+        tokens.push({ type: "punctuation", text: ":" });
+        re.lastIndex = colonIdx + 1;
+      }
+    } else if (m[2] !== undefined) {
+      tokens.push({ type: "string", text: m[2] });
+    } else if (m[3] !== undefined) {
+      tokens.push({ type: "number", text: m[3] });
+    } else if (m[4] !== undefined) {
+      tokens.push({ type: "boolean", text: m[4] });
+    } else if (m[5] !== undefined) {
+      tokens.push({ type: "null", text: m[5] });
+    } else if (m[6] !== undefined) {
+      tokens.push({ type: "punctuation", text: m[6] });
+    }
+
+    lastIndex = re.lastIndex;
+  }
+
+  // Trailing whitespace
+  if (lastIndex < json.length) {
+    tokens.push({ type: "punctuation", text: json.slice(lastIndex) });
+  }
+
+  return tokens;
+}
+
+interface Props {
+  json: string;
+  className?: string;
+  style?: React.CSSProperties;
+}
+
+export default function JsonHighlight({ json, className, style }: Props) {
+  const tokens = useMemo(() => tokenize(json), [json]);
+
+  return (
+    <pre className={className} style={style}>
+      {tokens.map((t, i) => (
+        <span key={i} style={{ color: TOKEN_COLORS[t.type] }}>
+          {t.text}
+        </span>
+      ))}
+    </pre>
+  );
+}

src/uipath/dev/server/frontend/src/components/traces/SpanDetails.tsx

@@ -1,5 +1,6 @@
 import { useState, useMemo, useCallback } from "react";
 import type { TraceSpan } from "../../types/run";
+import JsonHighlight from "../shared/JsonHighlight";
 
 const STATUS_CONFIG: Record<string, { color: string; label: string }> = {
   started: { color: "var(--info)", label: "Started" },
@@ -56,12 +57,19 @@ function AttributeValue({ value }: { value: unknown }) {
   const [expanded, setExpanded] = useState(false);
   const raw = stringifyValue(value);
   const jsonFormatted = useMemo(() => tryParseJson(value), [value]);
+  const isJson = jsonFormatted !== null;
   const displayValue = jsonFormatted ?? raw;
   const isLong = displayValue.length > TRUNCATE_LIMIT || displayValue.includes("\n");
   const toggle = useCallback(() => setExpanded((prev) => !prev), []);
 
   if (!isLong) {
-    return (
+    return isJson ? (
+      <JsonHighlight
+        json={displayValue}
+        className="font-mono text-[11px] break-all whitespace-pre-wrap"
+        style={{}}
+      />
+    ) : (
       <span className="font-mono text-[11px] break-all" style={{ color: "var(--text-primary)" }}>
         {displayValue}
       </span>
@@ -71,12 +79,20 @@ function AttributeValue({ value }: { value: unknown }) {
   return (
     <div>
       {expanded ? (
-        <pre
-          className="font-mono text-[11px] whitespace-pre-wrap break-all"
-          style={{ color: "var(--text-primary)" }}
-        >
-          {displayValue}
-        </pre>
+        isJson ? (
+          <JsonHighlight
+            json={displayValue}
+            className="font-mono text-[11px] whitespace-pre-wrap break-all"
+            style={{}}
+          />
+        ) : (
+          <pre
+            className="font-mono text-[11px] whitespace-pre-wrap break-all"
+            style={{ color: "var(--text-primary)" }}
+          >
+            {displayValue}
+          </pre>
+        )
       ) : (
         <span className="font-mono text-[11px] break-all" style={{ color: "var(--text-primary)" }}>
           {displayValue.slice(0, TRUNCATE_LIMIT)}...