improved tool calling, push to hub, etc.

frontend/src/components/Chat/ApprovalFlow.tsx

@@ -118,6 +118,10 @@ export default function ApprovalFlow({ message }: ApprovalFlowProps) {
 
   const currentTool = batch.tools[currentIndex];
 
+  // Check if script contains push_to_hub or upload_file
+  const args = currentTool.arguments as any;
+  const containsPushToHub = currentTool.tool === 'hf_jobs' && args.script && (args.script.includes('push_to_hub') || args.script.includes('upload_file'));
+
   const getToolDescription = (toolName: string, args: any) => {
     if (toolName === 'hf_jobs') {
       return (
@@ -217,6 +221,12 @@ export default function ApprovalFlow({ message }: ApprovalFlowProps) {
         <OpenInNewIcon sx={{ fontSize: 16, color: 'var(--muted-text)', opacity: 0.7 }} />
       </Box>
 
+      {containsPushToHub && (
+        <Typography variant="caption" sx={{ color: 'var(--accent-green)', fontSize: '0.75rem', opacity: 0.8, px: 0.5 }}>
+          We've detected the result will be pushed to hub.
+        </Typography>
+      )}
+
       {showLogsButton && (
         <Button
             variant="outlined"

frontend/src/components/Chat/MessageBubble.tsx

@@ -63,34 +63,6 @@ export default function MessageBubble({ message }: MessageBubbleProps) {
           </Box>
         )}
 
-        {/* Persisted Trace Logs */}
-        {message.trace && message.trace.length > 0 && (
-          <Box
-            sx={{
-              bgcolor: 'rgba(0,0,0,0.3)',
-              borderRadius: 1,
-              p: 1.5,
-              border: 1,
-              borderColor: 'rgba(255,255,255,0.05)',
-              width: '100%',
-              mb: 2,
-            }}
-          >
-            <Box sx={{ display: 'flex', flexDirection: 'column', gap: 0.5 }}>
-              {message.trace.map((log) => (
-                <Typography
-                  key={log.id}
-                  variant="caption"
-                  component="div"
-                  sx={{ color: 'var(--muted-text)', fontFamily: 'ui-monospace, SFMono-Regular, Menlo, Monaco, monospace', fontSize: '0.75rem' }}
-                >
-                  &gt; {log.text}
-                </Typography>
-              ))}
-            </Box>
-          </Box>
-        )}
-
         <Box
           sx={{
             '& p': { m: 0, color: isUser ? 'var(--text)' : 'var(--text)' }, // User might want different text color? Defaults to --text
@@ -147,7 +119,56 @@ export default function MessageBubble({ message }: MessageBubbleProps) {
         >
           <ReactMarkdown remarkPlugins={[remarkGfm]}>{message.content}</ReactMarkdown>
         </Box>
-        
+
+        {/* Persisted Trace Logs - Now at the bottom */}
+        {message.trace && message.trace.length > 0 && (
+          <Box
+            sx={{
+              bgcolor: 'rgba(0,0,0,0.3)',
+              borderRadius: 1,
+              p: 1.5,
+              border: 1,
+              borderColor: 'rgba(255,255,255,0.05)',
+              width: '100%',
+              mt: 2,
+            }}
+          >
+            <Box sx={{ display: 'flex', flexDirection: 'column', gap: 0.5 }}>
+              {message.trace.map((log) => {
+                // Extract tool name from text "Agent is executing {toolName}..."
+                const match = log.text.match(/Agent is executing (.+)\.\.\./);
+                const toolName = match ? match[1] : log.tool;
+
+                return (
+                  <Typography
+                    key={log.id}
+                    variant="caption"
+                    component="div"
+                    sx={{
+                      color: 'var(--muted-text)',
+                      fontFamily: 'ui-monospace, SFMono-Regular, Menlo, Monaco, monospace',
+                      fontSize: '0.75rem',
+                      display: 'flex',
+                      alignItems: 'center',
+                      gap: 0.5,
+                    }}
+                  >
+                    <span style={{ color: log.completed ? '#FDB022' : 'inherit' }}>*</span>
+                    <span>Agent is executing </span>
+                    <span style={{
+                      fontWeight: 600,
+                      color: 'rgba(255, 255, 255, 0.9)',
+                    }}>
+                      {toolName}
+                    </span>
+                    <span>...</span>
+                  </Typography>
+                );
+              })}
+            </Box>
+          </Box>
+        )}
+
         <Typography className="meta" variant="caption" sx={{ display: 'block', textAlign: 'right', mt: 1, fontSize: '11px', opacity: 0.5 }}>
             {new Date(message.timestamp).toLocaleTimeString()}
         </Typography>

frontend/src/components/Chat/MessageList.tsx

@@ -1,6 +1,5 @@
 import { useEffect, useRef } from 'react';
 import { Box, Typography } from '@mui/material';
-import { useAgentStore } from '@/store/agentStore';
 import { useSessionStore } from '@/store/sessionStore';
 import MessageBubble from './MessageBubble';
 import type { Message } from '@/types/agent';
@@ -40,21 +39,12 @@ const TechnicalIndicator = () => (
 
 export default function MessageList({ messages, isProcessing }: MessageListProps) {
   const bottomRef = useRef<HTMLDivElement>(null);
-  const traceBoxRef = useRef<HTMLDivElement>(null);
-  const { traceLogs } = useAgentStore();
   const { activeSessionId } = useSessionStore();
 
   // Auto-scroll to bottom when new messages arrive
   useEffect(() => {
     bottomRef.current?.scrollIntoView({ behavior: 'smooth' });
-  }, [messages, isProcessing, traceLogs]);
-
-  // Auto-scroll trace box
-  useEffect(() => {
-    if (traceBoxRef.current) {
-      traceBoxRef.current.scrollTop = traceBoxRef.current.scrollHeight;
-    }
-  }, [traceLogs]);
+  }, [messages, isProcessing]);
 
   return (
     <Box
@@ -67,7 +57,7 @@ export default function MessageList({ messages, isProcessing }: MessageListProps
       }}
     >
       <Box sx={{ maxWidth: 'md', mx: 'auto', width: '100%', display: 'flex', flexDirection: 'column', gap: 2 }}>
-        {messages.length === 0 && traceLogs.length === 0 && !isProcessing ? (
+        {messages.length === 0 && !isProcessing ? (
           <Box
             sx={{
               flex: 1,
@@ -95,37 +85,6 @@ export default function MessageList({ messages, isProcessing }: MessageListProps
               </Typography>
               <TechnicalIndicator />
             </Box>
-            
-            {traceLogs.length > 0 && (
-              <Box
-                sx={{
-                  bgcolor: 'background.default',
-                  borderRadius: 1,
-                  p: 2,
-                  border: 1,
-                  borderColor: 'divider',
-                  width: '100%',
-                  fontFamily: 'monospace',
-                  maxHeight: 120,
-                  overflowY: 'auto',
-                }}
-                ref={traceBoxRef}
-              >
-                <Box sx={{ display: 'flex', flexDirection: 'column', gap: 0.5 }}>
-                  {traceLogs.map((log) => (
-                    <Box key={log.id}>
-                      <Typography
-                        variant="caption"
-                        component="div"
-                        sx={{ color: 'common.white', fontFamily: 'monospace' }}
-                      >
-                        &gt; {log.text}
-                      </Typography>
-                    </Box>
-                  ))}
-                </Box>
-              </Box>
-            )}
           </Box>
         )}
 

frontend/src/components/SessionSidebar/SessionSidebar.tsx

@@ -47,18 +47,21 @@ const RunningIndicator = () => (
 export default function SessionSidebar({ onClose }: SessionSidebarProps) {
   const { sessions, activeSessionId, createSession, deleteSession, switchSession } =
     useSessionStore();
-  const { clearMessages, isConnected, isProcessing } = useAgentStore();
+  const { clearMessages, isConnected, isProcessing, setPlan, setPanelContent } = useAgentStore();
 
   const handleNewSession = useCallback(async () => {
     try {
       const response = await fetch('/api/session', { method: 'POST' });
       const data = await response.json();
       createSession(data.session_id);
+      // Clear plan and code panel for new session
+      setPlan([]);
+      setPanelContent(null);
       onClose?.();
     } catch (e) {
       console.error('Failed to create session:', e);
     }
-  }, [createSession, onClose]);
+  }, [createSession, setPlan, setPanelContent, onClose]);
 
   const handleDeleteSession = useCallback(
     async (sessionId: string, e: React.MouseEvent) => {
@@ -77,9 +80,12 @@ export default function SessionSidebar({ onClose }: SessionSidebarProps) {
   const handleSelectSession = useCallback(
     (sessionId: string) => {
       switchSession(sessionId);
+      // Clear plan and code panel when switching sessions
+      setPlan([]);
+      setPanelContent(null);
       onClose?.();
     },
-    [switchSession, onClose]
+    [switchSession, setPlan, setPanelContent, onClose]
   );
 
   const handleUndo = useCallback(async () => {

frontend/src/hooks/useAgentWebSocket.ts

@@ -25,15 +25,19 @@ export function useAgentWebSocket({
 
   const {
     addMessage,
+    updateMessage,
     setProcessing,
     setConnected,
     setPendingApprovals,
     setError,
     addTraceLog,
+    updateTraceLog,
     clearTraceLogs,
     setPanelContent,
     setPlan,
     traceLogs,
+    setCurrentTurnMessageId,
+    updateCurrentTurnTrace,
   } = useAgentStore();
 
   const { setRightPanelOpen, setLeftSidebarOpen } = useLayoutStore();
@@ -55,34 +59,61 @@ export function useAgentWebSocket({
         case 'processing':
           setProcessing(true);
           clearTraceLogs();
+          setCurrentTurnMessageId(null); // Start a new turn
           break;
 
         case 'assistant_message': {
           const content = (event.data?.content as string) || '';
           const currentTrace = useAgentStore.getState().traceLogs;
-          const message: Message = {
-            id: `msg_${Date.now()}`,
-            role: 'assistant',
-            content,
-            timestamp: new Date().toISOString(),
-            trace: currentTrace.length > 0 ? [...currentTrace] : undefined,
-          };
-          addMessage(sessionId, message);
+          const currentTurnMsgId = useAgentStore.getState().currentTurnMessageId;
+
+          if (currentTurnMsgId) {
+            // Update existing message - append content and update trace
+            const messages = useAgentStore.getState().getMessages(sessionId);
+            const existingMsg = messages.find(m => m.id === currentTurnMsgId);
+
+            if (existingMsg) {
+              const newContent = existingMsg.content ? existingMsg.content + '\n\n' + content : content;
+              updateMessage(sessionId, currentTurnMsgId, {
+                content: newContent,
+                trace: currentTrace.length > 0 ? [...currentTrace] : undefined,
+              });
+            }
+          } else {
+            // Create new message
+            const messageId = `msg_${Date.now()}`;
+            const message: Message = {
+              id: messageId,
+              role: 'assistant',
+              content,
+              timestamp: new Date().toISOString(),
+              trace: currentTrace.length > 0 ? [...currentTrace] : undefined,
+            };
+            addMessage(sessionId, message);
+            setCurrentTurnMessageId(messageId);
+          }
           break;
         }
 
         case 'tool_call': {
           const toolName = (event.data?.tool as string) || 'unknown';
           const args = (event.data?.arguments as Record<string, any>) || {};
-          const log: TraceLog = {
-            id: `tool_${Date.now()}`,
-            type: 'call',
-            text: `Calling ${toolName} with ${JSON.stringify(args)}`,
-            tool: toolName,
-            timestamp: new Date().toISOString(),
-          };
-          addTraceLog(log);
-          
+
+          // Don't display plan_tool in trace logs (it shows up elsewhere in the UI)
+          if (toolName !== 'plan_tool') {
+            const log: TraceLog = {
+              id: `tool_${Date.now()}`,
+              type: 'call',
+              text: `Agent is executing ${toolName}...`,
+              tool: toolName,
+              timestamp: new Date().toISOString(),
+              completed: false,
+            };
+            addTraceLog(log);
+            // Update the current turn message's trace in real-time
+            updateCurrentTurnTrace(sessionId);
+          }
+
           // Auto-expand Right Panel for specific tools
           if (toolName === 'hf_jobs' && (args.operation === 'run' || args.operation === 'scheduled run') && args.script) {
             setPanelContent({
@@ -111,18 +142,21 @@ export function useAgentWebSocket({
           const toolName = (event.data?.tool as string) || 'unknown';
           const output = (event.data?.output as string) || '';
           const success = event.data?.success as boolean;
-          
+
+          // Mark the corresponding trace log as completed
+          updateTraceLog(toolName, { completed: true });
+          // Update the current turn message's trace in real-time
+          updateCurrentTurnTrace(sessionId);
+
+          // Special handling for hf_jobs - update the approval message with output
           if (toolName === 'hf_jobs') {
-            // Find the last message with approval (likely the one that triggered this job)
             const messages = useAgentStore.getState().getMessages(sessionId);
-            // Reverse to find the most recent one
             const lastApprovalMsg = [...messages].reverse().find(m => m.approval);
-            
+
             if (lastApprovalMsg) {
-                // Append output if there's already some (for multiple jobs in batch)
                 const currentOutput = lastApprovalMsg.toolOutput || '';
                 const newOutput = currentOutput ? currentOutput + '\n\n' + output : output;
-                
+
                 useAgentStore.getState().updateMessage(sessionId, lastApprovalMsg.id, {
                     toolOutput: newOutput
                 });
@@ -130,19 +164,9 @@ export function useAgentWebSocket({
             } else {
                 console.warn('Received hf_jobs output but no approval message found to update.');
             }
-            // CRITICAL: Always break for hf_jobs to prevent a separate "Tool" bubble from appearing
-            break;
           }
-          
-          const message: Message = {
-            id: `msg_tool_${Date.now()}`,
-            role: 'tool',
-            content: output,
-            timestamp: new Date().toISOString(),
-            toolName: toolName,
-          };
-          addMessage(sessionId, message);
-          
+
+          // Don't create message bubbles for tool outputs - they only show in trace logs
           console.log('Tool output:', toolName, success);
           break;
         }
@@ -218,6 +242,7 @@ export function useAgentWebSocket({
 
         case 'turn_complete':
           setProcessing(false);
+          setCurrentTurnMessageId(null); // Clear the current turn
           break;
 
         case 'compacted': {

frontend/src/store/agentStore.ts

@@ -18,6 +18,7 @@ interface AgentStore {
   traceLogs: TraceLog[];
   panelContent: { title: string; content: string; language?: string; parameters?: any } | null;
   plan: PlanItem[];
+  currentTurnMessageId: string | null; // Track the current turn's assistant message
 
   // Actions
   addMessage: (sessionId: string, message: Message) => void;
@@ -30,9 +31,12 @@ interface AgentStore {
   setError: (error: string | null) => void;
   getMessages: (sessionId: string) => Message[];
   addTraceLog: (log: TraceLog) => void;
+  updateTraceLog: (toolName: string, updates: Partial<TraceLog>) => void;
   clearTraceLogs: () => void;
   setPanelContent: (content: { title: string; content: string; language?: string; parameters?: any } | null) => void;
   setPlan: (plan: PlanItem[]) => void;
+  setCurrentTurnMessageId: (id: string | null) => void;
+  updateCurrentTurnTrace: (sessionId: string) => void;
 }
 
 export const useAgentStore = create<AgentStore>((set, get) => ({
@@ -45,6 +49,7 @@ export const useAgentStore = create<AgentStore>((set, get) => ({
   traceLogs: [],
   panelContent: null,
   plan: [],
+  currentTurnMessageId: null,
 
   addMessage: (sessionId: string, message: Message) => {
     set((state) => {
@@ -112,6 +117,20 @@ export const useAgentStore = create<AgentStore>((set, get) => ({
     }));
   },
 
+  updateTraceLog: (toolName: string, updates: Partial<TraceLog>) => {
+    set((state) => {
+      // Find the last trace log with this tool name and update it
+      const traceLogs = [...state.traceLogs];
+      for (let i = traceLogs.length - 1; i >= 0; i--) {
+        if (traceLogs[i].tool === toolName && traceLogs[i].type === 'call') {
+          traceLogs[i] = { ...traceLogs[i], ...updates };
+          break;
+        }
+      }
+      return { traceLogs };
+    });
+  },
+
   clearTraceLogs: () => {
     set({ traceLogs: [] });
   },
@@ -123,4 +142,26 @@ export const useAgentStore = create<AgentStore>((set, get) => ({
   setPlan: (plan: PlanItem[]) => {
     set({ plan });
   },
+
+  setCurrentTurnMessageId: (id: string | null) => {
+    set({ currentTurnMessageId: id });
+  },
+
+  updateCurrentTurnTrace: (sessionId: string) => {
+    const state = get();
+    if (state.currentTurnMessageId) {
+      const currentMessages = state.messagesBySession[sessionId] || [];
+      const updatedMessages = currentMessages.map((msg) =>
+        msg.id === state.currentTurnMessageId
+          ? { ...msg, trace: state.traceLogs.length > 0 ? [...state.traceLogs] : undefined }
+          : msg
+      );
+      set({
+        messagesBySession: {
+          ...state.messagesBySession,
+          [sessionId]: updatedMessages,
+        },
+      });
+    }
+  },
 }));

frontend/src/types/agent.ts

@@ -54,6 +54,7 @@ export interface TraceLog {
   text: string;
   tool: string;
   timestamp: string;
+  completed?: boolean;
 }
 
 export interface User {