stack/apps/gateway/src/chat/embedded-chat.runtime.ts

import { ForbiddenException, Injectable, Logger, NotFoundException } from '@nestjs/common';
import type { AgentSessionEvent } from '@mariozechner/pi-coding-agent';
import { AgentService, type AgentSession } from '../agent/agent.service.js';
import type { ActorTenantScope } from '../auth/session-scope.js';
import type {
  ChatRuntime,
  LegacyBrowserMessagePayload,
  LegacyEmbeddedChatPort,
  LegacyRuntimeEvent,
  LegacyRuntimeResult,
  LegacySessionPresentation,
  LegacySocketTurnLease,
  LegacyUsage,
  OwnedConversationContext,
  VerifiedDiscordIngressContext,
  VerifiedDiscordTurnLease,
  LegacyRuntimeStream,
} from './chat-runtime.js';

/** Fixed timeout for a synchronous REST turn, matching the historical controller budget. */
const REST_TURN_TIMEOUT_MS = 120_000;

/**
 * The `legacy` chat runtime and the sole implementation of {@link LegacyEmbeddedChatPort}.
 *
 * It owns the embedded in-process execution path — the `AgentService` stack that the
 * `ChatController` and `ChatGateway` drove directly before Task Five. Once the
 * {@link import('./chat-runtime-router.js').ChatRuntimeRouter} fronts it, the browser
 * HTTP/WebSocket legacy path and verified-Discord ingress route through THIS runtime, so
 * neither the controller nor the gateway retains `AgentService`, `piSession`, session,
 * listener, channel, or metric access. Ownership (`userId`/`tenantId`) is re-checked by
 * `AgentService` on every operation; a missing, foreign, or no-longer-owned conversation
 * collapses to `conversation_unavailable` and never throws out of the port.
 */
@Injectable()
export class EmbeddedChatRuntime implements ChatRuntime, LegacyEmbeddedChatPort {
  readonly kind = 'embedded' as const;
  private readonly logger = new Logger(EmbeddedChatRuntime.name);

  constructor(readonly agentService: AgentService) {}

  // -------------------------------------------------------------------------
  // Legacy REST completion (op A)
  // -------------------------------------------------------------------------

  async completeLegacyRestTurn(
    context: OwnedConversationContext,
    input: Readonly<{ content: string }>,
  ): Promise<
    LegacyRuntimeResult<Readonly<{ text: string; presentation: LegacySessionPresentation }>>
  > {
    const scope = toScope(context.scope);
    const { conversationId } = context;

    const resolved = await this.resolveOrCreate(conversationId, scope, {});
    if (!resolved.ok) return resolved;

    let responseText = '';
    let timer: ReturnType<typeof setTimeout> | undefined;
    let detach: (() => void) | undefined;
    let disposed = false;
    // One idempotent teardown owned OUTSIDE the completion promise: it clears the timeout and
    // detaches the event listener exactly once, whichever of agent_end, timeout, or a prompt
    // rejection fires first. Without this, a prompt() rejection surfaced through the catch below
    // would return while leaving the listener attached (free to consume a later turn's events) and
    // the 120s timer live (its rejection later going unobserved).
    const dispose = (): void => {
      if (disposed) return;
      disposed = true;
      if (timer !== undefined) clearTimeout(timer);
      detach?.();
    };
    const done = new Promise<void>((resolve, reject) => {
      timer = setTimeout(() => {
        dispose();
        reject(new Error('Agent response timed out'));
      }, REST_TURN_TIMEOUT_MS);

      detach = this.agentService.onEvent(
        conversationId,
        (event: AgentSessionEvent) => {
          if (
            event.type === 'message_update' &&
            event.assistantMessageEvent.type === 'text_delta'
          ) {
            responseText += event.assistantMessageEvent.delta;
          }
          if (event.type === 'agent_end') {
            dispose();
            resolve();
          }
        },
        scope,
      );
    });

    try {
      await this.agentService.prompt(conversationId, input.content, scope);
      await done;
    } catch (err) {
      dispose();
      const message = err instanceof Error ? err.message : String(err);
      if (message.includes('timed out')) {
        return { ok: false, code: 'timeout', retryable: true };
      }
      this.logger.error(`Legacy REST turn failed for conversation=${conversationId}`, message);
      return { ok: false, code: 'operation_failed', retryable: false };
    }

    const presentation = this.presentationFor(conversationId, scope) ?? resolved.presentation;
    return { ok: true, value: { text: responseText, presentation } };
  }

