pi-sessions 0.6.0 → 0.7.0

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

@@ -1,29 +1,32 @@
 import { randomUUID } from "node:crypto";
 import type { ExtensionContext } from "@earendil-works/pi-coding-agent";
-import type { SessionLineageRelation } from "../../shared/session-index/index.ts";
 import {
-  getIndexStatus,
-  getLineageRelationMap,
-  INDEX_SCHEMA_VERSION,
-  openIndexDatabase,
-} from "../../shared/session-index/index.ts";
-import { spawnSessionMessagingBrokerIfNeeded } from "../broker/spawn.ts";
-import type { SessionMessagePayload, SessionMessagingSessionInfo } from "../shared/protocol.ts";
+  type ActiveSessionBrokerConnection,
+  setActiveSessionBrokerConnection,
+} from "../../shared/session-broker/active.ts";
 import {
   INCOMING_SESSION_MESSAGE_EVENT,
   type SessionMessageSendResult,
   SessionMessagingClient,
-} from "./client.ts";
+} from "../../shared/session-broker/client.ts";
+import type { SessionMessagePayload } from "../../shared/session-broker/protocol.ts";
+import type {
+  SessionLineageRelation,
+  SessionLineageRow,
+} from "../../shared/session-index/index.ts";
+import {
+  getLineageRelationMap,
+  getSessionById,
+  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;
 
-export interface LiveSessionRow extends SessionMessagingSessionInfo {
-  relation?: SessionLineageRelation | undefined;
-}
-
 export interface SendMessageRequest {
   target: string;
   body: string;
@@ -31,6 +34,12 @@ export interface SendMessageRequest {
   sourceToolCallId?: string | undefined;
 }
 
+interface IncomingMessageIndexContext {
+  source: ReceivedMessageEndpoint;
+  target: ReceivedMessageEndpoint;
+  relation?: SessionLineageRelation | undefined;
+}
+
 export class SessionMessagingService {
   private readonly indexPath: string;
   private readonly incomingRuntime: IncomingSessionMessageRuntime;
@@ -51,29 +60,18 @@ export class SessionMessagingService {
   }
 
   async start(ctx: ExtensionContext): Promise<void> {
-    this.refreshCachedRelations(ctx.sessionManager.getSessionId());
-    await this.connection.start(buildSessionInfo(ctx));
+    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 listLiveSessions(ctx: ExtensionContext): Promise<LiveSessionRow[]> {
-    const sessions = await this.connection.listSessions();
-    this.refreshCachedRelations(ctx.sessionManager.getSessionId());
-    const currentSessionId = ctx.sessionManager.getSessionId();
-    return sessions
-      .filter((session) => session.sessionId !== currentSessionId)
-      .map(
-        (session): LiveSessionRow => ({
-          ...session,
-          relation: this.getCachedRelationTo(session.sessionId),
-        }),
-      );
-  }
-
   async sendMessage(request: SendMessageRequest): Promise<SessionMessageSendResult> {
     const relation = this.getCachedRelationTo(request.target, true);
     const sentAt = new Date().toISOString();
@@ -103,8 +101,8 @@ export class SessionMessagingService {
   }
 
   private handleIncoming(requestId: string, message: SessionMessagePayload): void {
-    const relation = this.getCachedRelationTo(message.source.sessionId, true);
-    const result = this.incomingRuntime.deliver(message, relation);
+    const received = this.buildReceivedMessageEntry(message);
+    const result = this.incomingRuntime.deliver(received);
     this.connection.acknowledgeIncoming(requestId, message.messageId, result);
   }
 
@@ -124,12 +122,12 @@ export class SessionMessagingService {
       return undefined;
     }
 
-    const identity = this.connection.currentIdentity;
+    const identity = this.connection.currentSessionId;
     if (!identity) {
       throw new Error("Session messaging is not active.");
     }
 
-    this.refreshCachedRelations(identity.sessionId);
+    this.refreshCachedRelations(identity);
     if (!this.relationBySessionId.has(sessionId)) {
       this.relationBySessionId.set(sessionId, undefined);
     }
@@ -137,9 +135,67 @@ 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 getIncomingMessageIndexContext(
+    sourceSessionId: string,
+    targetSessionId: string,
+  ): IncomingMessageIndexContext {
+    const fallback = {
+      source: buildReceivedMessageEndpoint(sourceSessionId),
+      target: buildReceivedMessageEndpoint(targetSessionId),
+    };
+    return (
+      withSessionIndex(this.indexPath, { mode: "read", required: false }, ({ db }) => {
+        const source = getSessionById(db, sourceSessionId);
+        const target = getSessionById(db, targetSessionId);
+        const currentSessionId = this.connection.currentSessionId;
+        const relationBySessionId = currentSessionId
+          ? new Map(getLineageRelationMap(db, currentSessionId))
+          : new Map<string, SessionLineageRelation>();
+        const relation = relationBySessionId.get(sourceSessionId);
+
+        if (currentSessionId) {
+          this.replaceCachedRelations(relationBySessionId);
+          if (!this.relationBySessionId.has(sourceSessionId)) {
+            this.relationBySessionId.set(sourceSessionId, undefined);
+          }
+        }
+
+        return {
+          source: buildReceivedMessageEndpoint(sourceSessionId, source),
+          target: buildReceivedMessageEndpoint(targetSessionId, target),
+          ...(relation === undefined ? {} : { relation }),
+        };
+      }) ?? fallback
+    );
+  }
+
   private refreshCachedRelations(currentSessionId: string): void {
+    this.replaceCachedRelations(getRelationMapForSession(this.indexPath, currentSessionId));
+  }
+
+  private replaceCachedRelations(
+    nextRelations: Map<string, SessionLineageRelation | undefined>,
+  ): void {
     const previousRelations = this.relationBySessionId;
-    const nextRelations = getRelationMapForSession(this.indexPath, currentSessionId);
 
     for (const [sessionId, relation] of previousRelations) {
       if (relation === undefined && !nextRelations.has(sessionId)) {
@@ -151,9 +207,9 @@ export class SessionMessagingService {
   }
 }
 
-class BrokerConnection {
+class BrokerConnection implements ActiveSessionBrokerConnection {
   private client: SessionMessagingClient | undefined;
-  private identity: SessionMessagingSessionInfo | undefined;
+  private sessionId: string | undefined;
   private active = false;
   private reconnectTimer: NodeJS.Timeout | undefined;
   private reconnectDelayIndex = 0;
@@ -163,28 +219,28 @@ class BrokerConnection {
     private readonly onIncoming: (requestId: string, message: SessionMessagePayload) => void,
   ) {}
 
-  get currentIdentity(): SessionMessagingSessionInfo | undefined {
-    return this.identity;
+  get currentSessionId(): string | undefined {
+    return this.sessionId;
   }
 
-  async start(identity: SessionMessagingSessionInfo): Promise<void> {
+  async start(sessionId: string): Promise<void> {
     this.active = true;
-    this.identity = identity;
+    this.sessionId = sessionId;
     await this.ensureConnectedNow();
   }
 
   stop(): void {
     this.active = false;
-    this.identity = undefined;
+    this.sessionId = undefined;
     this.clearReconnectTimer();
     this.client?.disconnect();
     this.client = undefined;
     this.connectPromise = undefined;
   }
 
-  async listSessions(): Promise<SessionMessagingSessionInfo[]> {
+  async listSessionIds(): Promise<string[]> {
     await this.ensureConnectedNow();
-    return this.requireClient().listSessions();
+    return this.requireClient().listSessionIds();
   }
 
   async sendMessage(options: {
@@ -213,7 +269,7 @@ class BrokerConnection {
     if (this.client?.isConnected) {
       return;
     }
-    if (!this.active || !this.identity) {
+    if (!this.active || !this.sessionId) {
       throw new Error("Session messaging is not active.");
     }
     this.clearReconnectTimer();
@@ -233,8 +289,8 @@ class BrokerConnection {
   }
 
   private async connect(): Promise<void> {
-    const identity = this.identity;
-    if (!identity) {
+    const sessionId = this.sessionId;
+    if (!sessionId) {
       throw new Error("Session messaging is not active.");
     }
 
@@ -251,7 +307,7 @@ class BrokerConnection {
       }
     });
 
-    await client.connect(identity);
+    await client.connect(sessionId);
     this.client = client;
     this.reconnectDelayIndex = 0;
   }
@@ -288,12 +344,15 @@ class BrokerConnection {
   }
 }
 
-function buildSessionInfo(ctx: ExtensionContext): SessionMessagingSessionInfo {
-  const sessionName = ctx.sessionManager.getSessionName()?.trim();
+function buildReceivedMessageEndpoint(
+  sessionId: string,
+  session?: SessionLineageRow | undefined,
+): ReceivedMessageEndpoint {
+  const sessionName = session?.sessionName.trim();
   return {
-    sessionId: ctx.sessionManager.getSessionId(),
+    sessionId,
     ...(sessionName ? { sessionName } : {}),
-    cwd: ctx.cwd,
+    ...(session?.cwd ? { cwd: session.cwd } : {}),
   };
 }
 
@@ -301,15 +360,11 @@ function getRelationMapForSession(
   indexPath: string,
   sessionId: string,
 ): Map<string, SessionLineageRelation | undefined> {
-  const status = getIndexStatus(indexPath);
-  if (!status.exists || status.schemaVersion !== INDEX_SCHEMA_VERSION) {
-    return new Map();
-  }
-
-  const db = openIndexDatabase(status.dbPath, { create: false });
-  try {
-    return new Map(getLineageRelationMap(db, sessionId));
-  } finally {
-    db.close();
-  }
+  return (
+    withSessionIndex(
+      indexPath,
+      { mode: "read", required: false },
+      ({ db }) => new Map(getLineageRelationMap(db, sessionId)),
+    ) ?? new Map()
+  );
 }

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

@@ -1,52 +1,21 @@
 import type { ExtensionAPI } from "@earendil-works/pi-coding-agent";
 import { Text } from "@earendil-works/pi-tui";
-import { Type } from "typebox";
 import { SEND_MESSAGE_PARAMS, type SendMessageParams } from "./message-contracts.ts";
 import { formatSendMessageCall } from "./message-view.ts";
-import type { LiveSessionRow, SessionMessagingService } from "./service.ts";
-
-interface ListLiveDetails {
-  error?: boolean | undefined;
-  sessions: LiveSessionRow[];
-}
+import type { SessionMessagingService } from "./service.ts";
 
 export function registerSessionMessagingTools(
   pi: ExtensionAPI,
   service: SessionMessagingService,
 ): void {
-  pi.registerTool({
-    name: "session_list_live",
-    label: "List Live Pi Sessions",
-    description: "List other Pi sessions that are currently active.",
-    promptSnippet: "List other Pi sessions that are currently active.",
-    parameters: Type.Object({}),
-    renderResult(result, _options, theme, context) {
-      if (context.isError) {
-        return new Text(`\n${theme.fg("error", getFirstText(result))}`, 0, 0);
-      }
-
-      const details = result.details as ListLiveDetails | undefined;
-      return new Text(`\n${formatLiveSessionsForDisplay(details?.sessions ?? [], theme)}`, 0, 0);
-    },
-    async execute(_toolCallId, _params, _signal, _onUpdate, ctx) {
-      try {
-        const sessions = await service.listLiveSessions(ctx);
-        return {
-          content: [{ type: "text" as const, text: formatLiveSessions(sessions) }],
-          details: { sessions } satisfies ListLiveDetails,
-        };
-      } catch (error) {
-        throw new Error(formatError(error));
-      }
-    },
-  });
-
   pi.registerTool({
     name: "session_send_message",
     label: "Send Message to Session",
     description: "Send a message to another live pi session",
     promptSnippet: "Send a message to another live pi session",
-    promptGuidelines: ["Use session_list_live to discover targetable sessions."],
+    promptGuidelines: [
+      "Use session_search with live: true to discover targetable sessions. Usually it is best to exclude all other filters and simply get a list of all live sessions",
+    ],
     parameters: SEND_MESSAGE_PARAMS,
     renderCall(args, theme, context) {
       const component = (context.lastComponent as Text | undefined) ?? new Text("", 0, 0);
@@ -114,27 +83,6 @@ export function registerSessionMessagingTools(
   });
 }
 
-function formatLiveSessions(sessions: LiveSessionRow[]): string {
-  return JSON.stringify({ sessions }, null, 2);
-}
-
-function formatLiveSessionsForDisplay(
-  sessions: LiveSessionRow[],
-  theme: { fg(token: string, text: string): string },
-): string {
-  if (sessions.length === 0) {
-    return "No other live sessions found.";
-  }
-
-  return sessions
-    .map((session) => {
-      const title = session.sessionName?.trim() || "[untitled]";
-      const relation = session.relation ? ` (${session.relation})` : "";
-      return `${session.sessionId}${relation} - ${title}\n${theme.fg("dim", session.cwd)}`;
-    })
-    .join("\n\n");
-}
-
 function getFirstText(result: { content: Array<{ type: string; text?: string }> }): string {
   return result.content.find((item) => item.type === "text")?.text ?? "";
 }