apf_portal/apps/portal-bff/src/grpc/ai-bridge/ai-bridge.controller.ts

import {
  Body,
  Controller,
  Get,
  HttpStatus,
  Post,
  Query,
  Req,
  Res,
  UnauthorizedException,
} from '@nestjs/common';
import { ApiCookieAuth, ApiOperation, ApiTags } from '@nestjs/swagger';
import { status as GrpcStatus, type ServiceError } from '@grpc/grpc-js';
import type { Request, Response } from 'express';
import { ChatClient } from '../ai-client/chat.client';
import { ModelsClient } from '../ai-client/models.client';
import { PrincipalMapper } from '../ai-client/principal.mapper';
import { RagClient } from '../ai-client/rag.client';
import type { ChatEvent, ChatRequest } from '../gen/apf-ai/chat';
import { ChatRole } from '../gen/apf-ai/common';
import type { ListModelsResponse } from '../gen/apf-ai/models';
import type { RagSearchResponse } from '../gen/apf-ai/rag';
import { ChatRequestDto, type ChatMessageDto } from './dto/chat-request.dto';
import { RagSearchQueryDto } from './dto/rag-search-query.dto';
import { chatEventToSseFrame, relayErrorFrame } from './sse.writer';
import type { AuthenticatedUser } from '../../auth/auth.service';

/**
 * BFF-facing surface of the AI relay, per
 * [ADR-0024](../../../../../docs/decisions/0024-ai-service-relay-grpc-sse-bridge.md):
 *
 *   - `POST /api/ai/chat`        — streaming chat, bridged from
 *                                  `ChatService.Chat` (gRPC server-
 *                                  stream) to `text/event-stream`.
 *   - `GET  /api/ai/rag/search`  — unary RAG retrieval.
 *   - `GET  /api/ai/models`      — list configured providers.
 *
 * No `@RequireAdmin()` — AI features are end-user surfaces, gated
 * only by the active portal session (`req.session.user`). The
 * session + CSRF middleware mounted in `main.ts` covers the
 * authentication + double-submit token before the controller
 * runs; this class asserts presence of `session.user` for the
 * routes that need it and lets the global middleware handle the
 * rest.
 */
@ApiTags('ai')
@ApiCookieAuth('portal_session')
@Controller('ai')
export class AiBridgeController {
  constructor(
    private readonly chatClient: ChatClient,
    private readonly ragClient: RagClient,
    private readonly modelsClient: ModelsClient,
    private readonly principalMapper: PrincipalMapper,
  ) {}

  @ApiOperation({
    summary:
      'Streaming chat. Returns text/event-stream with one frame per ChatEvent (token / citation / agent-step / tool-call / error / done).',
  })
  @Post('chat')
  async chat(
    @Body() body: ChatRequestDto,
    @Req() req: Request,
    @Res() res: Response,
  ): Promise<void> {
    const user = requireSessionUser(req);

    const principal = this.principalMapper.fromInputs({
      oid: user.oid,
      tid: user.tid,
      roles: user.roles,
    });
    const protoRequest: ChatRequest = {
      messages: body.messages.map(toProtoMessage),
      conversationId: body.conversationId ?? '',
      model: body.model ?? '',
      provider: body.provider ?? '',
      toolsAvailable: [],
      // v1: tool registry is empty (per ADR-0024 §"Tool-dispatch
      // contract"). The AI service will never emit `tool_call` for
      // requests with `toolsAvailable: []`; the SSE writer still
      // covers the case in case a future tool lands.
      rag: { enabled: false, topK: 0 },
      principal,
    };

    writeSseHeaders(res);

    // Browser disconnect (tab close, fetch abort, network drop) →
    // AbortController → gRPC call.cancel() → upstream LLM stop.
    // Per ADR-0024 §"SSE bridge between BFF and SPA".
    const abort = new AbortController();
    req.on('close', () => {
      if (!res.writableEnded) {
        abort.abort();
      }
    });

    const stream = this.chatClient.chat(protoRequest, { signal: abort.signal });

    try {
      for await (const event of stream as AsyncIterable<ChatEvent>) {
        const frame = chatEventToSseFrame(event);
        if (frame !== null) {
          res.write(frame);
        }
      }
    } catch (err) {
      // gRPC `CANCELLED` after `abort.abort()` is the expected exit
      // path on browser disconnect — do not surface an error frame
      // because the response is already closing. Anything else
      // becomes a structured error event so the SPA's renderer
      // sees the failure rather than a torn-down connection.
      const code = (err as Partial<ServiceError>).code;
      const cancelled = code === GrpcStatus.CANCELLED;
      if (!cancelled && !res.writableEnded) {
        res.write(
          relayErrorFrame(
            mapServiceErrorCode(code),
            (err as Error).message ?? 'AI relay error',
            isRetriable(code),
          ),
        );
      }
    } finally {
      if (!res.writableEnded) {
        res.end();
      }
    }
  }

