metona-ai-desktop/src/hooks/useAgentStream.ts

/**
 * useAgentStream — Agent 流式事件监听 Hook
 *
 * 监听 window.metona.agent.onStreamEvent，将事件分发到 Zustand Store。
 * 负责流式文本追加、工具调用状态更新、思考内容处理。
 */

import { useEffect, useRef } from 'react';
import { useAgentStore, type ToolCallInfo } from '@renderer/stores/agent-store';
import { useSessionStore } from '@renderer/stores/session-store';

/**
 * Agent 流式事件监听 Hook
 *
 * 在 App 根组件调用一次，自动监听当前会话的流式事件。
 */
export function useAgentStream(): void {
  const cleanupRef = useRef<(() => void) | null>(null);

  useEffect(() => {
    // 检查 IPC 桥是否可用
    if (!window.metona?.agent?.onStreamEvent) return;

    const unsubscribe = window.metona.agent.onStreamEvent((event: unknown) => {
      const data = event as {
        type?: string;
        requestId?: string;
        sessionId?: string;
        iteration?: number;
        delta?: string;
        content?: string;
        toolCall?: { id: string; name: string; args: Record<string, unknown> };
        toolResult?: { toolCallId: string; success: boolean; result?: unknown; error?: string; durationMs?: number };
        usage?: { inputTokens?: number; outputTokens?: number; totalTokens?: number };
        error?: { code: string; message: string };
        state?: string;
      };

      // 每次都从 store 读取最新状态（避免闭包捕获过期快照）
      const getStore = () => useAgentStore.getState();

      // 每次收到迭代号时同步更新
      if (data.iteration != null && data.iteration !== getStore().currentIteration) {
        console.log(`[useAgentStream] iteration: ${getStore().currentIteration} → ${data.iteration} (event: ${data.type})`);
        getStore().setCurrentIteration(data.iteration);
      }

      switch (data.type) {
        // 思考开始
        case 'thinking_start':
          getStore().setAgentStatus('thinking');
          getStore().addTraceStep({
            iteration: data.iteration ?? getStore().currentIteration + 1,
            state: 'THINKING',
            startedAt: Date.now(),
          });
          break;

        // 推理内容增量
        case 'reasoning_delta':
          if (data.delta) {
            const messages = getStore().messages;
            const lastMsg = messages[messages.length - 1];
            if (lastMsg?.role === 'assistant') {
              // 追加到最后一条 assistant 消息
              getStore().updateMessage(lastMsg.id, {
                reasoningContent: (lastMsg.reasoningContent ?? '') + data.delta,
              });
            } else {
              // 还没有 assistant 消息，先创建一条（仅含思考内容）
              getStore().addMessage({
                id: `msg_${Date.now()}_assistant`,
                role: 'assistant',
                content: '',
                reasoningContent: data.delta,
                timestamp: Date.now(),
              });
            }
          }
          break;

        // 思考结束
        case 'thinking_end':
          if (data.iteration) {
            getStore().updateTraceStep(data.iteration, { completedAt: Date.now() });
          }
          break;

        // 文本增量
        case 'text_delta':
          if (data.delta) {
            getStore().setStreaming(true);
            getStore().updateLastAssistantMessage(data.delta);
          }
          break;

        // 工具调用完成
        case 'tool_call_complete':
          if (data.toolCall) {
            const tc: ToolCallInfo = {
              id: data.toolCall.id,
              name: data.toolCall.name,
              args: data.toolCall.args,
              status: 'pending',
            };

            // 追加到当前 assistant 消息
            const msgs = getStore().messages;
            const lastAssistant = msgs[msgs.length - 1];
            if (lastAssistant?.role === 'assistant') {
              getStore().updateMessage(lastAssistant.id, {
                toolCalls: [...(lastAssistant.toolCalls ?? []), tc],
              });
            }

            getStore().setAgentStatus('executing');
            const curIter = getStore().currentIteration;
            getStore().addTraceStep({
              iteration: data.iteration ?? curIter,
              state: 'EXECUTING',
              startedAt: Date.now(),
              toolCalls: [tc],
            });
          }
          break;

        // 工具执行结果
        case 'tool_result': {
          if (data.toolResult) {
            const msgs = getStore().messages;
            const lastMsg = msgs[msgs.length - 1];
            if (lastMsg?.toolCalls) {
              const updatedToolCalls = lastMsg.toolCalls.map((tc) =>
                tc.id === data.toolResult!.toolCallId
                  ? {
                      ...tc,
                      status: data.toolResult!.success ? 'success' as const : 'error' as const,
                      result: data.toolResult!.result,
                      error: data.toolResult!.error,
                      durationMs: data.toolResult!.durationMs,
                    }
                  : tc,
              );
              getStore().updateMessage(lastMsg.id, { toolCalls: updatedToolCalls });
            }
          }
          break;
        }

        // Token 使用统计
        case 'usage':
          if (data.usage) {
            getStore().updateTokenUsage({
              inputTokens: data.usage.inputTokens ?? 0,
              outputTokens: data.usage.outputTokens ?? 0,
              totalTokens: data.usage.totalTokens ?? 0,
            });
          }
          break;

        // 流结束
        case 'done':
          getStore().setStreaming(false);
          getStore().setAgentStatus('idle');
          getStore().saveTraceData();
          break;

        // 错误
        case 'error':
          getStore().setStreaming(false);
          getStore().setAgentStatus('error');
          getStore().addMessage({
            id: `msg_${Date.now()}_error`,
            role: 'system',
            content: `错误: ${data.error?.message ?? '未知错误'}`,
            timestamp: Date.now(),
          });
          break;

        // 状态变化
        case 'state_change':
          if (data.state) {
            getStore().addTraceStep({
              iteration: data.iteration ?? getStore().currentIteration,
              state: data.state,
              startedAt: Date.now(),
            });
          }
          break;
      }
    });

    cleanupRef.current = unsubscribe;

    return () => {
      cleanupRef.current?.();
    };
  }, []);

  // 监听状态变化事件（含迭代号）
  useEffect(() => {
    if (!window.metona?.agent?.onStateChange) return;

    const unsubscribe = window.metona.agent.onStateChange((state: unknown) => {
      const data = state as { sessionId?: string; iteration?: number; state?: string; previous?: string; current?: string };
      if (data.iteration != null) {
        const store = useAgentStore.getState();
        if (data.iteration !== store.currentIteration) {
          store.setCurrentIteration(data.iteration);
        }
      }
    });

    return () => unsubscribe();
  }, []);

  // 监听 Provider 切换通知
  useEffect(() => {
    if (!window.metona?.agent?.onProviderSwitched) return;

    const unsubscribe = window.metona.agent.onProviderSwitched((data: unknown) => {
      const { from, to, reason } = data as { from?: string; to?: string; reason?: string };
      useAgentStore.getState().addMessage({
        id: `msg_${Date.now()}_system`,
        role: 'system',
        content: `Provider 已切换: ${from ?? '未知'} → ${to ?? '未知'}${reason ? ` (${reason})` : ''}`,
        timestamp: Date.now(),
      });
    });

    return () => unsubscribe();
  }, []);
}