  // -------------------------------------------------------------------------
  // Legacy Socket streaming (op B)
  // -------------------------------------------------------------------------

  async prepareLegacySocketTurn(
    context: OwnedConversationContext,
    input: LegacyBrowserMessagePayload,
    stream: LegacyRuntimeStream,
  ): Promise<LegacyRuntimeResult<LegacySocketTurnLease>> {
    const scope = toScope(context.scope);
    const { conversationId } = context;

    const resolved = await this.resolveOrCreate(conversationId, scope, {
      ...(input.provider ? { provider: input.provider } : {}),
      ...(input.modelId ? { modelId: input.modelId } : {}),
      ...(input.agentId ? { agentConfigId: input.agentId } : {}),
    });
    if (!resolved.ok) return resolved;

    let detach: () => void;
    try {
      detach = this.subscribe(conversationId, scope, stream);
    } catch (err) {
      // A partial listener/channel setup rolled itself back inside subscribe(); surface a total
      // safe failure instead of throwing out of the port. Retryable — the attach is transient.
      this.logger.error(
        `Embedded socket subscription failed for conversation=${conversationId}`,
        err instanceof Error ? err.message : String(err),
      );
      return { ok: false, code: 'runtime_unavailable', retryable: true };
    }

    return {
      ok: true,
      value: this.buildLease(
        conversationId,
        scope,
        input.content,
        input.attachments,
        detach,
        resolved.presentation,
      ),
    };
  }

  // -------------------------------------------------------------------------
  // Thinking level (op C) — synchronous, total
  // -------------------------------------------------------------------------

  setLegacyThinking(
    context: OwnedConversationContext,
    level: string,
  ): LegacyRuntimeResult<LegacySessionPresentation> {
    const scope = toScope(context.scope);
    const session = this.agentService.getSession(context.conversationId, scope);
    if (!session) return CONVERSATION_UNAVAILABLE;

    const availableThinkingLevels = session.piSession.getAvailableThinkingLevels();
    if (!(availableThinkingLevels as readonly string[]).includes(level)) {
      return {
        ok: false,
        code: 'thinking_level_invalid',
        retryable: false,
        availableThinkingLevels,
      };
    }

    session.piSession.setThinkingLevel(level as never);
    return { ok: true, value: this.presentationForSession(session) };
  }

  // -------------------------------------------------------------------------
  // Abort (op D)
  // -------------------------------------------------------------------------

  async abortLegacyTurn(context: OwnedConversationContext): Promise<LegacyRuntimeResult<void>> {
    const scope = toScope(context.scope);
    const session = this.agentService.getSession(context.conversationId, scope);
    if (!session) return CONVERSATION_UNAVAILABLE;

    try {
      await session.piSession.abort();
    } catch (err) {
      this.logger.error(
        `Legacy abort failed for conversation=${context.conversationId}`,
        err instanceof Error ? err.message : String(err),
      );
      return { ok: false, code: 'operation_failed', retryable: false };
    }
    return { ok: true, value: undefined };
  }

  // -------------------------------------------------------------------------
  // Model override (synchronous, total)
  // -------------------------------------------------------------------------

  applyLegacyModelOverride(
    context: OwnedConversationContext,
    modelId: string,
  ): LegacyRuntimeResult<LegacySessionPresentation> {
    const scope = toScope(context.scope);
    const session = this.agentService.getSession(context.conversationId, scope);
    if (!session) return CONVERSATION_UNAVAILABLE;

    this.agentService.updateSessionModel(context.conversationId, modelId, scope);
    const refreshed = this.agentService.getSession(context.conversationId, scope) ?? session;
    return { ok: true, value: this.presentationForSession(refreshed) };
  }

  // -------------------------------------------------------------------------
  // Presentation read (synchronous, total)
  // -------------------------------------------------------------------------

  readLegacySessionPresentation(
    context: OwnedConversationContext,
  ): LegacyRuntimeResult<LegacySessionPresentation> {
    const scope = toScope(context.scope);
    const session = this.agentService.getSession(context.conversationId, scope);
    if (!session) return CONVERSATION_UNAVAILABLE;
    return { ok: true, value: this.presentationForSession(session) };
  }

  // -------------------------------------------------------------------------
  // Verified Discord ingress (embedded-only in both modes)
  // -------------------------------------------------------------------------