  @ApiOperation({ summary: 'Unary RAG retrieval bounded by the caller principal.' })
  @Get('rag/search')
  async ragSearch(
    @Query() query: RagSearchQueryDto,
    @Req() req: Request,
  ): Promise<RagSearchResponse> {
    const user = requireSessionUser(req);
    const principal = this.principalMapper.fromInputs({
      oid: user.oid,
      tid: user.tid,
      roles: user.roles,
    });

    return this.ragClient.search({
      query: query.query,
      topK: query.topK ?? 0,
      filters: {
        source: query.source ?? '',
        documentId: query.documentId ?? '',
      },
      principal,
    });
  }

  @ApiOperation({ summary: 'List configured AI providers and the active provider.' })
  @Get('models')
  async listModels(@Req() req: Request): Promise<ListModelsResponse> {
    requireSessionUser(req);
    return this.modelsClient.listModels({});
  }
}

function requireSessionUser(req: Request): AuthenticatedUser {
  const user = req.session.user;
  if (!user) {
    throw new UnauthorizedException({
      code: 'unauthenticated',
      message: 'The AI surface requires an authenticated portal session.',
    });
  }
  return user;
}

function writeSseHeaders(res: Response): void {
  res.statusCode = HttpStatus.OK;
  res.setHeader('Content-Type', 'text/event-stream; charset=utf-8');
  res.setHeader('Cache-Control', 'no-cache, no-transform');
  res.setHeader('Connection', 'keep-alive');
  // nginx-style buffering hint — keeps reverse-proxies from
  // accumulating chunks before forwarding, which would defeat the
  // streaming UX.
  res.setHeader('X-Accel-Buffering', 'no');
  res.flushHeaders();
}

function toProtoMessage(message: ChatMessageDto): {
  role: ChatRole;
  content: string;
  toolCallId: string;
  name: string;
} {
  return {
    role: ROLE_PROTO[message.role],
    content: message.content,
    toolCallId: '',
    name: '',
  };
}

const ROLE_PROTO: Record<ChatMessageDto['role'], ChatRole> = {
  system: ChatRole.CHAT_ROLE_SYSTEM,
  user: ChatRole.CHAT_ROLE_USER,
  assistant: ChatRole.CHAT_ROLE_ASSISTANT,
};

/**
 * gRPC status → SSE error code. Keeps the urn:apf-ai:* namespace
 * the AI service itself uses so a relay-side error and an upstream
 * error look the same to the SPA's renderer.
 */
function mapServiceErrorCode(code: number | undefined): string {
  switch (code) {
    case GrpcStatus.UNAVAILABLE:
      return 'urn:apf-ai:unavailable';
    case GrpcStatus.DEADLINE_EXCEEDED:
      return 'urn:apf-ai:timeout';
    case GrpcStatus.PERMISSION_DENIED:
      return 'urn:apf-ai:permission_denied';
    case GrpcStatus.RESOURCE_EXHAUSTED:
      return 'urn:apf-ai:rate_limited';
    case GrpcStatus.INVALID_ARGUMENT:
      return 'urn:apf-ai:invalid_argument';
    default:
      return 'urn:apf-ai:relay_error';
  }
}

function isRetriable(code: number | undefined): boolean {
  return code === GrpcStatus.UNAVAILABLE || code === GrpcStatus.DEADLINE_EXCEEDED;
}