TJWaterAgent/src/chat/sessionBridge.ts

import { randomUUID } from "node:crypto";

import { logger } from "../logger.js";
import { type OpencodeRuntimeAdapter } from "../runtime/opencode.js";
import {
  removeRuntimeSessionContext,
  setRuntimeSessionContext,
} from "../runtime/sessionContext.js";
import { toActorKey, toProjectKey } from "../utils/fileStore.js";

export type SessionBinding = {
  clientSessionId: string;
  sessionId: string;
  startedAt: number;
  workspaceDirectory?: string;
};

export type SessionContext = {
  clientSessionId: string;
  accessToken?: string;
  network?: string;
  projectId?: string;
  tokenExpiresAt?: string;
  userId?: string;
};

export type ChatRequestContext = SessionContext & {
  actorKey: string;
  projectKey: string;
  traceId: string;
};

export class ChatSessionBridge {
  private readonly abortControllers = new Map<string, AbortController>();

  constructor(private readonly runtime: OpencodeRuntimeAdapter) {}

  async resolve(context: {
    sessionId?: string;
    accessToken?: string;
    network?: string;
    projectId?: string;
    traceId?: string;
    tokenExpiresAt?: string;
    userId?: string;
  }): Promise<{
    binding: SessionBinding;
    requestContext: ChatRequestContext;
    created: boolean;
  }> {
    let requestContext = this.buildRequestContext(context);
    const existingSessionId = context.sessionId?.trim();
    await this.abortActiveRuntime(requestContext.clientSessionId, existingSessionId);

    let sessionId = existingSessionId;
    let runtimeSession;
    let created = false;
    if (!sessionId) {
      runtimeSession = await this.runtime.createSession(undefined, {
        conversationWorkspace: true,
      });
      sessionId = runtimeSession.id;
      requestContext = {
        ...requestContext,
        clientSessionId: sessionId,
      };
      created = true;
    } else {
      runtimeSession = await this.runtime.getSession(sessionId);
    }
    const binding: SessionBinding = {
      clientSessionId: requestContext.clientSessionId,
      sessionId,
      startedAt: Date.now(),
      workspaceDirectory: runtimeSession.directory,
    };
    setRuntimeSessionContext({
      accessToken: requestContext.accessToken,
      actorKey: requestContext.actorKey,
      allowLearningWrite: true,
      clientSessionId: requestContext.clientSessionId,
      learningMode: "interactive",
      network: requestContext.network,
      projectId: requestContext.projectId,
      projectKey: requestContext.projectKey,
      sessionId,
      tokenExpiresAt: requestContext.tokenExpiresAt,
      traceId: requestContext.traceId,
      workspaceDirectory: runtimeSession.directory,
    });

    return { binding, requestContext, created };
  }

  count(): number {
    return this.abortControllers.size;
  }

  createClientSessionId() {
    return `agent-${randomUUID().slice(0, 12)}`;
  }

  registerAbortController(clientSessionId: string, controller: AbortController) {
    this.abortControllers.set(clientSessionId, controller);
  }

  finalizeRequest(clientSessionId: string) {
    this.abortControllers.delete(clientSessionId);
  }

  async abort(context: {
    clientSessionId?: string;
    sessionId?: string;
  }): Promise<SessionBinding | null> {
    const clientSessionId = context.clientSessionId?.trim();
    const sessionId = context.sessionId?.trim();
    if (!clientSessionId || !sessionId) {
      return null;
    }

    await this.abortActiveRuntime(clientSessionId, sessionId);
    return {
      clientSessionId,
      sessionId,
      startedAt: Date.now(),
    };
  }

  async deleteSession(context: {
    clientSessionId: string;
    sessionId: string;
  }) {
    const clientSessionId = context.clientSessionId.trim();
    const sessionId = context.sessionId.trim();
    const controller = this.abortControllers.get(clientSessionId);
    if (controller) {
      this.abortControllers.delete(clientSessionId);
      controller.abort();
    }
    await this.runtime.abortSession(sessionId).catch((error) => {
      logger.warn(
        { clientSessionId, sessionId, err: error },
        "failed to abort runtime session",
      );
    });
    await this.runtime.waitForSessionIdle(sessionId).catch((error) => {
      logger.warn(
        { clientSessionId, sessionId, err: error },
        "failed while waiting for runtime session to become idle",
      );
    });
    removeRuntimeSessionContext(sessionId);
  }

  private buildRequestContext(context: {
    sessionId?: string;
    accessToken?: string;
    network?: string;
    projectId?: string;
    traceId?: string;
    tokenExpiresAt?: string;
    userId?: string;
  }): ChatRequestContext {
    const sessionId = context.sessionId?.trim();
    return {
      clientSessionId: sessionId || this.createClientSessionId(),
      accessToken: context.accessToken,
      actorKey: toActorKey(context.userId),
      network: context.network,
      projectId: context.projectId,
      projectKey: toProjectKey(context.projectId),
      tokenExpiresAt: context.tokenExpiresAt,
      traceId: context.traceId?.trim() || `trace-${randomUUID().slice(0, 12)}`,
      userId: context.userId?.trim(),
    };
  }

  private async abortActiveRuntime(clientSessionId: string, sessionId?: string) {
    const controller = this.abortControllers.get(clientSessionId);
    if (controller) {
      this.abortControllers.delete(clientSessionId);
      controller.abort();
    }

    if (!sessionId) {
      return;
    }
    await this.runtime.abortSession(sessionId).catch((error) => {
      logger.warn(
        { clientSessionId, sessionId, err: error },
        "failed to abort active runtime session",
      );
    });
    await this.runtime.waitForSessionIdle(sessionId).catch((error) => {
      logger.warn(
        { clientSessionId, sessionId, err: error },
        "failed while waiting for active runtime session to become idle",
      );
    });
  }
}