ai-terminal-assistant/packages/server/src/permission/handler.ts

/**
 * Server 端权限处理器
 * 通过 WebSocket 发送权限请求并等待客户端响应
 */

import { randomUUID } from 'crypto';
import { broadcastToSession } from '../ws.js';
import type {
  PermissionRequestPayload,
  PermissionDisplayContext,
  ServerMessage,
  DiffInfo,
  DiffHunkInfo,
} from '../types.js';
import { inferPermissionType } from '@ai-assistant/core';
import type { PermissionDecision, PermissionContext, PermissionType } from '@ai-assistant/core';

// 等待中的权限请求
interface PendingRequest {
  resolve: (decision: PermissionDecision) => void;
  reject: (error: Error) => void;
  timeout: ReturnType<typeof setTimeout>;
}

const pendingRequests = new Map<string, PendingRequest>();

// 会话级别的 auto-approve 配置
// key: sessionId, value: auto-approve 配置
const sessionAutoApprove = new Map<string, { file?: { write?: 'allow'; edit?: 'allow' } }>();

// 默认超时时间（60秒）
const PERMISSION_TIMEOUT = 60000;

/**
 * 设置会话的 auto-approve 配置
 */
export function setSessionAutoApprove(
  sessionId: string,
  config: { file?: { write?: 'allow'; edit?: 'allow' } } | null
): void {
  if (config) {
    sessionAutoApprove.set(sessionId, config);
  } else {
    sessionAutoApprove.delete(sessionId);
  }
}

/**
 * 获取会话的 auto-approve 配置
 */
export function getSessionAutoApprove(sessionId: string): { file?: { write?: 'allow'; edit?: 'allow' } } | null {
  return sessionAutoApprove.get(sessionId) ?? null;
}

/**
 * 检查操作是否被 auto-approve
 */
function isAutoApproved(sessionId: string, ctx: PermissionContext): boolean {
  const config = sessionAutoApprove.get(sessionId);
  if (!config) return false;

  // 优先使用结构化字段判断
  if (ctx.permissionType === 'file') {
    if (ctx.fileOperation === 'write' && config.file?.write === 'allow') {
      return true;
    }
    if (ctx.fileOperation === 'edit' && config.file?.edit === 'allow') {
      return true;
    }
  }

  // 向后兼容：解析 command 字符串
  const command = ctx.command.toLowerCase();
  if (command.startsWith('write ') && config.file?.write === 'allow') {
    return true;
  }
  if (command.startsWith('edit ') && config.file?.edit === 'allow') {
    return true;
  }

  return false;
}

/**
 * 生成简单的行级别 diff 信息
 * 用于在权限确认对话框中显示文件变更预览
 */
function generateDiffInfo(oldContent: string | undefined, newContent: string | undefined): DiffInfo | undefined {
  // 如果没有内容，无法生成 diff
  if (newContent === undefined) {
    return undefined;
  }

  const isNew = !oldContent;
  const oldLines = oldContent ? oldContent.split('\n') : [];
  const newLines = newContent.split('\n');

  // 简单的行级别 diff
  const lines: DiffHunkInfo['lines'] = [];
  let additions = 0;
  let deletions = 0;

  if (isNew) {
    // 新文件：所有行都是添加
    newLines.forEach((line, index) => {
      lines.push({
        type: 'add',
        lineNumber: index + 1,
        content: line,
      });
      additions++;
    });
  } else {
    // 编辑：显示删除和添加
    // 对于 edit_file（search-replace），oldContent 是被替换的部分，newContent 是替换后的部分
    oldLines.forEach((line, index) => {
      lines.push({
        type: 'remove',
        lineNumber: index + 1,
        content: line,
      });
      deletions++;
    });
    newLines.forEach((line, index) => {
      lines.push({
        type: 'add',
        lineNumber: index + 1,
        content: line,
      });
      additions++;
    });
  }

  // 构建单个 hunk
  const hunk: DiffHunkInfo = {
    oldStart: 1,
    oldCount: oldLines.length,
    newStart: 1,
    newCount: newLines.length,
    lines,
  };

  return {
    isNew,
    additions,
    deletions,
    hunks: [hunk],
  };
}

