pi-sessions 0.7.2 → 0.9.0

package/extensions/session-messaging/pi/service.ts

@@ -1,15 +1,20 @@
 import { randomUUID } from "node:crypto";
 import type { ExtensionContext } from "@earendil-works/pi-coding-agent";
+import { formatError } from "../../shared/errors.ts";
 import {
-  type ActiveSessionBrokerConnection,
-  setActiveSessionBrokerConnection,
-} from "../../shared/session-broker/active.ts";
-import {
-  INCOMING_SESSION_MESSAGE_EVENT,
-  type SessionMessageSendResult,
+  INCOMING_SESSION_ENVELOPE_EVENT,
+  type SessionEnvelopeSendResult,
   SessionMessagingClient,
 } from "../../shared/session-broker/client.ts";
-import type { SessionMessagePayload } from "../../shared/session-broker/protocol.ts";
+import type {
+  OutboundSessionEnvelope,
+  SessionCancelEnvelope,
+  SessionEnvelope,
+  SessionEnvelopeSendFailureReason,
+  SessionMessageEnvelope,
+  SessionSubagentReportEnvelope,
+  TaskReport,
+} from "../../shared/session-broker/protocol.ts";
 import type {
   SessionLineageRelation,
   SessionLineageRow,
@@ -20,12 +25,10 @@ import {
   withSessionIndex,
 } from "../../shared/session-index/index.ts";
 import { spawnSessionMessagingBrokerIfNeeded } from "../broker/spawn.ts";
-import type { IncomingSessionMessageRuntime } from "./incoming-runtime.ts";
 import type { ReceivedMessageEndpoint, ReceivedMessageEntry } from "./message-contracts.ts";
 
-export const MESSAGE_SENT_CUSTOM_TYPE = "pi-sessions.message_sent";
-
 const RECONNECT_DELAYS_MS = [250, 500, 1_000, 2_000, 5_000, 10_000, 30_000] as const;
+const SESSION_WAIT_POLL_MS = 100;
 
 export interface SendMessageRequest {
   target: string;
@@ -34,6 +37,49 @@ export interface SendMessageRequest {
   sourceToolCallId?: string | undefined;
 }
 
+export type SendMessageResult =
+  | {
+      messageId: string;
+      delivered: false;
+      reason?: SessionEnvelopeSendFailureReason | undefined;
+      error?: string | undefined;
+    }
+  | {
+      messageId: string;
+      delivered: true;
+      target: ReceivedMessageEndpoint;
+      relation?: SessionLineageRelation | undefined;
+    };
+
+export type CancelSessionResult =
+  | {
+      kind: "transport";
+      cancelId: string;
+      delivered: false;
+      reason?: SessionEnvelopeSendFailureReason | undefined;
+      error?: string | undefined;
+    }
+  | {
+      kind: "transport";
+      cancelId: string;
+      delivered: true;
+      target: ReceivedMessageEndpoint;
+      relation?: SessionLineageRelation | undefined;
+    };
+
+export type IncomingMessageHandler = (envelope: SessionMessageEnvelope) => Promise<void> | void;
+export type IncomingCancelHandler = (envelope: SessionCancelEnvelope) => Promise<void> | void;
+export type IncomingSubagentReportHandler = (
+  envelope: SessionSubagentReportEnvelope,
+) => Promise<void> | void;
+
+export type SendSubagentReportResult = SessionEnvelopeSendResult;
+
+export interface SendSubagentReportRequest extends TaskReport {
+  target: string;
+  reportId: string;
+}
+
 interface IncomingMessageIndexContext {
   source: ReceivedMessageEndpoint;
   target: ReceivedMessageEndpoint;
@@ -42,68 +88,171 @@ interface IncomingMessageIndexContext {
 
 export class SessionMessagingService {
   private readonly indexPath: string;
-  private readonly incomingRuntime: IncomingSessionMessageRuntime;
-  private readonly appendEntry: (customType: string, data: unknown) => void;
   private relationBySessionId = new Map<string, SessionLineageRelation | undefined>();
-  private readonly connection = new BrokerConnection((requestId, message) =>
-    this.handleIncoming(requestId, message),
-  );
+  private incomingMessageHandler: IncomingMessageHandler | undefined;
+  private incomingCancelHandler: IncomingCancelHandler | undefined;
+  private incomingSubagentReportHandler: IncomingSubagentReportHandler | undefined;
+  private readonly connection = new BrokerConnection((envelope) => this.acceptIncoming(envelope));
 
-  constructor(
-    indexPath: string,
-    incomingRuntime: IncomingSessionMessageRuntime,
-    appendEntry: (customType: string, data: unknown) => void,
-  ) {
+  constructor(indexPath: string) {
     this.indexPath = indexPath;
-    this.incomingRuntime = incomingRuntime;
-    this.appendEntry = appendEntry;
   }
 
   async start(ctx: ExtensionContext): Promise<void> {
     const sessionId = ctx.sessionManager.getSessionId();
     this.refreshCachedRelations(sessionId);
-    setActiveSessionBrokerConnection(this.connection);
     await this.connection.start(sessionId);
   }
 
   stop(): void {
-    setActiveSessionBrokerConnection(undefined);
     this.relationBySessionId.clear();
     this.connection.stop();
   }
 
-  async sendMessage(request: SendMessageRequest): Promise<SessionMessageSendResult> {
+  onIncomingMessage(handler: IncomingMessageHandler): void {
+    if (this.incomingMessageHandler) {
+      throw new Error("An incoming message handler is already registered.");
+    }
+    this.incomingMessageHandler = handler;
+  }
+
+  onIncomingCancel(handler: IncomingCancelHandler): void {
+    if (this.incomingCancelHandler) {
+      throw new Error("An incoming cancel handler is already registered.");
+    }
+    this.incomingCancelHandler = handler;
+  }
+
+  onIncomingSubagentReport(handler: IncomingSubagentReportHandler): void {
+    if (this.incomingSubagentReportHandler) {
+      throw new Error("An incoming subagent report handler is already registered.");
+    }
+    this.incomingSubagentReportHandler = handler;
+  }
+
+  async listSessions(): Promise<string[]> {
+    return this.connection.listSessionIds();
+  }
+
+  async waitForSession(sessionId: string, timeoutMs: number): Promise<boolean> {
+    if (!Number.isFinite(timeoutMs) || timeoutMs < 0) {
+      throw new Error("Session wait timeout must be a non-negative finite number.");
+    }
+
+    const deadline = Date.now() + timeoutMs;
+    while (true) {
+      if ((await this.listSessions()).includes(sessionId)) {
+        return true;
+      }
+
+      const remainingMs = deadline - Date.now();
+      if (remainingMs <= 0) {
+        return false;
+      }
+      await delay(Math.min(SESSION_WAIT_POLL_MS, remainingMs));
+    }
+  }
+
+  async sendMessage(request: SendMessageRequest): Promise<SendMessageResult> {
     const relation = this.getCachedRelationTo(request.target, true);
-    const sentAt = new Date().toISOString();
     const messageId = randomUUID();
-    const result = await this.connection.sendMessage({
+    const result = await this.connection.sendEnvelope({
+      target: request.target,
+      envelope: {
+        kind: "message",
+        messageId,
+        body: request.body,
+        ...(request.requestResponse === undefined
+          ? {}
+          : { requestResponse: request.requestResponse }),
+        ...(request.sourceToolCallId === undefined
+          ? {}
+          : { sourceToolCallId: request.sourceToolCallId }),
+        sentAt: new Date().toISOString(),
+      },
+    });
+
+    if (!result.delivered) {
+      return {
+        messageId,
+        delivered: false,
+        ...(result.reason === undefined ? {} : { reason: result.reason }),
+        ...(result.error === undefined ? {} : { error: result.error }),
+      };
+    }
+
+    return {
       messageId,
+      delivered: true,
+      target: this.getTargetEndpoint(request.target),
+      ...(relation === undefined ? {} : { relation }),
+    };
+  }
+
+  async sendSubagentReport(request: SendSubagentReportRequest): Promise<SendSubagentReportResult> {
+    return this.connection.sendEnvelope({
       target: request.target,
-      body: request.body,
-      sentAt,
-      requestResponse: request.requestResponse,
-      sourceToolCallId: request.sourceToolCallId,
+      envelope: {
+        kind: "subagent_report",
+        reportId: request.reportId,
+        status: request.status,
+        summary: request.summary,
+        ...(request.details ? { details: request.details } : {}),
+        ...(request.references ? { references: request.references } : {}),
+        ...(request.nextSteps ? { nextSteps: request.nextSteps } : {}),
+        sentAt: new Date().toISOString(),
+      },
     });
+  }
 
-    if (result.delivered) {
-      this.appendEntry(MESSAGE_SENT_CUSTOM_TYPE, {
-        messageId: result.messageId,
-        target: request.target,
-        body: request.body,
-        requestResponse: request.requestResponse,
-        sourceToolCallId: request.sourceToolCallId,
-        sentAt,
-        ...(relation === undefined ? {} : { relation }),
-      });
+  async cancelSession(sessionId: string): Promise<CancelSessionResult> {
+    const relation = this.getCachedRelationTo(sessionId, true);
+    const cancelId = randomUUID();
+    const result = await this.connection.sendEnvelope({
+      target: sessionId,
+      envelope: {
+        kind: "cancel",
+        cancelId,
+        sentAt: new Date().toISOString(),
+      },
+    });
+
+    if (!result.delivered) {
+      return {
+        kind: "transport",
+        cancelId,
+        delivered: false,
+        ...(result.reason === undefined ? {} : { reason: result.reason }),
+        ...(result.error === undefined ? {} : { error: result.error }),
+      };
     }
 
-    return result;
+    return {
+      kind: "transport",
+      cancelId,
+      delivered: true,
+      target: this.getTargetEndpoint(sessionId),
+      ...(relation === undefined ? {} : { relation }),
+    };
   }
 
-  private handleIncoming(requestId: string, message: SessionMessagePayload): void {
-    const received = this.buildReceivedMessageEntry(message);
-    const result = this.incomingRuntime.deliver(received);
-    this.connection.acknowledgeIncoming(requestId, message.messageId, result);
+  buildReceivedMessage(message: SessionMessageEnvelope): ReceivedMessageEntry {
+    const context = this.getIncomingMessageIndexContext(message.source, message.target);
+    return {
+      messageId: message.messageId,
+      source: context.source,
+      target: context.target,
+      body: message.body,
+      sentAt: message.sentAt,
+      receivedAt: new Date().toISOString(),
+      ...(message.requestResponse === undefined
+        ? {}
+        : { requestResponse: message.requestResponse }),
+      ...(message.sourceToolCallId === undefined
+        ? {}
+        : { sourceToolCallId: message.sourceToolCallId }),
+      ...(context.relation === undefined ? {} : { relation: context.relation }),
+    };
   }
 
   getCachedRelationTo(
@@ -135,23 +284,46 @@ export class SessionMessagingService {
     return this.relationBySessionId.get(sessionId);
   }
 
-  private buildReceivedMessageEntry(message: SessionMessagePayload): ReceivedMessageEntry {
-    const context = this.getIncomingMessageIndexContext(message.source, message.target);
-    return {
-      messageId: message.messageId,
-      source: context.source,
-      target: context.target,
-      body: message.body,
-      sentAt: message.sentAt,
-      receivedAt: new Date().toISOString(),
-      ...(message.requestResponse === undefined
-        ? {}
-        : { requestResponse: message.requestResponse }),
-      ...(message.sourceToolCallId === undefined
-        ? {}
-        : { sourceToolCallId: message.sourceToolCallId }),
-      ...(context.relation === undefined ? {} : { relation: context.relation }),
-    };
+  private async acceptIncoming(envelope: SessionEnvelope): Promise<IncomingAcceptance> {
+    try {
+      switch (envelope.kind) {
+        case "message": {
+          if (!this.incomingMessageHandler) {
+            throw new Error("Target session has no incoming message handler.");
+          }
+          await this.incomingMessageHandler(envelope);
+          break;
+        }
+        case "cancel": {
+          if (!this.incomingCancelHandler) {
+            throw new Error("Target session has no incoming cancel handler.");
+          }
+          await this.incomingCancelHandler(envelope);
+          break;
+        }
+        case "subagent_report": {
+          if (!this.incomingSubagentReportHandler) {
+            throw new Error("Target session has no incoming subagent report handler.");
+          }
+          await this.incomingSubagentReportHandler(envelope);
+          break;
+        }
+      }
+      return { delivered: true };
+    } catch (error) {
+      return {
+        delivered: false,
+        error: formatError(error),
+      };
+    }
+  }
+
+  private getTargetEndpoint(targetSessionId: string): ReceivedMessageEndpoint {
+    return (
+      withSessionIndex(this.indexPath, { mode: "read", required: false }, ({ db }) =>
+        buildReceivedMessageEndpoint(targetSessionId, getSessionById(db, targetSessionId)),
+      ) ?? buildReceivedMessageEndpoint(targetSessionId)
+    );
   }
 
   private getIncomingMessageIndexContext(
@@ -207,7 +379,9 @@ export class SessionMessagingService {
   }
 }
 
-class BrokerConnection implements ActiveSessionBrokerConnection {
+type IncomingAcceptance = { delivered: true } | { delivered: false; error: string };
+
+class BrokerConnection {
   private client: SessionMessagingClient | undefined;
   private sessionId: string | undefined;
   private active = false;
@@ -216,7 +390,7 @@ class BrokerConnection implements ActiveSessionBrokerConnection {
   private connectPromise: Promise<void> | undefined;
 
   constructor(
-    private readonly onIncoming: (requestId: string, message: SessionMessagePayload) => void,
+    private readonly onIncoming: (envelope: SessionEnvelope) => Promise<IncomingAcceptance>,
   ) {}
 
   get currentSessionId(): string | undefined {
@@ -243,25 +417,26 @@ class BrokerConnection implements ActiveSessionBrokerConnection {
     return this.requireClient().listSessionIds();
   }
 
-  async sendMessage(options: {
-    messageId: string;
+  async sendEnvelope(options: {
     target: string;
-    body: string;
-    sentAt: string;
-    requestResponse?: boolean | undefined;
-    sourceToolCallId?: string | undefined;
-  }): Promise<SessionMessageSendResult> {
+    envelope: OutboundSessionEnvelope;
+  }): Promise<SessionEnvelopeSendResult> {
     await this.ensureConnectedNow();
-    return this.requireClient().sendMessage(options);
+    return this.requireClient().sendEnvelope(options);
   }
 
-  acknowledgeIncoming(
+  private async acceptIncoming(
+    client: SessionMessagingClient,
     requestId: string,
-    messageId: string,
-    result: { delivered: true } | { delivered: false; error: string },
-  ): void {
+    envelope: SessionEnvelope,
+  ): Promise<void> {
+    const result = await this.onIncoming(envelope);
+    if (this.client !== client || !client.isConnected) {
+      return;
+    }
+
     try {
-      this.requireClient().acknowledgeIncoming(requestId, messageId, result);
+      client.acknowledgeIncoming(requestId, result);
     } catch {}
   }
 
@@ -299,7 +474,9 @@ class BrokerConnection implements ActiveSessionBrokerConnection {
 
     await spawnSessionMessagingBrokerIfNeeded();
     const client = new SessionMessagingClient();
-    client.on(INCOMING_SESSION_MESSAGE_EVENT, this.onIncoming);
+    client.on(INCOMING_SESSION_ENVELOPE_EVENT, (requestId: string, envelope: SessionEnvelope) => {
+      void this.acceptIncoming(client, requestId, envelope);
+    });
     client.on("disconnect", () => {
       if (this.client === client) {
         this.client = undefined;
@@ -317,7 +494,7 @@ class BrokerConnection implements ActiveSessionBrokerConnection {
       return;
     }
 
-    const delay = RECONNECT_DELAYS_MS[this.reconnectDelayIndex] ?? 30_000;
+    const delayMs = RECONNECT_DELAYS_MS[this.reconnectDelayIndex] ?? 30_000;
     this.reconnectDelayIndex = Math.min(
       this.reconnectDelayIndex + 1,
       RECONNECT_DELAYS_MS.length - 1,
@@ -325,7 +502,7 @@ class BrokerConnection implements ActiveSessionBrokerConnection {
     this.reconnectTimer = setTimeout(() => {
       this.reconnectTimer = undefined;
       this.ensureConnectedNow().catch(() => {});
-    }, delay);
+    }, delayMs);
   }
 
   private clearReconnectTimer(): void {
@@ -368,3 +545,7 @@ function getRelationMapForSession(
     ) ?? new Map()
   );
 }
+
+function delay(ms: number): Promise<void> {
+  return new Promise((resolve) => setTimeout(resolve, ms));
+}