packages/kap-server/src/services/transcript/coreBinding.ts

import {
  IAgentLifecycleService,
  IAgentConversationUndoParticipantRegistry,
  IAgentLoopService,
  IAgentScopeContext,
  IAgentTaskService,
  IEventBus,
  INTERACTION_TAG_AGENT_ID,
  INTERACTION_TAG_SESSION_ID,
  ISessionMetadata,
  MAIN_AGENT_ID,
  interactions,
  toDisposable,
  type AgentMeta,
  type IDisposable,
  type IAgentScopeHandle,
  type Interaction,
  type ISessionScopeHandle,
} from '@moonshot-ai/agent-core-v2';
import type { AgentDescriptor, TranscriptChangeEvent, TranscriptStore } from '@moonshot-ai/transcript';

import { legacyApprovalsOf } from '../legacyStatus/legacyActivity';
import {
  AgentTranscriptProjector,
  type ProjectorBusEvent,
  type ProjectorInteraction,
} from './coreEventMap';

export interface TranscriptBindingLogger {
  warn(obj: unknown, msg: string): void;
}

export interface TranscriptBinding extends IDisposable {
  seedPendingInteractions(agentId?: string): void;
}

export function bindSessionTranscript(
  store: TranscriptStore,
  session: ISessionScopeHandle,
  logger?: TranscriptBindingLogger,
  onOps?: (event: TranscriptChangeEvent) => void,
  reconcileAfterUndo?: (agentId: string) => Promise<void>,
): TranscriptBinding {
  const agents = session.accessor.get(IAgentLifecycleService);
  const pendingInteractions = (): readonly Interaction[] =>
    interactions.findAll({
      resolved: false,
      tags: { [INTERACTION_TAG_SESSION_ID]: session.id },
    });
  const disposables: IDisposable[] = [];
  const agentDisposables = new Map<string, IDisposable[]>();
  const subscribedAgents = new Set<string>();
  const projectors = new Map<string, AgentTranscriptProjector>();
  const interactionAgents = new Map<string, string>();
  const knownInteractions = new Set<string>();
  const unseeded = new Map<string, Interaction>();
  const earlyResolves = new Map<string, { agentId: string; response: unknown }>();
  const seededAgents = new Set<string>();
  let seededAll = false;
  const isSeeded = (agentId: string): boolean => seededAll || seededAgents.has(agentId);

  const applyOps = (agentId: string, ops: ReturnType<AgentTranscriptProjector['map']>): void => {
    if (ops.length === 0) return;
    const result = store.ensureAgent(agentId).apply(ops);
    if (result.gap !== undefined) {
      logger?.warn(
        { sessionId: store.sessionId, agentId, gap: result.gap },
        'transcript: append gap — producer/consumer skew',
      );
      return;
    }
    onOps?.({ agentId, ops });
  };

  const projectorFor = (agentId: string): AgentTranscriptProjector => {
    let projector = projectors.get(agentId);
    if (projector === undefined) {
      projector = new AgentTranscriptProjector(agentId, store.sessionId, {
        stepFrames: (turnId, stepId) =>
          store.getAgent(agentId)?.getTurn(turnId)?.steps.find((s) => s.stepId === stepId)?.frames,
        toolFrame: (toolCallId) => {
          const transcript = store.getAgent(agentId);
          if (transcript === undefined) return undefined;
          for (const item of transcript.getItems()) {
            if (item.kind !== 'turn') continue;
            for (const step of item.steps) {
              for (const frame of step.frames) {
                if (frame.kind === 'tool' && frame.toolCallId === toolCallId) {
                  return { turnId: item.turnId, stepId: step.stepId, frame };
                }
              }
            }
          }
          return undefined;
        },
        stepOrdinal: (turnId) => {
          const agentHandle = agents.handleOf(agentId);
          if (agentHandle === undefined) return undefined;
          const turn = agentHandle.accessor.get(IAgentLoopService)?.snapshot().turn;
          return turn === undefined || `t${turn.turnId}` !== turnId ? undefined : turn.step;
        },
        activitySnapshot: () =>
          agents.handleOf(agentId)?.accessor.get(IAgentLoopService)?.snapshot() ?? {},
        pendingApprovals: () => {
          const agentHandle = agents.handleOf(agentId);
          return agentHandle === undefined ? [] : legacyApprovalsOf(agentHandle);
        },
        turn: (turnId) => store.getAgent(agentId)?.getTurn(turnId),
        resolvePlanRevisionKey: (key) =>
          agents.handleOf(agentId)?.accessor.get(IAgentScopeContext).scope(key) ?? key,
      });
      const agentHandle = agents.handleOf(agentId);
      if (agentHandle !== undefined) {
        const tasks = agentHandle.accessor.get(IAgentTaskService)?.list() ?? [];
        for (const info of tasks) {
          if (info.kind === 'agent' && typeof info.agentId === 'string' && info.agentId.length > 0) {
            applyOps(
              agentId,
              projector.seedSubagentTask({
                taskId: info.taskId,
                agentId: info.agentId,
                description: info.description,
                status: info.status,
                detached: info.detached ?? false,
                startedAt: info.startedAt,
              }),
            );
          }
        }
      }
      projectors.set(agentId, projector);
    }
    return projector;
  };

  const subscribeAgent = (handle: IAgentScopeHandle): void => {
    if (subscribedAgents.has(handle.id)) return;
    subscribedAgents.add(handle.id);
    const projector = projectorFor(handle.id);
    store.ensureAgent(handle.id, { agentId: handle.id });
    const bus = handle.accessor.get(IEventBus);
    const busD = bus.subscribe((event) =>
      applyOps(handle.id, projectorFor(handle.id).map(event as ProjectorBusEvent)),
    );
    const loopStatus = handle.accessor.get(IAgentLoopService)?.snapshot();
    if (loopStatus?.state === 'running' && loopStatus.activeTurnId !== undefined) {
      const promptId = loopStatus.activePromptId;
      projector.seedActiveTurn({ turnId: loopStatus.activeTurnId, promptId });
    }
    const list = agentDisposables.get(handle.id) ?? [];
    list.push(busD);
    if (reconcileAfterUndo !== undefined) {
      list.push(handle.accessor.get(IAgentConversationUndoParticipantRegistry).register({
        id: 'transcript',
        phase: 'after-flush',
        reconcileAfterUndo: async () => {
          await reconcileAfterUndo(handle.id);
          projectors.delete(handle.id);
        },
      }));
    }
    agentDisposables.set(handle.id, list);
  };

  const interactionAgentId = (interaction: Interaction): string => {
    const payloadAgent = (interaction.payload as { agentId?: unknown }).agentId;
    const tag = interaction.tags[INTERACTION_TAG_AGENT_ID];
    return (
      (typeof tag === 'string' ? tag : undefined) ??
      (typeof payloadAgent === 'string' ? payloadAgent : undefined) ??
      MAIN_AGENT_ID
    );
  };

  const announceInteraction = (interaction: Interaction): void => {
    if (interaction.kind !== 'approval' && interaction.kind !== 'question') return;
    const agentId = interactionAgentId(interaction);
    interactionAgents.set(interaction.id, agentId);
    const request: ProjectorInteraction = {
      id: interaction.id,
      kind: interaction.kind,
      payload: interaction.payload,
      createdAt: interaction.createdAt,
    };
    applyOps(agentId, projectorFor(agentId).mapInteractionRequested(request));
  };

  const refreshDescriptors = (): void => {
    void session.accessor
      .get(ISessionMetadata)
      .read()
      .then((meta) => {
        for (const agentId of projectors.keys()) {
          store.describeAgent(descriptorFromMeta(agentId, meta.agents?.[agentId]));
        }
      })
      .catch(() => {
      });
  };

  for (const agent of agents.list()) {
    const handle = agents.handleOf(agent.agentId);
    if (handle !== undefined) subscribeAgent(handle);
  }
  disposables.push(
    agents.onDidCreate((context) => {
      const handle = agents.handleOf(context.agentId);
      if (handle !== undefined) subscribeAgent(handle);
      seededAgents.add(context.agentId);
      refreshDescriptors();
    }),
    agents.onDidClose((context) => {
      const agentId = context.agentId;
      for (const d of agentDisposables.get(agentId) ?? []) d.dispose();
      agentDisposables.delete(agentId);
      subscribedAgents.delete(agentId);
      projectors.delete(agentId);
      store.markDisposed(agentId, new Date().toISOString());
    }),
  );

  for (const pending of pendingInteractions()) {
    if (pending.kind !== 'approval' && pending.kind !== 'question') continue;
    if (knownInteractions.has(pending.id)) continue;
    knownInteractions.add(pending.id);
    interactionAgents.set(pending.id, interactionAgentId(pending));
    unseeded.set(pending.id, pending);
  }
  const seedPendingInteractions = (agentId?: string): void => {
    if (agentId === undefined) seededAll = true;
    else seededAgents.add(agentId);
    for (const [id, interaction] of unseeded) {
      if (agentId !== undefined && interactionAgents.get(id) !== agentId) continue;
      unseeded.delete(id);
      announceInteraction(interaction);
      const early = earlyResolves.get(id);
      if (early === undefined) continue;
      interactionAgents.delete(id);
      earlyResolves.delete(id);
      const projector = projectors.get(early.agentId);
      if (projector !== undefined) {
        applyOps(early.agentId, projector.mapInteractionResolved(id, early.response));
      }
    }
    for (const pending of pendingInteractions()) {
      if (knownInteractions.has(pending.id)) continue;
      if (agentId !== undefined && interactionAgentId(pending) !== agentId) continue;
      knownInteractions.add(pending.id);
      announceInteraction(pending);
    }
  };
  disposables.push(
    toDisposable(
      interactions.onDidChangePending(() => {
        for (const pending of pendingInteractions()) {
          if (knownInteractions.has(pending.id)) continue;
          const agentId = interactionAgentId(pending);
          knownInteractions.add(pending.id);
          if (!isSeeded(agentId)) {
            interactionAgents.set(pending.id, agentId);
            unseeded.set(pending.id, pending);
            continue;
          }
          announceInteraction(pending);
        }
      }),
    ),
    toDisposable(
      interactions.onDidResolve(({ id, response }) => {
        knownInteractions.delete(id);
        const agentId = interactionAgents.get(id);
        if (agentId === undefined) return;
        interactionAgents.delete(id);
        if (unseeded.has(id)) {
          earlyResolves.set(id, { agentId, response });
          return;
        }
        const projector = projectors.get(agentId);
        if (projector === undefined) return;
        applyOps(agentId, projector.mapInteractionResolved(id, response));
      }),
    ),
  );

  refreshDescriptors();

  return {
    seedPendingInteractions,
    dispose: () => {
      for (const d of disposables) d.dispose();
      for (const list of agentDisposables.values()) {
        for (const d of list) d.dispose();
      }
      agentDisposables.clear();
      projectors.clear();
      interactionAgents.clear();
      knownInteractions.clear();
      unseeded.clear();
      earlyResolves.clear();
    },
  };
}

export function descriptorFromMeta(agentId: string, meta: AgentMeta | undefined): AgentDescriptor {
  const parentFromLabels = meta?.labels?.['parentAgentId'];
  const swarmItem = meta?.labels?.['swarmItem'] ?? meta?.swarmItem;
  return {
    agentId,
    type: meta?.type ?? (agentId === MAIN_AGENT_ID ? 'main' : 'sub'),
    parentAgentId:
      parentFromLabels !== undefined && parentFromLabels.length > 0
        ? parentFromLabels
        : (meta?.parentAgentId ?? undefined),
    label: swarmItem !== undefined && swarmItem.length > 0 ? swarmItem : undefined,
  };
}