AstroResearch/dashboard/src/components/agent/subagentRouting.ts

// dashboard/src/components/agent/subagentRouting.ts
// 子代理路由逻辑（纯函数，TimelineItem[] 操作）
// ResearchAgentPanel 和 AIAssistantPanel 共用
import type { TimelineItem } from './types';

// ── 前缀检测 ──

export const SUBAGENT_CONTAINER_NAMES = ['subagent', 'delegate_research'];

export function isSubagentThought(content: string): boolean {
  return content.startsWith('[子代理] ');
}

export function cleanSubagentThought(content: string): string {
  return content.slice(6); // 去掉 "[子代理] " 前缀
}

export function isSubagentContainerCall(name: string): boolean {
  return SUBAGENT_CONTAINER_NAMES.includes(name);
}

export function isSubagentChildCall(name: string): boolean {
  return name.startsWith('[sub] ');
}

export function cleanSubagentChildName(name: string): string {
  return name.slice(6);
}

export function isSubagentChildResult(name: string): boolean {
  return name.startsWith('[sub] ');
}

// ── 容器查找 ──

// 找到最后一个 streaming 状态的 subagent_container 索引
export function findStreamingSubAgent(timeline: TimelineItem[]): number {
  for (let i = timeline.length - 1; i >= 0; i--) {
    const t = timeline[i];
    if (t.type === 'subagent_container' && t.status === 'streaming') {
      return i;
    }
  }
  return -1;
}

// ── 操作函数：接受 timeline → 返回 new timeline（不可变风格）──

// 路由 thought 事件
export function routeThought(
  timeline: TimelineItem[],
  step: number,
  content: string
): TimelineItem[] {
  // 子代理思考 → 路由到活跃容器
  if (isSubagentThought(content)) {
    const subContent = cleanSubagentThought(content);
    const saIdx = findStreamingSubAgent(timeline);
    if (saIdx >= 0) {
      return updateContainerChild(timeline, saIdx, (children) => {
        const existing = children.find(
          (c) => c.type === 'thought' && c.step === step
        );
        if (existing) {
          return children.map((c) =>
            c.type === 'thought' && c.step === step
              ? ({ ...c, content: subContent } as TimelineItem)
              : c
          );
        }
        return [...children, { type: 'thought', step, content: subContent }];
      });
    }
    // 无活跃容器 → fallback 为父代理 thought
    return upsertParentThought(timeline, step, subContent);
  }
  // 父代理思考
  return upsertParentThought(timeline, step, content);
}

// 路由 tool_call 事件
export function routeToolCall(
  timeline: TimelineItem[],
  step: number,
  id: string,
  name: string,
  args: Record<string, unknown>
): TimelineItem[] {
  // 创建 subagent_container
  if (isSubagentContainerCall(name)) {
    const existing = timeline.find(
      (t) => t.type === 'subagent_container' && t.id === id
    );
    if (existing) return timeline;
    return [
      ...timeline,
      {
        type: 'subagent_container' as const,
        id,
        status: 'streaming' as const,
        children: [],
      },
    ];
  }

  // 子代理工具 → 路由到活跃容器
  if (isSubagentChildCall(name)) {
    const cleanName = cleanSubagentChildName(name);
    const saIdx = findStreamingSubAgent(timeline);
    if (saIdx >= 0) {
      return updateContainerChild(timeline, saIdx, (children) => {
        const dup = children.find(
          (c) => c.type === 'tool_call' && 'id' in c && c.id === id
        );
        if (dup) return children;
        return [
          ...children,
          {
            type: 'tool_call' as const,
            step,
            id,
            name: cleanName,
            arguments: args,
          },
        ];
      });
    }
    // 无活跃容器 → fallback
    return upsertParentToolCall(timeline, step, id, cleanName, args);
  }

  // 父代理工具调用
  return upsertParentToolCall(timeline, step, id, name, args);
}

// 路由 tool_result 事件
export function routeToolResult(
  timeline: TimelineItem[],
  toolCallId: string,
  name: string,
  output: string,
  isError: boolean,
  metadata?: Record<string, unknown>
): TimelineItem[] {
  // 匹配 subagent_container id → 完成容器
  const saCompleteIdx = timeline.findIndex(
    (t) => t.type === 'subagent_container' && t.id === toolCallId
  );
  if (saCompleteIdx >= 0) {
    const newTimeline = [...timeline];
    const container = newTimeline[saCompleteIdx] as Extract<
      TimelineItem,
      { type: 'subagent_container' }
    >;
    newTimeline[saCompleteIdx] = {
      ...container,
      status: 'complete' as const,
      summary: output,
    };
    return newTimeline;
  }

  // 子代理工具结果 → 路由到活跃容器
  if (isSubagentChildResult(name)) {
    const saIdx = findStreamingSubAgent(timeline);
    if (saIdx >= 0) {
      return updateContainerChild(timeline, saIdx, (children) =>
        children.map((c) => {
          if (c.type === 'tool_call' && c.id === toolCallId && !c.result) {
            return { ...c, result: { output, isError, metadata } };
          }
          return c;
        })
      );
    }
  }

  // 父代理工具结果
  return timeline.map((t) => {
    if (t.type === 'tool_call' && t.id === toolCallId) {
      return { ...t, result: { output, isError, metadata } };
    }
    return t;
  });
}

// 路由 text_delta 事件（可携带 tool_call_id 流式输出到工具结果）
export function routeTextDelta(
  timeline: TimelineItem[],
  content: string,
  toolCallId?: string
): { timeline: TimelineItem[]; answerDelta: string } {
  if (toolCallId) {
    return {
      timeline: timeline.map((t) => {
        if (t.type === 'tool_call' && t.id === toolCallId) {
          const prevOutput = t.result?.output || '';
          return {
            ...t,
            result: { output: prevOutput + content, isError: false },
          };
        }
        return t;
      }),
      answerDelta: '',
    };
  }
  return { timeline, answerDelta: content };
}

// ── 内部 helpers ──

function upsertParentThought(
  timeline: TimelineItem[],
  step: number,
  content: string
): TimelineItem[] {
  const existing = timeline.find(
    (t) => t.type === 'thought' && t.step === step
  );
  if (existing) {
    return timeline.map((t) =>
      t.type === 'thought' && t.step === step ? { ...t, content } : t
    );
  }
  return [...timeline, { type: 'thought', step, content }];
}

function upsertParentToolCall(
  timeline: TimelineItem[],
  step: number,
  id: string,
  name: string,
  args: Record<string, unknown>
): TimelineItem[] {
  const existing = timeline.find(
    (t) => t.type === 'tool_call' && 'id' in t && t.id === id
  );
  if (existing) return timeline;
  return [
    ...timeline,
    { type: 'tool_call' as const, step, id, name, arguments: args },
  ];
}

function updateContainerChild(
  timeline: TimelineItem[],
  containerIdx: number,
  updateFn: (children: TimelineItem[]) => TimelineItem[]
): TimelineItem[] {
  const newTimeline = [...timeline];
  const container = newTimeline[containerIdx] as Extract<
    TimelineItem,
    { type: 'subagent_container' }
  >;
  newTimeline[containerIdx] = {
    ...container,
    children: updateFn(container.children),
  };
  return newTimeline;
}