popagent

Latest change f6fedafe4a36ed6468d5b1b819bbcdcbed5a0d53 - Address a peer by display capitalization, and correct the 9Router runbook by AkurAI Build

import type { AgentExecutionSource } from "./api-types";

export type TaskAgentParticipantState = "idle" | "active" | "waiting" | "terminal";

export type TaskAgentPeer = {
  id: string;
  state: TaskAgentParticipantState;
};

export type TaskAgentCommunicationMessage = {
  id: string;
  taskId: string;
  workspaceId: string;
  resourceId: string;
  from: string;
  to: string;
  body: string;
  replyTo?: string;
  createdAt: string;
};

export type TaskAgentDeliveryReceipt = {
  to: string;
  outcome: "delivered" | "queued" | "woken" | "failed";
  messageId?: string;
  error?: string;
};

type Waiter = {
  from?: string;
  resolve: (message: TaskAgentCommunicationMessage | null) => void;
  reject: (error: Error) => void;
};

export type TaskAgentWakeHandler = (
  recipientId: string,
  message: TaskAgentCommunicationMessage,
) => Promise<void>;

const DEFAULT_MAILBOX_CAP = 100;

export class TaskAgentCommunicationSession {
  readonly taskId: string;
  readonly workspaceId: string;
  readonly resourceId: string;
  readonly #mailboxCap: number;
  readonly #persist?: (message: TaskAgentCommunicationMessage) => Promise<void>;
  readonly #wake?: TaskAgentWakeHandler;
  readonly #participants = new Map<string, TaskAgentParticipantState>();
  readonly #mailboxes = new Map<string, TaskAgentCommunicationMessage[]>();
  readonly #waiters = new Map<string, Waiter[]>();
  readonly #messages = new Map<string, TaskAgentCommunicationMessage>();
  readonly #wakes = new Set<Promise<void>>();
  #executionContext?: {
    executionSource: AgentExecutionSource;
    selfUpdateWorkspacePath?: string;
  };
  #closed = false;
  #closeReason = new Error("Task communication session closed");

  constructor(options: {
    taskId: string;
    workspaceId: string;
    resourceId: string;
    participants: readonly string[];
    mailboxCap?: number;
    signal?: AbortSignal;
    persist?: (message: TaskAgentCommunicationMessage) => Promise<void>;
    wake?: TaskAgentWakeHandler;
  }) {
    const mailboxCap = options.mailboxCap ?? DEFAULT_MAILBOX_CAP;
    if (!Number.isInteger(mailboxCap) || mailboxCap < 1) {
      throw new RangeError("Task communication mailbox cap must be a positive integer");
    }
    this.taskId = options.taskId;
    this.workspaceId = options.workspaceId;
    this.resourceId = options.resourceId;
    this.#mailboxCap = mailboxCap;
    this.#persist = options.persist;
    this.#wake = options.wake;
    for (const id of options.participants) {
      if (id && !this.#participants.has(id)) this.#participants.set(id, "idle");
    }
    if (options.signal?.aborted) {
      this.close(options.signal.reason);
    } else {
      options.signal?.addEventListener("abort", () => this.close(options.signal!.reason), { once: true });
    }
  }

  list(agentId: string): TaskAgentPeer[] {
    this.#activate(agentId);
    return [...this.#participants]
      .filter(([id, state]) => id !== agentId && state !== "terminal")
      .map(([id, state]) => ({ id, state }));
  }

  configureExecution(context: {
    executionSource: AgentExecutionSource;
    selfUpdateWorkspacePath?: string;
  }): void {
    if (this.#closed) return;
    this.#executionContext = context;
  }

  executionContext(): {
    executionSource: AgentExecutionSource;
    selfUpdateWorkspacePath?: string;
  } | undefined {
    return this.#executionContext;
  }

  hasReply(messageId: string, from: string): boolean {
    return [...this.#messages.values()].some((message) =>
      message.replyTo === messageId && message.from === from);
  }

  async send(rawInput: {
    from: string;
    to: string;
    body: string;
    replyTo?: string;
  }): Promise<TaskAgentDeliveryReceipt> {
    // Agents read display names ("Orchistrator") from the channel and address
    // them verbatim, while participants are keyed by agent id ("orchistrator").
    const input = { ...rawInput, from: this.#canonical(rawInput.from), to: this.#canonical(rawInput.to) };
    const invalid = this.#validateSend(input.from, input.to, input.replyTo);
    if (invalid) return { to: input.to, outcome: "failed", error: invalid };
    const body = input.body.trim();
    if (!body || body.length > 4_000) {
      return { to: input.to, outcome: "failed", error: "Message body must contain 1 to 4000 characters" };
    }
    this.#activate(input.from);
    const message: TaskAgentCommunicationMessage = {
      id: crypto.randomUUID(),
      taskId: this.taskId,
      workspaceId: this.workspaceId,
      resourceId: this.resourceId,
      from: input.from,
      to: input.to,
      body,
      ...(input.replyTo ? { replyTo: input.replyTo } : {}),
      createdAt: new Date().toISOString(),
    };
    try {
      await this.#persist?.(message);
    } catch (error) {
      return {
        to: input.to,
        outcome: "failed",
        error: error instanceof Error ? error.message : String(error),
      };
    }
    if (this.#closed) {
      return { to: input.to, outcome: "failed", error: this.#closeReason.message };
    }
    this.#messages.set(message.id, message);
    const waiter = this.#takeWaiter(input.to, input.from);
    if (waiter) {
      waiter.resolve(message);
      return { to: input.to, outcome: "delivered", messageId: message.id };
    }
    const mailbox = this.#mailboxes.get(input.to) ?? [];
    mailbox.push(message);
    if (mailbox.length > this.#mailboxCap) mailbox.shift();
    this.#mailboxes.set(input.to, mailbox);
    if (!input.replyTo && this.#participants.get(input.to) === "idle" && this.#wake) {
      this.#scheduleWake(message);
      return { to: input.to, outcome: "woken", messageId: message.id };
    }
    return { to: input.to, outcome: "queued", messageId: message.id };
  }

  async broadcast(from: string, body: string): Promise<TaskAgentDeliveryReceipt[]> {
    const senderError = this.#participantError(from);
    if (senderError) return [{ to: "all", outcome: "failed", error: senderError }];
    const recipients = [...this.#participants]
      .filter(([id, state]) => id !== from && state !== "terminal")
      .map(([id]) => id);
    const receipts: TaskAgentDeliveryReceipt[] = [];
    for (const to of recipients) receipts.push(await this.send({ from, to, body }));
    return receipts;
  }

  inbox(agentId: string, options: { peek?: boolean } = {}): TaskAgentCommunicationMessage[] {
    if (this.#participantError(agentId)) return [];
    this.#activate(agentId);
    const mailbox = this.#mailboxes.get(agentId) ?? [];
    if (options.peek) return [...mailbox];
    this.#mailboxes.delete(agentId);
    return mailbox;
  }

  takeContext(agentId: string, limit = 8): TaskAgentCommunicationMessage[] {
    if (this.#participantError(agentId)) return [];
    if (!Number.isInteger(limit) || limit < 1 || limit > 8) {
      throw new RangeError("Task communication context limit must be between 1 and 8");
    }
    this.#activate(agentId);
    const mailbox = this.#mailboxes.get(agentId);
    if (!mailbox?.length) return [];
    const messages = mailbox.splice(0, limit);
    if (!mailbox.length) this.#mailboxes.delete(agentId);
    return messages;
  }

  idle(agentId: string): void {
    if (this.#closed || this.#participants.get(agentId) === "terminal") return;
    this.#participants.set(agentId, "idle");
    const pending = this.#mailboxes.get(agentId)?.[0];
    if (pending && this.#wake) this.#scheduleWake(pending);
  }

  async settle(): Promise<void> {
    while (this.#wakes.size) {
      const wakes = [...this.#wakes];
      try {
        await Promise.all(wakes);
      } finally {
        for (const wake of wakes) this.#wakes.delete(wake);
      }
    }
  }

  async wait(
    agentId: string,
    options: { from?: string; timeoutMs?: number } = {},
  ): Promise<TaskAgentCommunicationMessage | null> {
    const participantError = this.#participantError(agentId);
    if (participantError) throw new Error(participantError);
    if (options.from) {
      const fromError = this.#participantError(options.from);
      if (fromError) throw new Error(fromError);
    }
    const pending = this.#takeMailbox(agentId, options.from);
    if (pending) {
      this.#activate(agentId);

      return pending;
    }
    const timeoutMs = options.timeoutMs ?? 10_000;
    if (!Number.isInteger(timeoutMs) || timeoutMs < 0 || timeoutMs > 30_000) {
      throw new RangeError("Wait timeout must be between 0 and 30000 milliseconds");
    }
    this.#participants.set(agentId, "waiting");
    const { promise, resolve, reject } = Promise.withResolvers<TaskAgentCommunicationMessage | null>();
    let timer: Timer | undefined;
    const cleanup = () => {
      clearTimeout(timer);
      const waiters = this.#waiters.get(agentId);
      if (!waiters) return;
      const index = waiters.indexOf(waiter);
      if (index >= 0) waiters.splice(index, 1);
      if (!waiters.length) this.#waiters.delete(agentId);
    };
    const settleState = () => {
      if (!this.#closed && this.#participants.get(agentId) !== "terminal") {
        this.#participants.set(agentId, "active");
      }
    };
    const waiter: Waiter = {
      from: options.from,
      resolve: (message) => { cleanup(); settleState(); resolve(message); },
      reject: (error) => { cleanup(); settleState(); reject(error); },
    };
    const waiters = this.#waiters.get(agentId) ?? [];
    waiters.push(waiter);
    this.#waiters.set(agentId, waiters);
    if (timeoutMs > 0) {
      timer = setTimeout(() => waiter.resolve(null), timeoutMs);
      timer.unref?.();
    }
    return promise;
  }

  terminal(agentId: string): void {
    if (!this.#participants.has(agentId)) return;
    this.#participants.set(agentId, "terminal");
    this.#mailboxes.delete(agentId);
    const error = new Error(`Agent "${agentId}" is terminal`);
    for (const waiter of [...this.#waiters.get(agentId) ?? []]) waiter.reject(error);
  }

  close(reason?: unknown): void {
    if (this.#closed) return;
    this.#closed = true;
    this.#closeReason = reason instanceof Error
      ? reason
      : new Error(String(reason ?? "Task communication session closed"));
    for (const id of this.#participants.keys()) this.#participants.set(id, "terminal");
    this.#mailboxes.clear();
    for (const waiters of [...this.#waiters.values()]) {
      for (const waiter of [...waiters]) waiter.reject(this.#closeReason);
    }
    this.#waiters.clear();
  }

  #activate(agentId: string): void {
    if (this.#participants.get(agentId) !== "terminal") this.#participants.set(agentId, "active");
  }

  /** Registered id for a case variant, so identity checks still run on one key. */
  #canonical(agentId: string): string {
    if (this.#participants.has(agentId)) return agentId;
    const wanted = agentId.trim().toLowerCase();
    for (const id of this.#participants.keys()) if (id.toLowerCase() === wanted) return id;
    return agentId;
  }

  #participantError(agentId: string): string | undefined {
    if (this.#closed) return `Task communication session is closed: ${this.#closeReason.message}`;
    const state = this.#participants.get(agentId);
    if (!state) return `Unknown task agent "${agentId}"`;
    if (state === "terminal") return `Task agent "${agentId}" is terminal`;
    return undefined;
  }

  #validateSend(from: string, to: string, replyTo?: string): string | undefined {
    const senderError = this.#participantError(from);
    if (senderError) return senderError;
    const recipientError = this.#participantError(to);
    if (recipientError) return recipientError;
    if (from === to) return "Task agents cannot message themselves";
    if (replyTo) {
      const original = this.#messages.get(replyTo);
      if (!original) return `Unknown reply message "${replyTo}"`;
      if (original.from !== to || original.to !== from) {
        return `Only the addressed recipient "${original.to}" can reply to message "${replyTo}"`;
      }
    }
    return undefined;
  }

  #scheduleWake(message: TaskAgentCommunicationMessage): void {
    if (!this.#wake || this.#closed) return;
    const mailbox = this.#mailboxes.get(message.to);
    const index = mailbox?.findIndex((pending) => pending.id === message.id) ?? -1;
    if (mailbox && index >= 0) mailbox.splice(index, 1);
    if (mailbox && !mailbox.length) this.#mailboxes.delete(message.to);
    const wake = this.#wake(message.to, message)
      .finally(() => this.idle(message.to));
    this.#wakes.add(wake);
    void wake.catch(() => undefined);
  }

  #takeMailbox(agentId: string, from?: string): TaskAgentCommunicationMessage | undefined {
    const mailbox = this.#mailboxes.get(agentId);
    if (!mailbox?.length) return undefined;
    const index = from ? mailbox.findIndex((message) => message.from === from) : 0;
    if (index < 0) return undefined;
    const [message] = mailbox.splice(index, 1);
    if (!mailbox.length) this.#mailboxes.delete(agentId);
    return message;
  }

  #takeWaiter(agentId: string, from: string): Waiter | undefined {
    const waiters = this.#waiters.get(agentId);
    if (!waiters) return undefined;
    const index = waiters.findIndex((waiter) => !waiter.from || waiter.from === from);
    if (index < 0) return undefined;
    const [waiter] = waiters.splice(index, 1);
    if (!waiters.length) this.#waiters.delete(agentId);
    return waiter;
  }
}