tower/apps/api/src/modules/ai-stream/ai-stream.service.ts

import { Injectable, NotFoundException, BadRequestException, Logger } from '@nestjs/common';
import { ConfigService } from '@nestjs/config';
import { MessageEvent } from '@nestjs/common';
import { Observable } from 'rxjs';
import Redis from 'ioredis';
import { PrismaService } from '../../prisma/prisma.service';
import { CreateDraftDto } from './dto';

const STREAM_KEY = 'tower:stream:messages';
// Block this long per XREAD; on timeout we emit a heartbeat to keep the
// connection warm through proxies, then loop again.
const BLOCK_MS = 15000;

interface StreamRecord {
  messageId: string;
  tenantId: string;
  content: string;
  senderJid: string;
  senderName: string | null;
  sourceGroupId: string;
  sourceGroupName: string | null;
  organizationId: string | null;
  replyToMessageId: string | null;
  tags: string[];
  effectiveAction: string | null;
  traceId: string | null;
  timestamp: string;
}

@Injectable()
export class AiStreamService {
  private readonly logger = new Logger(AiStreamService.name);
  private readonly redisUrl: string;

  constructor(
    private readonly config: ConfigService,
    private readonly prisma: PrismaService,
  ) {
    this.redisUrl = this.config.get<string>('REDIS_URL', 'redis://localhost:6379');
  }

  /**
   * SSE source. Reads the Redis stream from `afterId` and emits one MessageEvent
   * per message, with the stream entry id set as the SSE event id so the client's
   * Last-Event-ID resumes exactly where it dropped off.
   *
   * afterId semantics (Redis stream cursor):
   *   '$'  → only messages that arrive after connect (default for a fresh client)
   *   '0'  → replay the entire retained stream (~50k messages) then follow live
   *   '<id>' → resume strictly after that entry (used on reconnect)
   */
  streamMessages(afterId: string, filters?: { tenantId?: string; sourceGroupId?: string }): Observable<MessageEvent> {
    return new Observable<MessageEvent>((subscriber) => {
      const redis = new Redis(this.redisUrl, { maxRetriesPerRequest: null });
      let active = true;
      let cursor = afterId;

      const loop = async () => {
        while (active) {
          try {
            const res = (await redis.xread(
              'COUNT', 100,
              'BLOCK', BLOCK_MS,
              'STREAMS', STREAM_KEY, cursor,
            )) as [string, [string, string[]][]][] | null;

            if (!res) {
              // Idle timeout — heartbeat keeps the pipe alive. EventSource.onmessage
              // ignores typed events, so this won't reach the dev's message handler.
              if (active) subscriber.next({ type: 'heartbeat', data: '' } as MessageEvent);
              continue;
            }

            for (const [, entries] of res) {
              for (const [id, fields] of entries) {
                const record = parseFields(fields);
                cursor = id;
                if (filters?.tenantId && record.tenantId !== filters.tenantId) continue;
                if (filters?.sourceGroupId && record.sourceGroupId !== filters.sourceGroupId) continue;
                subscriber.next({ id, data: record } as MessageEvent);
              }
            }
          } catch (err) {
            if (active) subscriber.error(err);
            break;
          }
        }
      };

      loop();

      return () => {
        active = false;
        redis.disconnect();
      };
    });
  }

  /**
   * Writeback: the AI dev posts an extracted draft. Creates a ContentDraft in
   * PENDING_REVIEW status — it then shows up in the admin /admin/drafts UI for
   * approve/reject, which publishes the real Event/SevaOpportunity/KnowledgeItem.
   */
  async createDraft(dto: CreateDraftDto): Promise<{ draftId: string }> {
    const message = await this.prisma.message.findUnique({
      where: { id: dto.sourceMessageId },
      select: { id: true, tenantId: true },
    });
    if (!message) {
      throw new NotFoundException('sourceMessageId does not exist');
    }
    if (message.tenantId !== dto.tenantId) {
      throw new BadRequestException('tenantId does not match the source message');
    }

    const draft = await this.prisma.contentDraft.create({
      data: {
        tenantId: dto.tenantId,
        type: dto.type,
        status: 'PENDING_REVIEW',
        sourceMessageId: dto.sourceMessageId,
        proposedJson: dto.proposedJson as object,
        confidence: dto.confidence ?? null,
      },
      select: { id: true },
    });

    this.logger.log(`AI draft created: ${draft.id} (${dto.type}) from message ${dto.sourceMessageId}`);
    return { draftId: draft.id };
  }
}

function parseFields(fields: string[]): StreamRecord {
  const obj: Record<string, string> = {};
  for (let i = 0; i < fields.length; i += 2) {
    obj[fields[i]] = fields[i + 1];
  }
  let tags: string[] = [];
  try {
    tags = obj.tags ? JSON.parse(obj.tags) : [];
  } catch {
    tags = [];
  }
  return {
    messageId: obj.messageId ?? '',
    tenantId: obj.tenantId ?? '',
    content: obj.content ?? '',
    senderJid: obj.senderJid ?? '',
    senderName: obj.senderName || null,
    sourceGroupId: obj.sourceGroupId ?? '',
    sourceGroupName: obj.sourceGroupName || null,
    organizationId: obj.organizationId || null,
    replyToMessageId: obj.replyToMessageId || null,
    tags,
    effectiveAction: obj.effectiveAction || null,
    traceId: obj.traceId || null,
    timestamp: obj.timestamp ?? '',
  };
}