/**
 * 构建权限请求显示上下文
 * 将 Core 的完整 PermissionContext 转换为用于前端显示的精简格式
 */
function buildDisplayContext(ctx: PermissionContext): PermissionDisplayContext {
  const permType = inferPermissionType(ctx);

  switch (permType) {
    case 'file':
      return {
        command: ctx.command,
        operation: ctx.fileOperation || ctx.command.split(' ')[0],
        path: ctx.filePath || ctx.command.split(' ').slice(1).join(' '),
        patterns: ctx.patterns,
        externalPaths: ctx.externalPaths,
      };

    case 'git':
      return {
        command: ctx.command,
        gitOperation: ctx.gitOperation || ctx.command.replace(/^git\s+/, '').split(' ')[0],
      };

    case 'web':
      return {
        command: ctx.command,
        query: ctx.webQuery || ctx.command,
      };

    default:
      return {
        command: ctx.command,
        patterns: ctx.patterns,
        externalPaths: ctx.externalPaths,
      };
  }
}

/**
 * 创建 Server 端权限回调函数
 * 用于在 Agent 创建时设置
 */
export function createServerPermissionCallback(sessionId: string) {
  return async (ctx: unknown): Promise<PermissionDecision> => {
    const permCtx = ctx as PermissionContext;

    // 检查 auto-approve 配置
    if (isAutoApproved(sessionId, permCtx)) {
      console.log(`[Permission] Auto-approved: ${permCtx.command}`);
      return { allow: true, remember: false };
    }

    const requestId = randomUUID();
    const permissionType = inferPermissionType(permCtx);
    const context = buildDisplayContext(permCtx);

    // 为文件操作生成 diff
    let diff: DiffInfo | undefined;
    if (permissionType === 'file' && (permCtx.fileOperation === 'write' || permCtx.fileOperation === 'edit')) {
      diff = generateDiffInfo(permCtx.oldContent, permCtx.newContent);
    }

    // 构建请求 payload
    const payload: PermissionRequestPayload = {
      requestId,
      permissionType,
      context,
      diff,
    };

    // 发送权限请求到客户端
    const message: ServerMessage = {
      type: 'permission_request',
      sessionId,
      payload,
    };

    broadcastToSession(sessionId, message);

    // 等待响应（带超时）
    return new Promise((resolve, reject) => {
      const timeout = setTimeout(() => {
        pendingRequests.delete(requestId);
        // 超时默认拒绝
        resolve({ allow: false, remember: false });
      }, PERMISSION_TIMEOUT);

      pendingRequests.set(requestId, { resolve, reject, timeout });
    });
  };
}

/**
 * 处理权限响应
 * 由 WebSocket 消息处理器调用
 */
export function handlePermissionResponse(
  requestId: string,
  allow: boolean,
  remember?: boolean
): boolean {
  const pending = pendingRequests.get(requestId);

  if (!pending) {
    console.warn(`[Permission] No pending request found for: ${requestId}`);
    return false;
  }

  clearTimeout(pending.timeout);
  pendingRequests.delete(requestId);
  pending.resolve({ allow, remember });

  return true;
}

/**
 * 取消会话的所有待处理权限请求
 * 用于会话断开时清理
 */
export function cancelPendingRequests(sessionId: string): void {
  // 由于我们没有存储 sessionId -> requestId 的映射，
  // 这个函数目前只是一个占位符
  // 在实际使用中，如果需要按会话取消请求，需要维护额外的映射
  console.log(`[Permission] Cancelling pending requests for session: ${sessionId}`);
}

/**
 * 获取待处理请求数量（用于调试）
 */
export function getPendingRequestCount(): number {
  return pendingRequests.size;
}