TJWaterAgent/src/chat/sessionBridge.ts

import { randomUUID } from "node:crypto";

import { logger } from "../logger.js";
import { type OpencodeRuntimeAdapter } from "../runtime/opencode.js";
import {
  buildToolSessionScopeKey,
  ToolSessionContextStore,
} from "../session/toolContextStore.js";
import { toActorKey, toProjectKey } from "../utils/fileStore.js";

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

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

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

export class ChatSessionBridge {
  // runtime session 仅在单次请求生命周期内有效；线程连续性由 clientSessionId 对应的持久状态承担。
  private readonly activeRuntimeSessions = new Map<string, string>();
  private readonly activeSensitiveContexts = new Map<string, ChatRequestContext>();
  private readonly abortControllers = new Map<string, AbortController>();
  private readonly toolContextStore = new ToolSessionContextStore();

  constructor(private readonly runtime: OpencodeRuntimeAdapter) {}

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

    const session = await this.runtime.createSession(requestContext.clientSessionId);
    const binding: SessionBinding = {
      clientSessionId: requestContext.clientSessionId,
      sessionId: session.id,
      startedAt: Date.now(),
    };
    const sessionScopeKey = buildToolSessionScopeKey(
      requestContext.actorKey,
      requestContext.projectKey,
      requestContext.clientSessionId,
    );
    this.activeRuntimeSessions.set(requestContext.clientSessionId, session.id);
    this.activeSensitiveContexts.set(sessionScopeKey, requestContext);
    await this.toolContextStore.write({
      actorKey: requestContext.actorKey,
      allowLearningWrite: true,
      clientSessionId: requestContext.clientSessionId,
      learningMode: "interactive",
      projectId: requestContext.projectId,
      projectKey: requestContext.projectKey,
      sessionId: session.id,
      sessionScopeKey,
      traceId: requestContext.traceId,
    });

    return { binding, requestContext, created: true };
  }

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

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

  getActiveSensitiveContext(sessionScopeKey: string) {
    return this.activeSensitiveContexts.get(sessionScopeKey) ?? null;
  }

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

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

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

    const sessionId = this.activeRuntimeSessions.get(clientSessionId);
    if (!sessionId) {
      return null;
    }

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

  async releaseRuntimeSession(clientSessionId: string, sessionId: string) {
    const activeSessionId = this.activeRuntimeSessions.get(clientSessionId);
    if (activeSessionId === sessionId) {
      this.activeRuntimeSessions.delete(clientSessionId);
    }
    this.activeSensitiveContexts.delete(findScopeKey(this.activeSensitiveContexts, clientSessionId));
    await this.toolContextStore.remove(sessionId).catch((error) => {
      logger.debug({ sessionId, err: error }, "failed to cleanup runtime tool context");
    });
    await this.runtime.abortSession(sessionId).catch((error) => {
      logger.debug({ sessionId, err: error }, "failed to cleanup runtime session");
    });
  }

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

  private async abortActiveRuntime(clientSessionId: string) {
    const activeSessionId = this.activeRuntimeSessions.get(clientSessionId);
    this.activeRuntimeSessions.delete(clientSessionId);
    this.activeSensitiveContexts.delete(findScopeKey(this.activeSensitiveContexts, clientSessionId));

    const controller = this.abortControllers.get(clientSessionId);
    if (controller) {
      this.abortControllers.delete(clientSessionId);
      controller.abort();
    }

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

const findScopeKey = (
  contexts: Map<string, ChatRequestContext>,
  clientSessionId: string,
) => {
  for (const [scopeKey, context] of contexts.entries()) {
    if (context.clientSessionId === clientSessionId) {
      return scopeKey;
    }
  }
  return clientSessionId;
};