diff --git a/app/agent_base.py b/app/agent_base.py index 618aa4e7..6fba7b2a 100644 --- a/app/agent_base.py +++ b/app/agent_base.py @@ -1525,6 +1525,59 @@ async def _handle_chat_message(self, payload: Dict): # Local TUI/CLI message platform = "CraftBot TUI" + # Direct reply bypass - skip routing LLM when target_session_id is provided + target_session_id = payload.get("target_session_id") + if target_session_id: + logger.info(f"[CHAT] Direct reply to session {target_session_id}") + + # Fire the target trigger directly, bypassing routing LLM + fired = await self.triggers.fire( + target_session_id, message=chat_content, platform=platform + ) + + if fired: + logger.info(f"[CHAT] Successfully resumed session {target_session_id}") + + # Reset task status from "waiting" to "running" when user replies + if self.ui_controller: + from app.ui_layer.events import UIEvent, UIEventType + + self.ui_controller.event_bus.emit( + UIEvent( + type=UIEventType.TASK_UPDATE, + data={ + "task_id": target_session_id, + "status": "running", + }, + ) + ) + + # Check if there are still other tasks waiting + triggers = await self.triggers.list_triggers() + has_waiting_tasks = any( + getattr(t, 'waiting_for_reply', False) + for t in triggers + if t.session_id != target_session_id + ) + if not has_waiting_tasks: + self.ui_controller.event_bus.emit( + UIEvent( + type=UIEventType.AGENT_STATE_CHANGED, + data={ + "state": "working", + "status_message": "Agent is working...", + }, + ) + ) + + return # Task will resume with user message in event stream + + # If fire() returns False, no waiting trigger found for this session + # Fall through to normal routing (conversation mode) + logger.warning( + f"[CHAT] Session {target_session_id} not found or expired, falling through to normal routing" + ) + # Check active tasks — route message to matching session if possible # Use active_task_ids from state_manager (not just triggers in queue) to ensure # all running tasks are visible for routing, not just those waiting in queue diff --git a/app/ui_layer/adapters/base.py b/app/ui_layer/adapters/base.py index 620d074e..f8f9aa8f 100644 --- a/app/ui_layer/adapters/base.py +++ b/app/ui_layer/adapters/base.py @@ -273,7 +273,10 @@ def _handle_agent_message(self, event: UIEvent) -> None: agent_name = onboarding_manager.state.agent_name or "Agent" asyncio.create_task( self._display_chat_message( - agent_name, event.data.get("message", ""), "agent" + agent_name, + event.data.get("message", ""), + "agent", + task_session_id=event.task_id, ) ) @@ -438,6 +441,7 @@ async def _display_chat_message( label: str, message: str, style: str, + task_session_id: Optional[str] = None, ) -> None: """ Display a chat message. @@ -446,6 +450,7 @@ async def _display_chat_message( label: Message sender label message: Message content style: Style identifier + task_session_id: Optional task session ID for reply feature """ import time @@ -455,6 +460,7 @@ async def _display_chat_message( content=message, style=style, timestamp=time.time(), + task_session_id=task_session_id, ) ) diff --git a/app/ui_layer/adapters/browser_adapter.py b/app/ui_layer/adapters/browser_adapter.py index 4d1536fa..78b737f9 100644 --- a/app/ui_layer/adapters/browser_adapter.py +++ b/app/ui_layer/adapters/browser_adapter.py @@ -238,6 +238,7 @@ async def append_message(self, message: ChatMessage) -> None: style=message.style, timestamp=message.timestamp, attachments=attachments_data, + task_session_id=message.task_session_id, ) self._storage.insert_message(stored) except Exception: @@ -265,6 +266,10 @@ async def append_message(self, message: ChatMessage) -> None: for att in message.attachments ] + # Include task session ID for reply feature + if message.task_session_id: + message_data["taskSessionId"] = message.task_session_id + await self._adapter._broadcast({ "type": "chat_message", "data": message_data, @@ -676,6 +681,36 @@ def metrics_collector(self) -> MetricsCollector: """Get the metrics collector for dashboard data.""" return self._metrics_collector + async def submit_message( + self, + message: str, + reply_context: Optional[Dict[str, Any]] = None + ) -> None: + """ + Submit a message from the user with optional reply context. + + Overrides base class to handle reply-to-chat/task feature. + Appends reply context to the message before routing to the agent. + + Args: + message: The user's input message + reply_context: Optional dict with {sessionId?: str, originalMessage: str} + """ + agent_context = message + + # Add reply context note (similar to attachment_note pattern) + if reply_context and reply_context.get("originalMessage"): + reply_note = f"\n\n[REPLYING TO PREVIOUS AGENT MESSAGE]:\n{reply_context['originalMessage']}" + agent_context = message + reply_note + + # Pass to controller with target session ID if replying + target_session_id = reply_context.get("sessionId") if reply_context else None + await self._controller.submit_message( + agent_context, + self._adapter_id, + target_session_id=target_session_id + ) + def _handle_task_start(self, event: UIEvent) -> None: """Handle task start event with metrics tracking.""" # Call parent implementation @@ -874,16 +909,17 @@ async def _handle_ws_message(self, data: Dict[str, Any]) -> None: msg_type = data.get("type") if msg_type == "message": - # User sent a message (may include attachments) + # User sent a message (may include attachments and/or reply context) content = data.get("content", "") attachments = data.get("attachments", []) + reply_context = data.get("replyContext") # {sessionId?: str, originalMessage: str} if attachments: # Message with attachments - use custom handler - await self._handle_chat_message_with_attachments(content, attachments) + await self._handle_chat_message_with_attachments(content, attachments, reply_context) elif content: # Regular message without attachments - use normal flow - await self.submit_message(content) + await self.submit_message(content, reply_context) elif msg_type == "chat_attachment_upload": # Upload attachment for chat message @@ -3660,8 +3696,13 @@ async def _handle_file_download(self, file_path: str) -> None: }, }) - async def _handle_chat_message_with_attachments(self, content: str, attachments: List[Dict[str, Any]]) -> None: - """Handle user chat message with attachments.""" + async def _handle_chat_message_with_attachments( + self, + content: str, + attachments: List[Dict[str, Any]], + reply_context: Optional[Dict[str, Any]] = None + ) -> None: + """Handle user chat message with attachments and optional reply context.""" import uuid from app.ui_layer.state.ui_state import AgentStateType from app.ui_layer.events import UIEvent, UIEventType @@ -3726,6 +3767,11 @@ async def _handle_chat_message_with_attachments(self, content: str, attachments: # (This is what the agent sees in the event stream - includes file paths) agent_context = content + attachment_note + # Add reply context note (similar to attachment_note pattern) + if reply_context and reply_context.get("originalMessage"): + reply_note = f"\n\n[REPLYING TO PREVIOUS AGENT MESSAGE]:\n{reply_context['originalMessage']}" + agent_context = agent_context + reply_note + if not agent_context.strip(): return @@ -3751,6 +3797,10 @@ async def _handle_chat_message_with_attachments(self, content: str, attachments: "sender": {"id": self._adapter_id or "user", "type": "user"}, "gui_mode": self._controller._state_store.state.gui_mode, } + # Include target session ID if replying to a specific session + if reply_context and reply_context.get("sessionId"): + payload["target_session_id"] = reply_context["sessionId"] + await self._controller._agent._handle_chat_message(payload) except Exception as e: diff --git a/app/ui_layer/browser/frontend/src/contexts/WebSocketContext.tsx b/app/ui_layer/browser/frontend/src/contexts/WebSocketContext.tsx index b4808353..8ad9a226 100644 --- a/app/ui_layer/browser/frontend/src/contexts/WebSocketContext.tsx +++ b/app/ui_layer/browser/frontend/src/contexts/WebSocketContext.tsx @@ -10,6 +10,20 @@ interface PendingAttachment { content: string // base64 } +// Reply target for reply-to-chat/task feature +interface ReplyTarget { + type: 'chat' | 'task' + sessionId?: string // May be undefined for old messages without session tracking + displayName: string // Truncated preview for UI display + originalContent: string // Full content for agent context +} + +// Reply context sent with message +interface ReplyContext { + sessionId?: string + originalMessage: string +} + interface WebSocketState { connected: boolean messages: ChatMessage[] @@ -29,10 +43,12 @@ interface WebSocketState { onboardingLoading: boolean // Unread message tracking lastSeenMessageId: string | null + // Reply state for reply-to-chat/task feature + replyTarget: ReplyTarget | null } interface WebSocketContextType extends WebSocketState { - sendMessage: (content: string, attachments?: PendingAttachment[]) => void + sendMessage: (content: string, attachments?: PendingAttachment[], replyContext?: ReplyContext) => void sendCommand: (command: string) => void clearMessages: () => void cancelTask: (taskId: string) => void @@ -46,6 +62,9 @@ interface WebSocketContextType extends WebSocketState { goBackOnboardingStep: () => void // Unread message tracking markMessagesAsSeen: () => void + // Reply-to-chat/task methods + setReplyTarget: (target: ReplyTarget) => void + clearReplyTarget: () => void } // Initialize lastSeenMessageId from localStorage @@ -86,6 +105,8 @@ const defaultState: WebSocketState = { onboardingLoading: false, // Unread message tracking lastSeenMessageId: getInitialLastSeenMessageId(), + // Reply state + replyTarget: null, } const WebSocketContext = createContext(undefined) @@ -462,12 +483,13 @@ export function WebSocketProvider({ children }: { children: ReactNode }) { } }, [connect]) - const sendMessage = useCallback((content: string, attachments?: PendingAttachment[]) => { + const sendMessage = useCallback((content: string, attachments?: PendingAttachment[], replyContext?: ReplyContext) => { if (wsRef.current?.readyState === WebSocket.OPEN) { wsRef.current.send(JSON.stringify({ type: 'message', content, - attachments: attachments || [] + attachments: attachments || [], + replyContext: replyContext || null, })) } }, []) @@ -557,6 +579,16 @@ export function WebSocketProvider({ children }: { children: ReactNode }) { }) }, []) + // Set reply target for reply-to-chat/task feature + const setReplyTarget = useCallback((target: ReplyTarget) => { + setState(prev => ({ ...prev, replyTarget: target })) + }, []) + + // Clear reply target + const clearReplyTarget = useCallback(() => { + setState(prev => ({ ...prev, replyTarget: null })) + }, []) + return ( {children} diff --git a/app/ui_layer/browser/frontend/src/pages/Chat/ChatMessage.tsx b/app/ui_layer/browser/frontend/src/pages/Chat/ChatMessage.tsx index 10e954b7..0c5e26c0 100644 --- a/app/ui_layer/browser/frontend/src/pages/Chat/ChatMessage.tsx +++ b/app/ui_layer/browser/frontend/src/pages/Chat/ChatMessage.tsx @@ -1,5 +1,6 @@ -import React, { memo } from 'react' -import { MarkdownContent, AttachmentDisplay } from '../../components/ui' +import React, { memo, useState, useMemo } from 'react' +import { Reply } from 'lucide-react' +import { MarkdownContent, AttachmentDisplay, IconButton } from '../../components/ui' import type { ChatMessage as ChatMessageType } from '../../types' import styles from './ChatPage.module.css' @@ -7,25 +8,92 @@ interface ChatMessageProps { message: ChatMessageType onOpenFile: (path: string) => void onOpenFolder: (path: string) => void + onReply?: ( + sessionId: string | undefined, + displayName: string, + fullContent: string + ) => void +} + +// Parse reply context from message content +const REPLY_MARKER = '[REPLYING TO PREVIOUS AGENT MESSAGE]:' + +function parseReplyContext(content: string): { userMessage: string; replyContext: string | null } { + const markerIndex = content.indexOf(REPLY_MARKER) + if (markerIndex === -1) { + return { userMessage: content, replyContext: null } + } + const userMessage = content.slice(0, markerIndex).trim() + const replyContext = content.slice(markerIndex + REPLY_MARKER.length).trim() + return { userMessage, replyContext } } export const ChatMessageItem = memo(function ChatMessageItem({ message, onOpenFile, - onOpenFolder + onOpenFolder, + onReply, }: ChatMessageProps) { + const [isHovered, setIsHovered] = useState(false) + + // Show reply for ALL agent messages + const canReply = message.style === 'agent' && onReply + + // Parse reply context for user messages + const { userMessage, replyContext } = useMemo(() => { + if (message.style === 'user') { + return parseReplyContext(message.content) + } + return { userMessage: message.content, replyContext: null } + }, [message.content, message.style]) + + const handleReply = (e: React.MouseEvent) => { + e.stopPropagation() + if (canReply) { + // Truncate content for display preview + const displayName = message.content.length > 50 + ? message.content.slice(0, 50) + '...' + : message.content + onReply(message.taskSessionId, displayName, message.content) + } + } + return ( -
-
-
- {message.sender} - - {new Date(message.timestamp * 1000).toLocaleTimeString()} - -
-
- +
setIsHovered(true)} + onMouseLeave={() => setIsHovered(false)} + > + {/* Message bubble container - for positioning reply button outside */} +
+
+
+ {message.sender} + + {new Date(message.timestamp * 1000).toLocaleTimeString()} + +
+ {/* Reply context callout - shown above user message when replying */} + {replyContext && ( +
+ +
+ )} +
+ +
+ {/* Reply button - positioned outside the bubble at top-right */} + {canReply && isHovered && ( + } + variant="ghost" + size="sm" + onClick={handleReply} + tooltip="Reply to this message" + className={styles.replyButtonOutside} + /> + )}
{message.attachments && message.attachments.length > 0 && (
diff --git a/app/ui_layer/browser/frontend/src/pages/Chat/ChatPage.module.css b/app/ui_layer/browser/frontend/src/pages/Chat/ChatPage.module.css index 18550bde..92db79d1 100644 --- a/app/ui_layer/browser/frontend/src/pages/Chat/ChatPage.module.css +++ b/app/ui_layer/browser/frontend/src/pages/Chat/ChatPage.module.css @@ -505,3 +505,107 @@ .dismissError:hover { opacity: 1; } + +/* ───────────────────────────────────────────────────────────────────── + Reply UI Styles + ───────────────────────────────────────────────────────────────────── */ + +/* Agent wrapper needs padding-right for the reply button */ +.agentWrapper { + padding-right: var(--space-8); +} + +/* Message bubble container - wraps bubble + reply button */ +.messageBubbleContainer { + position: relative; +} + +/* Reply button outside the bubble - positioned in the padding area */ +.replyButtonOutside { + position: absolute; + top: var(--space-1); + left: 100%; + margin-left: var(--space-1); + opacity: 0; + transition: opacity var(--transition-fast); +} + +.messageWrapper:hover .replyButtonOutside { + opacity: 1; +} + +/* Reply bar above input - styled like pending attachments */ +.replyBar { + display: flex; + align-items: center; + gap: var(--space-2); + padding: var(--space-2) var(--space-3); + background: var(--bg-tertiary); + border: 1px solid var(--border-primary); + border-radius: var(--radius-sm); + font-size: var(--text-xs); + color: var(--text-primary); +} + +.replyText { + flex: 1; + overflow: hidden; + text-overflow: ellipsis; + white-space: nowrap; +} + +.replyCancel { + display: flex; + align-items: center; + justify-content: center; + background: none; + border: none; + padding: 2px; + cursor: pointer; + color: var(--text-muted); + transition: color var(--transition-fast); + flex-shrink: 0; +} + +.replyCancel:hover { + color: var(--color-error); +} + +/* Task reply button - shown on hover */ +.taskReplyBtn { + opacity: 0; + flex-shrink: 0; + color: var(--text-muted); + transition: opacity var(--transition-fast), color var(--transition-fast); +} + +.taskItem:hover .taskReplyBtn { + opacity: 1; +} + +.taskReplyBtn:hover { + color: var(--color-primary); + background: var(--color-primary-light); +} + +/* Reply context callout - shown above user message when replying */ +.replyContextCallout { + margin-bottom: var(--space-2); + padding: var(--space-2) var(--space-3); + background: rgba(0, 0, 0, 0.05); + border-left: 3px solid rgba(0, 0, 0, 0.15); + border-radius: var(--radius-sm); + font-size: var(--text-xs); + color: inherit; + opacity: 0.85; +} + +.replyContextCallout p { + margin: 0; +} + +.replyContextCallout ul, +.replyContextCallout ol { + margin: var(--space-1) 0; + padding-left: var(--space-4); +} diff --git a/app/ui_layer/browser/frontend/src/pages/Chat/ChatPage.tsx b/app/ui_layer/browser/frontend/src/pages/Chat/ChatPage.tsx index 1a1d9f64..43839257 100644 --- a/app/ui_layer/browser/frontend/src/pages/Chat/ChatPage.tsx +++ b/app/ui_layer/browser/frontend/src/pages/Chat/ChatPage.tsx @@ -1,5 +1,5 @@ import React, { useState, useRef, useEffect, useLayoutEffect, KeyboardEvent, useCallback, ChangeEvent, useMemo } from 'react' -import { Send, Paperclip, X, Loader2, File, AlertCircle } from 'lucide-react' +import { Send, Paperclip, X, Loader2, File, AlertCircle, Reply } from 'lucide-react' import { useVirtualizer } from '@tanstack/react-virtual' import { useLocation } from 'react-router-dom' import { useWebSocket } from '../../contexts/WebSocketContext' @@ -35,7 +35,7 @@ const formatFileSize = (bytes: number): string => { } export function ChatPage() { - const { messages, actions, connected, sendMessage, cancelTask, cancellingTaskId, openFile, openFolder, lastSeenMessageId, markMessagesAsSeen } = useWebSocket() + const { messages, actions, connected, sendMessage, cancelTask, cancellingTaskId, openFile, openFolder, lastSeenMessageId, markMessagesAsSeen, replyTarget, setReplyTarget, clearReplyTarget } = useWebSocket() // Derive agent status from actions and messages const status = useDerivedAgentStatus({ @@ -213,15 +213,52 @@ export function ChatPage() { } }, [isResizing]) + // Handle reply from chat message + const handleChatReply = useCallback(( + sessionId: string | undefined, + displayName: string, + fullContent: string + ) => { + setReplyTarget({ + type: 'chat', + sessionId, + displayName, + originalContent: fullContent, + }) + inputRef.current?.focus() + }, [setReplyTarget]) + + // Handle reply from task panel + const handleTaskReply = useCallback((taskId: string, taskName: string) => { + setReplyTarget({ + type: 'task', + sessionId: taskId, + displayName: taskName, + originalContent: `Task: ${taskName}`, + }) + inputRef.current?.focus() + }, [setReplyTarget]) + const handleSend = () => { // Don't send if there are validation errors if (!attachmentValidation.valid) return if (input.trim() || pendingAttachments.length > 0) { - sendMessage(input.trim(), pendingAttachments.length > 0 ? pendingAttachments : undefined) + // Include reply context if replying to a message/task + const replyContext = replyTarget ? { + sessionId: replyTarget.sessionId, + originalMessage: replyTarget.originalContent, + } : undefined + + sendMessage( + input.trim(), + pendingAttachments.length > 0 ? pendingAttachments : undefined, + replyContext + ) setInput('') setPendingAttachments([]) setAttachmentError(null) + clearReplyTarget() // Clear reply target after sending // Reset textarea height after clearing input if (inputRef.current) { inputRef.current.style.height = 'auto' @@ -370,6 +407,7 @@ export function ChatPage() { message={message} onOpenFile={openFile} onOpenFolder={openFolder} + onReply={handleChatReply} />
) @@ -418,6 +456,23 @@ export function ChatPage() {
)} + {/* Reply bar - shows when replying to a message/task */} + {replyTarget && ( +
+ + + Replying to: {replyTarget.displayName} + + +
+ )} + {/* Pending attachments preview */} {pendingAttachments.length > 0 && (
@@ -489,24 +544,37 @@ export function ChatPage() { {task.name} {(task.status === 'running' || task.status === 'waiting') && ( - { - e.stopPropagation() - cancelTask(task.id) - }} - disabled={cancellingTaskId === task.id} - title="Cancel Task" - icon={ - cancellingTaskId === task.id ? ( - - ) : ( - - ) - } - /> + <> + { + e.stopPropagation() + handleTaskReply(task.id, task.name) + }} + title="Reply to Task" + icon={} + /> + { + e.stopPropagation() + cancelTask(task.id) + }} + disabled={cancellingTaskId === task.id} + title="Cancel Task" + icon={ + cancellingTaskId === task.id ? ( + + ) : ( + + ) + } + /> + )}
{selectedTaskId === task.id && ( diff --git a/app/ui_layer/browser/frontend/src/pages/Tasks/TasksPage.module.css b/app/ui_layer/browser/frontend/src/pages/Tasks/TasksPage.module.css index 5a58d460..b9892244 100644 --- a/app/ui_layer/browser/frontend/src/pages/Tasks/TasksPage.module.css +++ b/app/ui_layer/browser/frontend/src/pages/Tasks/TasksPage.module.css @@ -442,3 +442,24 @@ word-break: break-word; } } + +/* ───────────────────────────────────────────────────────────────────── + Reply Button Styles + ───────────────────────────────────────────────────────────────────── */ + +/* Task reply button - shown on hover */ +.taskReplyBtn { + opacity: 0; + flex-shrink: 0; + color: var(--text-muted); + transition: opacity var(--transition-fast), color var(--transition-fast); +} + +.taskItem:hover .taskReplyBtn { + opacity: 1; +} + +.taskReplyBtn:hover { + color: var(--color-primary); + background: var(--color-primary-light); +} diff --git a/app/ui_layer/browser/frontend/src/pages/Tasks/TasksPage.tsx b/app/ui_layer/browser/frontend/src/pages/Tasks/TasksPage.tsx index 646eac2b..3acea84d 100644 --- a/app/ui_layer/browser/frontend/src/pages/Tasks/TasksPage.tsx +++ b/app/ui_layer/browser/frontend/src/pages/Tasks/TasksPage.tsx @@ -1,5 +1,6 @@ import React, { useState, useRef, useEffect, useCallback } from 'react' -import { ChevronRight, XCircle, ArrowLeft } from 'lucide-react' +import { ChevronRight, XCircle, ArrowLeft, Reply } from 'lucide-react' +import { useNavigate } from 'react-router-dom' import { useWebSocket } from '../../contexts/WebSocketContext' import { StatusIndicator, Badge, Button, IconButton } from '../../components/ui' import type { ActionItem } from '../../types' @@ -271,7 +272,8 @@ const MIN_PANEL_WIDTH = 200 const MAX_PANEL_WIDTH = 600 export function TasksPage() { - const { actions, cancelTask, cancellingTaskId } = useWebSocket() + const { actions, cancelTask, cancellingTaskId, setReplyTarget } = useWebSocket() + const navigate = useNavigate() const [selectedItem, setSelectedItem] = useState(null) const [mobileShowDetail, setMobileShowDetail] = useState(false) @@ -282,6 +284,17 @@ export function TasksPage() { const tasks = actions.filter(a => a.itemType === 'task') + // Handle reply to task - set reply target and navigate to chat + const handleTaskReply = useCallback((task: ActionItem) => { + setReplyTarget({ + type: 'task', + sessionId: task.id, + displayName: task.name, + originalContent: `Task: ${task.name}`, + }) + navigate('/chat') + }, [setReplyTarget, navigate]) + // Get all items (actions + reasoning) for a task const getItemsForTask = (taskId: string) => actions.filter(a => (a.itemType === 'action' || a.itemType === 'reasoning') && a.parentId === taskId) @@ -387,6 +400,19 @@ export function TasksPage() { /> {task.name} + {(task.status === 'running' || task.status === 'waiting') && ( + { + e.stopPropagation() + handleTaskReply(task) + }} + title="Reply to Task" + icon={} + /> + )} None: """Generate message_id if not provided.""" diff --git a/app/ui_layer/controller/ui_controller.py b/app/ui_layer/controller/ui_controller.py index 09e4dc05..e2f7f343 100644 --- a/app/ui_layer/controller/ui_controller.py +++ b/app/ui_layer/controller/ui_controller.py @@ -219,7 +219,12 @@ def unregister_adapter(self) -> None: # Message Handling # ───────────────────────────────────────────────────────────────────── - async def submit_message(self, message: str, adapter_id: str = "") -> None: + async def submit_message( + self, + message: str, + adapter_id: str = "", + target_session_id: Optional[str] = None + ) -> None: """ Handle user input from any interface. @@ -228,6 +233,7 @@ async def submit_message(self, message: str, adapter_id: str = "") -> None: Args: message: The user's input message adapter_id: ID of the adapter that sent the message + target_session_id: Optional session ID for direct reply (bypasses routing) """ if not message.strip(): return @@ -268,6 +274,10 @@ async def submit_message(self, message: str, adapter_id: str = "") -> None: "sender": {"id": adapter_id or "user", "type": "user"}, "gui_mode": self._state_store.state.gui_mode, } + # Include target session ID for direct reply (bypasses routing LLM) + if target_session_id: + payload["target_session_id"] = target_session_id + await self._agent._handle_chat_message(payload) # ───────────────────────────────────────────────────────────────────── diff --git a/app/usage/chat_storage.py b/app/usage/chat_storage.py index bb1c179e..6a4ff9d9 100644 --- a/app/usage/chat_storage.py +++ b/app/usage/chat_storage.py @@ -33,6 +33,7 @@ class StoredChatMessage: style: str timestamp: float attachments: Optional[List[Dict[str, Any]]] = None + task_session_id: Optional[str] = None def to_dict(self) -> Dict[str, Any]: """Convert to dictionary for JSON serialization.""" @@ -45,6 +46,8 @@ def to_dict(self) -> Dict[str, Any]: } if self.attachments: result["attachments"] = self.attachments + if self.task_session_id: + result["taskSessionId"] = self.task_session_id return result @@ -101,6 +104,16 @@ def _init_db(self) -> None: ON chat_messages(message_id) """) + # Migration: Add task_session_id column if it doesn't exist + cursor.execute("PRAGMA table_info(chat_messages)") + columns = [col[1] for col in cursor.fetchall()] + if "task_session_id" not in columns: + cursor.execute(""" + ALTER TABLE chat_messages + ADD COLUMN task_session_id TEXT + """) + logger.info("[ChatStorage] Migrated: added task_session_id column") + conn.commit() def insert_message(self, message: StoredChatMessage) -> int: @@ -117,8 +130,8 @@ def insert_message(self, message: StoredChatMessage) -> int: cursor = conn.cursor() cursor.execute(""" INSERT OR REPLACE INTO chat_messages - (message_id, sender, content, style, timestamp, attachments) - VALUES (?, ?, ?, ?, ?, ?) + (message_id, sender, content, style, timestamp, attachments, task_session_id) + VALUES (?, ?, ?, ?, ?, ?, ?) """, ( message.message_id, message.sender, @@ -126,6 +139,7 @@ def insert_message(self, message: StoredChatMessage) -> int: message.style, message.timestamp, json.dumps(message.attachments) if message.attachments else None, + message.task_session_id, )) conn.commit() return cursor.lastrowid @@ -148,7 +162,7 @@ def get_messages( with sqlite3.connect(self._db_path) as conn: cursor = conn.cursor() cursor.execute(""" - SELECT message_id, sender, content, style, timestamp, attachments + SELECT message_id, sender, content, style, timestamp, attachments, task_session_id FROM chat_messages ORDER BY timestamp ASC LIMIT ? OFFSET ? @@ -163,6 +177,7 @@ def get_messages( style=row[3], timestamp=row[4], attachments=json.loads(row[5]) if row[5] else None, + task_session_id=row[6], ) for row in rows ] @@ -181,7 +196,7 @@ def get_recent_messages(self, limit: int = 100) -> List[StoredChatMessage]: cursor = conn.cursor() # Get last N messages ordered by timestamp DESC, then reverse cursor.execute(""" - SELECT message_id, sender, content, style, timestamp, attachments + SELECT message_id, sender, content, style, timestamp, attachments, task_session_id FROM chat_messages ORDER BY timestamp DESC LIMIT ? @@ -196,6 +211,7 @@ def get_recent_messages(self, limit: int = 100) -> List[StoredChatMessage]: style=row[3], timestamp=row[4], attachments=json.loads(row[5]) if row[5] else None, + task_session_id=row[6], ) for row in rows ]