From 703f2bca7c26f9f464264d07d899f747de56918d Mon Sep 17 00:00:00 2001 From: sean-eskerium Date: Wed, 20 Aug 2025 23:31:47 -0400 Subject: [PATCH] Fixed the socket optimistic updates. And the MCP for update task. --- .../src/components/project-tasks/TasksTab.tsx | 96 ++++++++++++++----- archon-ui-main/src/hooks/useTaskSocket.ts | 65 ++++++++++++- .../src/services/socketIOService.ts | 33 ++----- .../src/services/taskSocketService.ts | 21 ++-- .../mcp_server/features/tasks/task_tools.py | 63 ++++++++---- 5 files changed, 200 insertions(+), 78 deletions(-) diff --git a/archon-ui-main/src/components/project-tasks/TasksTab.tsx b/archon-ui-main/src/components/project-tasks/TasksTab.tsx index 60f862ad..288ff265 100644 --- a/archon-ui-main/src/components/project-tasks/TasksTab.tsx +++ b/archon-ui-main/src/components/project-tasks/TasksTab.tsx @@ -511,33 +511,50 @@ export const TasksTab = ({ debouncedPersistBatchReorder(tasksToUpdate); }, [tasks, updateTasks, debouncedPersistBatchReorder]); - // Task move function (for board view) + // Task move function (for board view) with optimistic UI update const moveTask = async (taskId: string, newStatus: Task['status']) => { console.log(`[TasksTab] Attempting to move task ${taskId} to new status: ${newStatus}`); + + const movingTask = tasks.find(task => task.id === taskId); + if (!movingTask) { + console.warn(`[TasksTab] Task ${taskId} not found for move operation.`); + return; + } + + const oldStatus = movingTask.status; + const newOrder = getNextOrderForStatus(newStatus); + const updatedTask = { ...movingTask, status: newStatus, task_order: newOrder }; + + console.log(`[TasksTab] Moving task ${movingTask.title} from ${oldStatus} to ${newStatus} with order ${newOrder}`); + + // OPTIMISTIC UPDATE: Update UI immediately + setTasks(prev => { + const updated = prev.map(task => task.id === taskId ? updatedTask : task); + setTimeout(() => onTasksChange(updated), 0); + return updated; + }); + console.log(`[TasksTab] Optimistically updated UI for task ${taskId}`); + try { - const movingTask = tasks.find(task => task.id === taskId); - if (!movingTask) { - console.warn(`[TasksTab] Task ${taskId} not found for move operation.`); - return; - } - - const oldStatus = movingTask.status; - const newOrder = getNextOrderForStatus(newStatus); - - console.log(`[TasksTab] Moving task ${movingTask.title} from ${oldStatus} to ${newStatus} with order ${newOrder}`); - - // Update the task with new status and order + // Then update the backend await projectService.updateTask(taskId, { status: mapUIStatusToDBStatus(newStatus), task_order: newOrder }); console.log(`[TasksTab] Successfully updated task ${taskId} status in backend.`); - // Don't update local state immediately - let socket handle it - console.log(`[TasksTab] Waiting for socket update for task ${taskId}.`); + // Socket will confirm the update, but UI is already updated } catch (error) { - console.error(`[TasksTab] Failed to move task ${taskId}:`, error); + console.error(`[TasksTab] Failed to move task ${taskId}, rolling back:`, error); + + // ROLLBACK on error - restore original task + setTasks(prev => { + const updated = prev.map(task => task.id === taskId ? movingTask : task); + setTimeout(() => onTasksChange(updated), 0); + return updated; + }); + alert(`Failed to move task: ${error instanceof Error ? error.message : 'Unknown error'}`); } }; @@ -557,15 +574,50 @@ export const TasksTab = ({ if (!taskToDelete) return; try { - // Delete (actually archives) the task - backend will emit socket event - await projectService.deleteTask(taskToDelete.id); - console.log(`[TasksTab] Task ${taskToDelete.id} archival sent to backend`); + // Add to recently deleted cache to prevent race conditions + setRecentlyDeletedIds(prev => new Set(prev).add(taskToDelete.id)); - // Don't update local state - let socket handle it + // OPTIMISTIC UPDATE: Remove task from UI immediately + setTasks(prev => { + const updated = prev.filter(t => t.id !== taskToDelete.id); + setTimeout(() => onTasksChange(updated), 0); + return updated; + }); + console.log(`[TasksTab] Optimistically removed task ${taskToDelete.id} from UI`); + + // Then delete from backend + await projectService.deleteTask(taskToDelete.id); + console.log(`[TasksTab] Task ${taskToDelete.id} deletion confirmed by backend`); + + // Clear from recently deleted cache after a delay (to catch any lingering socket events) + setTimeout(() => { + setRecentlyDeletedIds(prev => { + const newSet = new Set(prev); + newSet.delete(taskToDelete.id); + return newSet; + }); + }, 3000); // 3 second window to ignore stale socket events } catch (error) { - console.error('Failed to archive task:', error); - // Note: The toast notification for deletion is now handled by TaskBoardView and TaskTableView + console.error('Failed to delete task:', error); + + // Remove from recently deleted cache on error + setRecentlyDeletedIds(prev => { + const newSet = new Set(prev); + newSet.delete(taskToDelete.id); + return newSet; + }); + + // ROLLBACK on error - restore the task + setTasks(prev => { + const updated = [...prev, taskToDelete].sort((a, b) => a.task_order - b.task_order); + setTimeout(() => onTasksChange(updated), 0); + return updated; + }); + console.log(`[TasksTab] Rolled back task deletion for ${taskToDelete.id}`); + + // Re-throw to let the calling component handle the error display + throw error; } finally { setTaskToDelete(null); setShowDeleteConfirm(false); diff --git a/archon-ui-main/src/hooks/useTaskSocket.ts b/archon-ui-main/src/hooks/useTaskSocket.ts index 05b3aecc..376bb501 100644 --- a/archon-ui-main/src/hooks/useTaskSocket.ts +++ b/archon-ui-main/src/hooks/useTaskSocket.ts @@ -6,7 +6,7 @@ * approach that avoids conflicts and connection issues. */ -import { useEffect, useRef, useCallback } from 'react'; +import { useEffect, useRef, useCallback, useState } from 'react'; import { taskSocketService, TaskSocketEvents } from '../services/taskSocketService'; import { WebSocketState } from '../services/socketIOService'; @@ -36,6 +36,10 @@ export function useTaskSocket(options: UseTaskSocketOptions) { const componentIdRef = useRef(`task-socket-${Math.random().toString(36).substring(7)}`); const currentProjectIdRef = useRef(null); const isInitializedRef = useRef(false); + + // Add reactive state for connection status + const [isConnected, setIsConnected] = useState(false); + const [connectionState, setConnectionState] = useState(WebSocketState.DISCONNECTED); // Memoized handlers to prevent unnecessary re-registrations const memoizedHandlers = useCallback((): Partial => { @@ -58,6 +62,44 @@ export function useTaskSocket(options: UseTaskSocketOptions) { onConnectionStateChange ]); + // Subscribe to connection state changes + useEffect(() => { + const checkConnection = () => { + const connected = taskSocketService.isConnected(); + const state = taskSocketService.getConnectionState(); + setIsConnected(connected); + setConnectionState(state); + }; + + // Check initial state + checkConnection(); + + // Poll for connection state changes (since the service doesn't expose event emitters) + const interval = setInterval(checkConnection, 500); + + // Also trigger when connection state handler is called + const wrappedOnConnectionStateChange = onConnectionStateChange ? (state: WebSocketState) => { + setConnectionState(state); + setIsConnected(state === WebSocketState.CONNECTED); + onConnectionStateChange(state); + } : (state: WebSocketState) => { + setConnectionState(state); + setIsConnected(state === WebSocketState.CONNECTED); + }; + + // Update the handler + if (componentIdRef.current && taskSocketService) { + taskSocketService.registerHandlers(componentIdRef.current, { + ...memoizedHandlers(), + onConnectionStateChange: wrappedOnConnectionStateChange + }); + } + + return () => { + clearInterval(interval); + }; + }, [onConnectionStateChange, memoizedHandlers]); + // Initialize connection once and register handlers useEffect(() => { if (!projectId || isInitializedRef.current) return; @@ -65,6 +107,7 @@ export function useTaskSocket(options: UseTaskSocketOptions) { const initializeConnection = async () => { try { console.log(`[USE_TASK_SOCKET] Initializing connection for project: ${projectId}`); + setConnectionState(WebSocketState.CONNECTING); // Register handlers first taskSocketService.registerHandlers(componentIdRef.current, memoizedHandlers()); @@ -76,8 +119,14 @@ export function useTaskSocket(options: UseTaskSocketOptions) { isInitializedRef.current = true; console.log(`[USE_TASK_SOCKET] Successfully initialized for project: ${projectId}`); + // Update connection state after successful connection + setIsConnected(taskSocketService.isConnected()); + setConnectionState(taskSocketService.getConnectionState()); + } catch (error) { console.error(`[USE_TASK_SOCKET] Failed to initialize for project ${projectId}:`, error); + setConnectionState(WebSocketState.DISCONNECTED); + setIsConnected(false); } }; @@ -103,6 +152,8 @@ export function useTaskSocket(options: UseTaskSocketOptions) { const switchProject = async () => { try { + setConnectionState(WebSocketState.CONNECTING); + // Update handlers for new project taskSocketService.registerHandlers(componentIdRef.current, memoizedHandlers()); @@ -112,8 +163,14 @@ export function useTaskSocket(options: UseTaskSocketOptions) { currentProjectIdRef.current = projectId; console.log(`[USE_TASK_SOCKET] Successfully switched to project: ${projectId}`); + // Update connection state + setIsConnected(taskSocketService.isConnected()); + setConnectionState(taskSocketService.getConnectionState()); + } catch (error) { console.error(`[USE_TASK_SOCKET] Failed to switch to project ${projectId}:`, error); + setConnectionState(WebSocketState.DISCONNECTED); + setIsConnected(false); } }; @@ -132,10 +189,10 @@ export function useTaskSocket(options: UseTaskSocketOptions) { }; }, []); - // Return utility functions + // Return reactive state and utility functions return { - isConnected: taskSocketService.isConnected(), - connectionState: taskSocketService.getConnectionState(), + isConnected, // Now reactive! + connectionState, // Now reactive! reconnect: taskSocketService.reconnect.bind(taskSocketService), getCurrentProjectId: taskSocketService.getCurrentProjectId.bind(taskSocketService) }; diff --git a/archon-ui-main/src/services/socketIOService.ts b/archon-ui-main/src/services/socketIOService.ts index 35f2cccc..47db04e6 100644 --- a/archon-ui-main/src/services/socketIOService.ts +++ b/archon-ui-main/src/services/socketIOService.ts @@ -678,28 +678,15 @@ export function createWebSocketService(config?: WebSocketConfig): WebSocketServi return new WebSocketService(config); } -// Create SEPARATE WebSocket instances for different features -// This prevents a failure in one feature (like a long crawl) from breaking the entire site -export const knowledgeSocketIO = new WebSocketService({ - maxReconnectAttempts: 10, // More attempts for crawls that can take a long time - reconnectInterval: 2000, - heartbeatInterval: 30000, - enableAutoReconnect: true -}); +// Create a SINGLE shared WebSocket instance to prevent multiple connections +// This fixes the socket disconnection issue when switching tabs +const sharedSocketInstance = new WebSocketService(); -export const taskUpdateSocketIO = new WebSocketService({ - maxReconnectAttempts: 5, - reconnectInterval: 1000, - heartbeatInterval: 30000, - enableAutoReconnect: true -}); +// Export the SAME instance with different names for backward compatibility +// This ensures only ONE Socket.IO connection is created and shared across all features +export const knowledgeSocketIO = sharedSocketInstance; +export const taskUpdateSocketIO = sharedSocketInstance; +export const projectListSocketIO = sharedSocketInstance; -export const projectListSocketIO = new WebSocketService({ - maxReconnectAttempts: 5, - reconnectInterval: 1000, - heartbeatInterval: 30000, - enableAutoReconnect: true -}); - -// Export knowledgeSocketIO as default for backward compatibility -export default knowledgeSocketIO; \ No newline at end of file +// Export as default for new code +export default sharedSocketInstance; \ No newline at end of file diff --git a/archon-ui-main/src/services/taskSocketService.ts b/archon-ui-main/src/services/taskSocketService.ts index 51c0a9df..caca6586 100644 --- a/archon-ui-main/src/services/taskSocketService.ts +++ b/archon-ui-main/src/services/taskSocketService.ts @@ -13,7 +13,8 @@ * - Proper session identification */ -import { WebSocketService, WebSocketState } from './socketIOService'; +import { WebSocketState } from './socketIOService'; +import sharedSocketInstance from './socketIOService'; export interface Task { id: string; @@ -38,7 +39,7 @@ export interface TaskSocketEvents { class TaskSocketService { private static instance: TaskSocketService | null = null; - private socketService: WebSocketService; + private socketService: typeof sharedSocketInstance; private currentProjectId: string | null = null; private eventHandlers: Map = new Map(); private connectionPromise: Promise | null = null; @@ -47,13 +48,11 @@ class TaskSocketService { private connectionCooldown = 1000; // 1 second cooldown between connection attempts private constructor() { - this.socketService = new WebSocketService({ - maxReconnectAttempts: 5, - reconnectInterval: 1000, - heartbeatInterval: 30000, - enableAutoReconnect: true, - enableHeartbeat: true - }); + // Use the shared socket instance instead of creating a new one + this.socketService = sharedSocketInstance; + + // Enable operation tracking for echo suppression + this.socketService.enableOperationTracking(); // Set up global event handlers this.setupGlobalHandlers(); @@ -191,7 +190,7 @@ class TaskSocketService { const joinSuccess = this.socketService.send({ type: 'join_project', project_id: projectId - }); + }, true); // Enable operation tracking if (!joinSuccess) { throw new Error('Failed to send join_project message'); @@ -214,7 +213,7 @@ class TaskSocketService { this.socketService.send({ type: 'leave_project', project_id: this.currentProjectId - }); + }, true); // Enable operation tracking this.currentProjectId = null; } diff --git a/python/src/mcp_server/features/tasks/task_tools.py b/python/src/mcp_server/features/tasks/task_tools.py index 024f44ed..bc1d9ed3 100644 --- a/python/src/mcp_server/features/tasks/task_tools.py +++ b/python/src/mcp_server/features/tasks/task_tools.py @@ -7,7 +7,7 @@ Mirrors the functionality of the original manage_task tool but with individual t import json import logging -from typing import Any, Dict, List, Optional, TypedDict +from typing import Any, Dict, List, Optional from urllib.parse import urljoin import httpx @@ -20,19 +20,6 @@ from src.server.config.service_discovery import get_api_url logger = logging.getLogger(__name__) -class TaskUpdateFields(TypedDict, total=False): - """Valid fields that can be updated on a task.""" - - title: str - description: str - status: str # "todo" | "doing" | "review" | "done" - assignee: str # "User" | "Archon" | "AI IDE Agent" | "prp-executor" | "prp-validator" - task_order: int # 0-100, higher = more priority - feature: Optional[str] - sources: Optional[List[Dict[str, str]]] - code_examples: Optional[List[Dict[str, str]]] - - def register_task_tools(mcp: FastMCP): """Register individual task management tools with the MCP server.""" @@ -315,26 +302,66 @@ def register_task_tools(mcp: FastMCP): async def update_task( ctx: Context, task_id: str, - update_fields: TaskUpdateFields, + title: Optional[str] = None, + description: Optional[str] = None, + status: Optional[str] = None, + assignee: Optional[str] = None, + task_order: Optional[int] = None, + feature: Optional[str] = None, + sources: Optional[List[Dict[str, str]]] = None, + code_examples: Optional[List[Dict[str, str]]] = None, ) -> str: """ Update a task's properties. Args: task_id: UUID of the task to update - update_fields: Dict of fields to update (e.g., {"status": "doing", "assignee": "AI IDE Agent"}) + title: New task title (optional) + description: New task description (optional) + status: New status - "todo" | "doing" | "review" | "done" (optional) + assignee: New assignee (optional) + task_order: New priority order (optional) + feature: New feature label (optional) + sources: New source references (optional) + code_examples: New code examples (optional) Returns: JSON with updated task details Examples: - update_task(task_id="uuid", update_fields={"status": "doing"}) - update_task(task_id="uuid", update_fields={"title": "New Title", "description": "Updated description"}) + update_task(task_id="uuid", status="doing") + update_task(task_id="uuid", title="New Title", description="Updated description") """ try: api_url = get_api_url() timeout = get_default_timeout() + # Build update_fields dict from provided parameters + update_fields = {} + if title is not None: + update_fields["title"] = title + if description is not None: + update_fields["description"] = description + if status is not None: + update_fields["status"] = status + if assignee is not None: + update_fields["assignee"] = assignee + if task_order is not None: + update_fields["task_order"] = task_order + if feature is not None: + update_fields["feature"] = feature + if sources is not None: + update_fields["sources"] = sources + if code_examples is not None: + update_fields["code_examples"] = code_examples + + if not update_fields: + return MCPErrorFormatter.format_error( + error_type="validation_error", + message="No fields provided to update", + suggestion="Provide at least one field to update", + ) + async with httpx.AsyncClient(timeout=timeout) as client: response = await client.put( urljoin(api_url, f"/api/tasks/{task_id}"), json=update_fields