2 files changed, 310 insertions, 0 deletions
diff --git a/makima/frontend/src/hooks/useDirectiveMemories.ts b/makima/frontend/src/hooks/useDirectiveMemories.ts
new file mode 100644
index 0000000..3844c44
--- /dev/null
+++ b/makima/frontend/src/hooks/useDirectiveMemories.ts
@@ -0,0 +1,119 @@
+import { useState, useEffect, useCallback } from "react";
+import {
+  type DirectiveMemoryEntry,
+  type DirectiveMemoryConfig,
+  type MemoryCategory,
+  type CreateDirectiveMemoryRequest,
+  getDirectiveMemoryConfig,
+  setDirectiveMemoryEnabled,
+  listDirectiveMemories,
+  addDirectiveMemory,
+  deleteDirectiveMemory,
+  clearDirectiveMemories,
+} from "../lib/api";
+
+export function useDirectiveMemories(directiveId: string | undefined) {
+  const [memories, setMemories] = useState<DirectiveMemoryEntry[]>([]);
+  const [config, setConfig] = useState<DirectiveMemoryConfig | null>(null);
+  const [loading, setLoading] = useState(false);
+  const [error, setError] = useState<string | null>(null);
+
+  const refreshConfig = useCallback(async () => {
+    if (!directiveId) return;
+    try {
+      const c = await getDirectiveMemoryConfig(directiveId);
+      setConfig(c);
+    } catch (e) {
+      // Config may not exist yet — treat as disabled
+      setConfig({ directiveId, enabled: false, updatedAt: new Date().toISOString() });
+    }
+  }, [directiveId]);
+
+  const refreshMemories = useCallback(async () => {
+    if (!directiveId) return;
+    try {
+      setLoading(true);
+      setError(null);
+      const entries = await listDirectiveMemories(directiveId);
+      setMemories(entries);
+    } catch (e) {
+      setError(e instanceof Error ? e.message : "Failed to load memories");
+    } finally {
+      setLoading(false);
+    }
+  }, [directiveId]);
+
+  const refresh = useCallback(async () => {
+    await Promise.all([refreshConfig(), refreshMemories()]);
+  }, [refreshConfig, refreshMemories]);
+
+  useEffect(() => {
+    refresh();
+  }, [refresh]);
+
+  const toggleEnabled = useCallback(async (enabled: boolean) => {
+    if (!directiveId) return;
+    try {
+      setError(null);
+      const c = await setDirectiveMemoryEnabled(directiveId, enabled);
+      setConfig(c);
+    } catch (e) {
+      setError(e instanceof Error ? e.message : "Failed to toggle memory");
+    }
+  }, [directiveId]);
+
+  const add = useCallback(async (req: CreateDirectiveMemoryRequest) => {
+    if (!directiveId) return;
+    try {
+      setError(null);
+      await addDirectiveMemory(directiveId, req);
+      await refreshMemories();
+    } catch (e) {
+      setError(e instanceof Error ? e.message : "Failed to add memory");
+    }
+  }, [directiveId, refreshMemories]);
+
+  const remove = useCallback(async (memoryId: string) => {
+    if (!directiveId) return;
+    try {
+      setError(null);
+      await deleteDirectiveMemory(directiveId, memoryId);
+      await refreshMemories();
+    } catch (e) {
+      setError(e instanceof Error ? e.message : "Failed to delete memory");
+    }
+  }, [directiveId, refreshMemories]);
+
+  const clearAll = useCallback(async () => {
+    if (!directiveId) return;
+    try {
+      setError(null);
+      await clearDirectiveMemories(directiveId);
+      setMemories([]);
+    } catch (e) {
+      setError(e instanceof Error ? e.message : "Failed to clear memories");
+    }
+  }, [directiveId]);
+
+  /** Group entries by category */
+  const grouped = memories.reduce<Record<MemoryCategory, DirectiveMemoryEntry[]>>(
+    (acc, entry) => {
+      acc[entry.category].push(entry);
+      return acc;
+    },
+    { decision: [], context: [], preference: [], learning: [], other: [] },
+  );
+
+  return {
+    memories,
+    grouped,
+    config,
+    loading,
+    error,
+    refresh,
+    toggleEnabled,
+    add,
+    remove,
+    clearAll,
+  };
+}
diff --git a/makima/frontend/src/hooks/useMultiTaskSubscription.ts b/makima/frontend/src/hooks/useMultiTaskSubscription.ts
new file mode 100644
index 0000000..19d6dea
--- /dev/null
+++ b/makima/frontend/src/hooks/useMultiTaskSubscription.ts
@@ -0,0 +1,191 @@
+import { useState, useCallback, useRef, useEffect, useMemo } from "react";
+import { TASK_SUBSCRIBE_ENDPOINT } from "../lib/api";
+import type { TaskOutputEvent } from "./useTaskSubscription";
+
+export interface MultiTaskOutputEntry extends TaskOutputEvent {
+  /** Label for the task (e.g. step name or "Orchestrator") */
+  taskLabel: string;
+  /** Timestamp when the entry was received */
+  receivedAt: number;
+}
+
+interface UseMultiTaskSubscriptionOptions {
+  /** Map of taskId -> label */
+  taskMap: Map<string, string>;
+  /** Whether to actively subscribe */
+  enabled?: boolean;
+  /** Max entries to keep in buffer */
+  maxEntries?: number;
+}
+
+export function useMultiTaskSubscription(options: UseMultiTaskSubscriptionOptions) {
+  const { taskMap, enabled = true, maxEntries = 2000 } = options;
+
+  const [connected, setConnected] = useState(false);
+  const [entries, setEntries] = useState<MultiTaskOutputEntry[]>([]);
+  const wsRef = useRef<WebSocket | null>(null);
+  const reconnectTimeoutRef = useRef<number | null>(null);
+  const subscribedTasksRef = useRef<Set<string>>(new Set());
+  const taskMapRef = useRef(taskMap);
+  const enabledRef = useRef(enabled);
+
+  // Keep refs in sync
+  useEffect(() => {
+    taskMapRef.current = taskMap;
+  }, [taskMap]);
+
+  useEffect(() => {
+    enabledRef.current = enabled;
+  }, [enabled]);
+
+  // Derive task IDs from the map
+  const taskIds = useMemo(() => Array.from(taskMap.keys()), [taskMap]);
+
+  const subscribeToTask = useCallback((ws: WebSocket, taskId: string) => {
+    if (ws.readyState === WebSocket.OPEN) {
+      ws.send(JSON.stringify({ type: "subscribeOutput", taskId }));
+      subscribedTasksRef.current.add(taskId);
+    }
+  }, []);
+
+  const unsubscribeFromTask = useCallback((ws: WebSocket, taskId: string) => {
+    if (ws.readyState === WebSocket.OPEN) {
+      ws.send(JSON.stringify({ type: "unsubscribeOutput", taskId }));
+      subscribedTasksRef.current.delete(taskId);
+    }
+  }, []);
+
+  const connect = useCallback(() => {
+    const currentState = wsRef.current?.readyState;
+    if (currentState === WebSocket.OPEN || currentState === WebSocket.CONNECTING) {
+      return;
+    }
+
+    if (wsRef.current && currentState === WebSocket.CLOSING) {
+      wsRef.current = null;
+    }
+
+    try {
+      const ws = new WebSocket(TASK_SUBSCRIBE_ENDPOINT);
+      wsRef.current = ws;
+
+      ws.onopen = () => {
+        setConnected(true);
+        // Re-subscribe to all tasks
+        for (const taskId of subscribedTasksRef.current) {
+          ws.send(JSON.stringify({ type: "subscribeOutput", taskId }));
+        }
+      };
+
+      ws.onmessage = (event) => {
+        try {
+          const message = JSON.parse(event.data);
+
+          if (message.type === "taskOutput") {
+            const label = taskMapRef.current.get(message.taskId) || message.taskId;
+            const entry: MultiTaskOutputEntry = {
+              taskId: message.taskId,
+              messageType: message.messageType,
+              content: message.content,
+              toolName: message.toolName,
+              toolInput: message.toolInput,
+              isError: message.isError,
+              costUsd: message.costUsd,
+              durationMs: message.durationMs,
+              isPartial: message.isPartial,
+              taskLabel: label,
+              receivedAt: Date.now(),
+            };
+
+            setEntries((prev) => {
+              const next = [...prev, entry];
+              if (next.length > maxEntries) {
+                return next.slice(next.length - maxEntries);
+              }
+              return next;
+            });
+          }
+        } catch (e) {
+          console.error("Failed to parse multi-task subscription message:", e);
+        }
+      };
+
+      ws.onerror = () => {
+        console.error("Multi-task WebSocket connection error");
+      };
+
+      ws.onclose = () => {
+        setConnected(false);
+        wsRef.current = null;
+
+        // Reconnect if we still have subscriptions
+        if (subscribedTasksRef.current.size > 0 && enabledRef.current) {
+          reconnectTimeoutRef.current = window.setTimeout(() => {
+            connect();
+          }, 3000);
+        }
+      };
+    } catch (e) {
+      console.error("Failed to connect multi-task subscription:", e);
+    }
+  }, [maxEntries]);
+
+  // Manage subscriptions when task IDs change
+  useEffect(() => {
+    if (!enabled || taskIds.length === 0) {
+      // Close connection if no tasks
+      if (wsRef.current) {
+        subscribedTasksRef.current.clear();
+        wsRef.current.close();
+        wsRef.current = null;
+      }
+      return;
+    }
+
+    const newTaskIds = new Set(taskIds);
+    const ws = wsRef.current;
+
+    if (!ws || ws.readyState !== WebSocket.OPEN) {
+      // Set desired subscriptions and connect
+      subscribedTasksRef.current = newTaskIds;
+      connect();
+      return;
+    }
+
+    // Unsubscribe from removed tasks
+    for (const existingId of subscribedTasksRef.current) {
+      if (!newTaskIds.has(existingId)) {
+        unsubscribeFromTask(ws, existingId);
+      }
+    }
+
+    // Subscribe to new tasks
+    for (const newId of newTaskIds) {
+      if (!subscribedTasksRef.current.has(newId)) {
+        subscribeToTask(ws, newId);
+      }
+    }
+  }, [taskIds, enabled, connect, subscribeToTask, unsubscribeFromTask]);
+
+  // Cleanup on unmount
+  useEffect(() => {
+    return () => {
+      if (reconnectTimeoutRef.current) {
+        clearTimeout(reconnectTimeoutRef.current);
+      }
+      if (wsRef.current) {
+        wsRef.current.close();
+      }
+    };
+  }, []);
+
+  const clearEntries = useCallback(() => {
+    setEntries([]);
+  }, []);
+
+  return {
+    connected,
+    entries,
+    clearEntries,
+  };
+}