  async dispatchVerifiedDiscordIngress(
    context: VerifiedDiscordIngressContext,
    stream: LegacyRuntimeStream,
  ): Promise<LegacyRuntimeResult<VerifiedDiscordTurnLease>> {
    const scope = toScope(context.scope);
    const { conversationId } = context;

    const resolved = await this.resolveOrCreate(
      conversationId,
      scope,
      { agentConfigId: context.configuredAgent.agentConfigId },
      {
        agentConfigId: context.configuredAgent.agentConfigId,
        instanceId: context.configuredAgent.instanceId,
      },
    );
    if (!resolved.ok) return resolved;

    let detach: () => void;
    try {
      detach = this.subscribe(conversationId, scope, stream);
    } catch (err) {
      // A partial listener/channel setup rolled itself back inside subscribe(); surface a total
      // safe failure instead of throwing out of the port. Retryable — the attach is transient.
      this.logger.error(
        `Embedded Discord subscription failed for conversation=${conversationId}`,
        err instanceof Error ? err.message : String(err),
      );
      return { ok: false, code: 'runtime_unavailable', retryable: true };
    }

    return {
      ok: true,
      value: this.buildLease(
        conversationId,
        scope,
        context.content,
        context.attachments,
        detach,
        resolved.presentation,
      ),
    };
  }

  // -------------------------------------------------------------------------
  // Shared helpers
  // -------------------------------------------------------------------------

  /**
   * Resolves the owned session, creating it on first use. Ownership/scope rejections
   * (`Forbidden`/`NotFound`) collapse to `conversation_unavailable`; any other creation
   * failure surfaces as the retryable `runtime_unavailable`. On success returns the
   * session presentation so callers avoid a redundant `getSession`.
   */
  private async resolveOrCreate(
    conversationId: string,
    scope: ActorTenantScope,
    extraOptions: Readonly<{ provider?: string; modelId?: string; agentConfigId?: string }>,
    expectedAgent?: Readonly<{ agentConfigId: string; instanceId: string }>,
  ): Promise<
    | { readonly ok: true; readonly presentation: LegacySessionPresentation }
    | Exclude<LegacyRuntimeResult<never>, { ok: true }>
  > {
    let session = this.agentService.getSession(conversationId, scope);
    if (session && expectedAgent) {
      // A verified-Discord turn reuses an existing in-process session only when its configured
      // identity still matches the reconciled agent record. A pre-existing same-scope session
      // carrying a different configured agent (config id or resolved name) is rejected rather than
      // executed under the verified label — never silently run a different prompt/model/tool policy.
      if (
        session.agentConfigId !== expectedAgent.agentConfigId ||
        session.agentName !== expectedAgent.instanceId
      ) {
        return CONVERSATION_UNAVAILABLE;
      }
    }
    if (!session) {
      try {
        session = await this.agentService.createSession(conversationId, {
          userId: scope.userId,
          tenantId: scope.tenantId,
          ...extraOptions,
        });
      } catch (err) {
        if (err instanceof ForbiddenException || err instanceof NotFoundException) {
          return CONVERSATION_UNAVAILABLE;
        }
        this.logger.error(
          `Embedded session creation failed for conversation=${conversationId}`,
          err instanceof Error ? err.stack : String(err),
        );
        return { ok: false, code: 'runtime_unavailable', retryable: true };
      }
    }
    return { ok: true, presentation: this.presentationForSession(session) };
  }

  /** Installs a normalizing event listener that forwards to the server-owned stream. */
  private subscribe(
    conversationId: string,
    scope: ActorTenantScope,
    stream: LegacyRuntimeStream,
  ): () => void {
    const unsubscribe = this.agentService.onEvent(
      conversationId,
      (event: AgentSessionEvent) => {
        const normalized = this.normalizeEvent(conversationId, scope, event);
        if (normalized) stream.onEvent(normalized);
      },
      scope,
    );
    try {
      this.agentService.addChannel(conversationId, stream.channelId, scope);
    } catch (err) {
      // Partial setup: the listener was acquired but the channel attach failed. Roll back
      // exactly what was acquired (the listener) before the failure escapes, so no leaked
      // subscription survives; the caller converts the rethrow into a total safe failure.
      try {
        unsubscribe();
      } catch {
        /* idempotent teardown */
      }
      throw err;
    }
    return () => {
      try {
        unsubscribe();
      } catch {
        /* idempotent teardown */
      }
      try {
        this.agentService.removeChannel(conversationId, stream.channelId, scope);
      } catch {
        /* idempotent teardown */
      }
    };
  }

  /** Builds an atomically one-shot, scope-rechecking dispatch lease. */
  private buildLease(
    conversationId: string,
    scope: ActorTenantScope,
    content: string,
    attachments: VerifiedDiscordIngressContext['attachments'],
    detach: () => void,
    presentation: LegacySessionPresentation,
  ): LegacySocketTurnLease & VerifiedDiscordTurnLease {
    let dispatched = false;
    let disposed = false;
    return {
      presentation,
      dispatch: async (): Promise<LegacyRuntimeResult<void>> => {
        if (dispatched) {
          return { ok: false, code: 'turn_already_dispatched', retryable: false };
        }
        dispatched = true;
        try {
          await this.agentService.prompt(conversationId, content, scope, attachments);
        } catch (err) {
          this.logger.error(
            `Legacy dispatch failed for conversation=${conversationId}`,
            err instanceof Error ? err.message : String(err),
          );
          return { ok: false, code: 'operation_failed', retryable: false };
        }
        return { ok: true, value: undefined };
      },
      dispose: async (): Promise<void> => {
        if (disposed) return;
        disposed = true;
        detach();
      },
    };
  }

  /** Normalizes a raw agent event into the redaction-agnostic transport event, or drops it. */
  private normalizeEvent(
    conversationId: string,
    scope: ActorTenantScope,
    event: AgentSessionEvent,
  ): LegacyRuntimeEvent | undefined {
    switch (event.type) {
      case 'agent_start':
        return { type: 'started' };
      case 'agent_end':
        return { type: 'settled', ...this.usageFor(conversationId, scope) };
      case 'message_update': {
        const assistant = event.assistantMessageEvent;
        if (assistant.type === 'text_delta') return { type: 'text_delta', text: assistant.delta };
        if (assistant.type === 'thinking_delta') {
          return { type: 'thinking_delta', text: assistant.delta };
        }
        return undefined;
      }
      case 'tool_execution_start':
        return { type: 'tool_started', toolCallId: event.toolCallId, toolName: event.toolName };
      case 'tool_execution_end':
        return {
          type: 'tool_finished',
          toolCallId: event.toolCallId,
          toolName: event.toolName,
          isError: event.isError,
        };
      default:
        return undefined;
    }
  }

  /**
   * Gathers terminal usage from the Pi session and records it into session metrics.
   * Embedded owns AgentService metrics; the gateway never touches `piSession` stats.
   */
  private usageFor(conversationId: string, scope: ActorTenantScope): { usage?: LegacyUsage } {
    const session = this.agentService.getSession(conversationId, scope);
    const piSession = session?.piSession;
    const stats = piSession?.getSessionStats();
    if (!session || !stats) return {};
    const contextUsage = piSession?.getContextUsage();

    const tokens = {
      input: stats.tokens?.input ?? 0,
      output: stats.tokens?.output ?? 0,
      cacheRead: stats.tokens?.cacheRead ?? 0,
      cacheWrite: stats.tokens?.cacheWrite ?? 0,
      total: stats.tokens?.total ?? 0,
    };

    this.agentService.recordTokenUsage(conversationId, { ...tokens });

    return {
      usage: {
        provider: session.provider,
        modelId: session.modelId,
        thinkingLevel: piSession?.thinkingLevel ?? 'off',
        tokens,
        cost: stats.cost ?? 0,
        context: {
          percent: contextUsage?.percent ?? null,
          window: contextUsage?.contextWindow ?? 0,
        },
      },
    };
  }

  /** Presentation from a live session id, or undefined when no owned session exists. */
  private presentationFor(
    conversationId: string,
    scope: ActorTenantScope,
  ): LegacySessionPresentation | undefined {
    const session = this.agentService.getSession(conversationId, scope);
    return session ? this.presentationForSession(session) : undefined;
  }

  /** User-facing projection carrying no session handle, credential, or raw stats. */
  private presentationForSession(session: AgentSession): LegacySessionPresentation {
    return {
      provider: session.provider,
      modelId: session.modelId,
      thinkingLevel: session.piSession.thinkingLevel,
      availableThinkingLevels: session.piSession.getAvailableThinkingLevels(),
      ...(session.agentName ? { agentName: session.agentName } : {}),
    };
  }
}

/** The shared terminal `conversation_unavailable` failure (missing/foreign/lost ownership). */
const CONVERSATION_UNAVAILABLE = {
  ok: false as const,
  code: 'conversation_unavailable' as const,
  retryable: false as const,
};

/** Narrows a branded context scope to the `AgentService` actor/tenant scope (identical shape). */
function toScope(scope: Readonly<{ userId: string; tenantId: string }>): ActorTenantScope {
  return { userId: scope.userId, tenantId: scope.tenantId };
}