stack/apps/gateway/src/agent/runtime-provider-registry.service.ts

import { ForbiddenException, Inject, Injectable, Logger, NotFoundException } from '@nestjs/common';
import { AgentRuntimeProviderRegistry } from '@mosaicstack/agent';
import {
  createRuntimeAuditLogEntry,
  type LogService,
  type RuntimeAuditErrorCode,
} from '@mosaicstack/log';
import type {
  AgentRuntimeProvider,
  RuntimeAttachHandle,
  RuntimeAttachMode,
  RuntimeCapability,
  RuntimeCapabilitySet,
  RuntimeHealth,
  RuntimeMessage,
  RuntimeScope,
  RuntimeSession,
  RuntimeSessionTree,
  RuntimeStreamEvent,
} from '@mosaicstack/types';
import type { ActorTenantScope } from '../auth/session-scope.js';
import { LOG_SERVICE } from '../log/log.tokens.js';

export const AGENT_RUNTIME_PROVIDER_REGISTRY = Symbol('AGENT_RUNTIME_PROVIDER_REGISTRY');
export const RUNTIME_PROVIDER_AUDIT_SINK = Symbol('RUNTIME_PROVIDER_AUDIT_SINK');
export const RUNTIME_APPROVAL_VERIFIER = Symbol('RUNTIME_APPROVAL_VERIFIER');

export type RuntimeProviderOperation =
  | RuntimeCapability
  | 'runtime.capabilities'
  | 'runtime.health';
export type RuntimeProviderAuditOutcome = 'requested' | 'succeeded' | 'denied' | 'failed';

/** Trusted server-side context only; it intentionally excludes client-provided identity fields. */
export interface RuntimeProviderRequestContext {
  actorScope: ActorTenantScope;
  channelId: string;
  correlationId: string;
}

/** Metadata-only audit record. Message bodies, idempotency keys, and approval refs are excluded. */
export interface RuntimeAuditEvent {
  providerId: string;
  operation: RuntimeProviderOperation;
  outcome: RuntimeProviderAuditOutcome;
  actorId: string;
  tenantId: string;
  channelId: string;
  correlationId: string;
  resourceId?: string;
  durationMs?: number;
  errorCode?: RuntimeAuditErrorCode;
}

export interface RuntimeAuditSink {
  record(event: RuntimeAuditEvent): Promise<void>;
}

/** Exact action shape that a durable approval implementation must consume once. */
export interface RuntimeTerminationAction {
  providerId: string;
  sessionId: string;
  actorId: string;
  tenantId: string;
  channelId: string;
  correlationId: string;
  agentName: string;
}

export interface RuntimeApprovalVerifier {
  consume(approvalRef: string, action: RuntimeTerminationAction): Promise<boolean>;
}

function configuredAgentName(): string {
  const agentName = process.env['MOSAIC_AGENT_NAME']?.trim();
  if (!agentName) throw new RuntimeApprovalDeniedError();
  return agentName;
}

class RuntimeApprovalDeniedError extends Error {
  constructor() {
    super('Runtime termination approval denied');
  }
}

/**
 * The default denies all runtime termination until a durable, exact-action
 * approval implementation is configured. This is safer than a permissive stub.
 */
@Injectable()
export class DenyRuntimeApprovalVerifier implements RuntimeApprovalVerifier {
  async consume(_approvalRef: string, _action: RuntimeTerminationAction): Promise<boolean> {
    return false;
  }
}

/**
 * Temporary metadata-only audit sink. M1 observability can replace this token
 * with a durable audit writer without changing provider call sites.
 */
@Injectable()
export class RuntimeProviderAuditService implements RuntimeAuditSink {
  private readonly logger = new Logger(RuntimeProviderAuditService.name);

  constructor(@Inject(LOG_SERVICE) private readonly logService: LogService) {}

  async record(event: RuntimeAuditEvent): Promise<void> {
    const entry = createRuntimeAuditLogEntry(event);
    await this.logService.logs.ingest(entry);
    this.logger.log(JSON.stringify({ event: entry.content, metadata: entry.metadata }));
  }
}

@Injectable()
export class RuntimeProviderService {
  private readonly logger = new Logger(RuntimeProviderService.name);

  constructor(
    @Inject(AGENT_RUNTIME_PROVIDER_REGISTRY)
    private readonly registry: AgentRuntimeProviderRegistry,
    @Inject(RUNTIME_PROVIDER_AUDIT_SINK)
    private readonly audit: RuntimeAuditSink,
    @Inject(RUNTIME_APPROVAL_VERIFIER)
    private readonly approvals: RuntimeApprovalVerifier,
  ) {}

  async capabilities(
    providerId: string,
    context: RuntimeProviderRequestContext,
  ): Promise<RuntimeCapabilitySet> {
    return this.execute(
      providerId,
      'runtime.capabilities',
      undefined,
      undefined,
      context,
      (provider: AgentRuntimeProvider, scope: RuntimeScope): Promise<RuntimeCapabilitySet> =>
        provider.capabilities(scope),
    );
  }

  async health(providerId: string, context: RuntimeProviderRequestContext): Promise<RuntimeHealth> {
    return this.execute(
      providerId,
      'runtime.health',
      undefined,
      undefined,
      context,
      (provider: AgentRuntimeProvider, scope: RuntimeScope): Promise<RuntimeHealth> =>
        provider.health(scope),
    );
  }

  async listSessions(
    providerId: string,
    context: RuntimeProviderRequestContext,
  ): Promise<RuntimeSession[]> {
    return this.execute(
      providerId,
      'session.list',
      'session.list',
      undefined,
      context,
      (provider: AgentRuntimeProvider, scope: RuntimeScope): Promise<RuntimeSession[]> =>
        provider.listSessions(scope),
    );
  }

  async getSessionTree(
    providerId: string,
    context: RuntimeProviderRequestContext,
  ): Promise<RuntimeSessionTree[]> {
    return this.execute(
      providerId,
      'session.tree',
      'session.tree',
      undefined,
      context,
      (provider: AgentRuntimeProvider, scope: RuntimeScope): Promise<RuntimeSessionTree[]> =>
        provider.getSessionTree(scope),
    );
  }

  streamSession(
    providerId: string,
    sessionId: string,
    cursor: string | undefined,
    context: RuntimeProviderRequestContext,
  ): AsyncIterable<RuntimeStreamEvent> {
    return this.stream(
      providerId,
      'session.stream',
      'session.stream',
      sessionId,
      context,
      (provider: AgentRuntimeProvider, scope: RuntimeScope): AsyncIterable<RuntimeStreamEvent> =>
        provider.streamSession(sessionId, cursor, scope),
    );
  }

  async sendMessage(
    providerId: string,
    sessionId: string,
    message: RuntimeMessage,
    context: RuntimeProviderRequestContext,
  ): Promise<void> {
    await this.execute(
      providerId,
      'session.send',
      'session.send',
      sessionId,
      context,
      (provider: AgentRuntimeProvider, scope: RuntimeScope): Promise<void> =>
        provider.sendMessage(sessionId, message, scope),
    );
  }

  async attach(
    providerId: string,
    sessionId: string,
    mode: RuntimeAttachMode,
    context: RuntimeProviderRequestContext,
  ): Promise<RuntimeAttachHandle> {
    return this.execute(
      providerId,
      'session.attach',
      'session.attach',
      sessionId,
      context,
      (provider: AgentRuntimeProvider, scope: RuntimeScope): Promise<RuntimeAttachHandle> =>
        provider.attach(sessionId, mode, scope),
    );
  }

  async detach(
    providerId: string,
    attachmentId: string,
    context: RuntimeProviderRequestContext,
  ): Promise<void> {
    await this.execute(
      providerId,
      'session.attach',
      'session.attach',
      attachmentId,
      context,
      (provider: AgentRuntimeProvider, scope: RuntimeScope): Promise<void> =>
        provider.detach(attachmentId, scope),
    );
  }

  async terminate(
    providerId: string,
    sessionId: string,
    approvalRef: string,
    context: RuntimeProviderRequestContext,
  ): Promise<void> {
    await this.execute(
      providerId,
      'session.terminate',
      'session.terminate',
      sessionId,
      context,
      async (provider: AgentRuntimeProvider, scope: RuntimeScope): Promise<void> => {
        const approved = await this.approvals.consume(approvalRef, {
          providerId,
          sessionId,
          actorId: scope.actorId,
          tenantId: scope.tenantId,
          channelId: scope.channelId,
          correlationId: scope.correlationId,
          agentName: configuredAgentName(),
        });
        if (!approved) {
          throw new RuntimeApprovalDeniedError();
        }
        await provider.terminate(sessionId, approvalRef, scope);
      },
    );
  }

  private async execute<T>(
    providerId: string,
    operation: RuntimeProviderOperation,
    requiredCapability: RuntimeCapability | undefined,
    resourceId: string | undefined,
    context: RuntimeProviderRequestContext,
    invoke: (provider: AgentRuntimeProvider, scope: RuntimeScope) => Promise<T>,
  ): Promise<T> {
    const scope = this.deriveScope(context);
    const startedAt = Date.now();
    await this.record(providerId, operation, 'requested', scope, resourceId);
    let invocationStarted = false;
    try {
      const provider = this.provider(providerId);
      if (requiredCapability) {
        await this.assertCapability(provider, requiredCapability, scope);
      }
      invocationStarted = true;
      const result = await invoke(provider, scope);
      await this.recordCompletion(providerId, operation, scope, resourceId, Date.now() - startedAt);
      return result;
    } catch (error: unknown) {
      const durationMs = Date.now() - startedAt;
      if (invocationStarted && !(error instanceof RuntimeApprovalDeniedError)) {
        await this.recordFailure(providerId, operation, scope, resourceId, durationMs);
      } else {
        await this.record(
          providerId,
          operation,
          'denied',
          scope,
          resourceId,
          durationMs,
          'policy_denied',
        );
      }
      throw error;
    }
  }

  private async *stream(
    providerId: string,
    operation: RuntimeProviderOperation,
    requiredCapability: RuntimeCapability,
    resourceId: string,
    context: RuntimeProviderRequestContext,
    invoke: (
      provider: AgentRuntimeProvider,
      scope: RuntimeScope,
    ) => AsyncIterable<RuntimeStreamEvent>,
  ): AsyncIterable<RuntimeStreamEvent> {
    const scope = this.deriveScope(context);
    const startedAt = Date.now();
    await this.record(providerId, operation, 'requested', scope, resourceId);
    let invocationStarted = false;
    try {
      const provider = this.provider(providerId);
      await this.assertCapability(provider, requiredCapability, scope);
      invocationStarted = true;
      for await (const event of invoke(provider, scope)) {
        yield event;
      }
      await this.recordCompletion(providerId, operation, scope, resourceId, Date.now() - startedAt);
    } catch (error: unknown) {
      const durationMs = Date.now() - startedAt;
      if (invocationStarted) {
        await this.recordFailure(providerId, operation, scope, resourceId, durationMs);
      } else {
        await this.record(
          providerId,
          operation,
          'denied',
          scope,
          resourceId,
          durationMs,
          'policy_denied',
        );
      }
      throw error;
    }
  }

  private provider(providerId: string): AgentRuntimeProvider {
    try {
      return this.registry.require(providerId);
    } catch (error: unknown) {
      const message = error instanceof Error ? error.message : 'Runtime provider is not registered';
      throw new NotFoundException(message);
    }
  }

  private async assertCapability(
    provider: AgentRuntimeProvider,
    requiredCapability: RuntimeCapability,
    scope: RuntimeScope,
  ): Promise<void> {
    const capabilities = await provider.capabilities(scope);
    if (!capabilities.supported.includes(requiredCapability)) {
      throw new ForbiddenException(`Runtime provider capability denied: ${requiredCapability}`);
    }
  }

  private deriveScope(context: RuntimeProviderRequestContext): RuntimeScope {
    const actorId = context.actorScope.userId.trim();
    const tenantId = context.actorScope.tenantId.trim();
    const channelId = context.channelId.trim();
    const correlationId = context.correlationId.trim();
    if (!actorId || !tenantId || !channelId || !correlationId) {
      throw new ForbiddenException(
        'Authenticated runtime actor scope and correlation are required',
      );
    }
    return Object.freeze({ actorId, tenantId, channelId, correlationId });
  }

  private async recordFailure(
    providerId: string,
    operation: RuntimeProviderOperation,
    scope: RuntimeScope,
    resourceId: string | undefined,
    durationMs: number,
  ): Promise<void> {
    try {
      await this.record(
        providerId,
        operation,
        'failed',
        scope,
        resourceId,
        durationMs,
        'provider_error',
      );
    } catch {
      this.logger.error(
        `Runtime provider failure audit failed provider=${providerId} operation=${operation} correlation=${scope.correlationId}`,
      );
    }
  }

  private async recordCompletion(
    providerId: string,
    operation: RuntimeProviderOperation,
    scope: RuntimeScope,
    resourceId: string | undefined,
    durationMs: number,
  ): Promise<void> {
    try {
      await this.record(providerId, operation, 'succeeded', scope, resourceId, durationMs);
    } catch {
      this.logger.error(
        `Runtime provider completion audit failed provider=${providerId} operation=${operation} correlation=${scope.correlationId}`,
      );
    }
  }

  private async record(
    providerId: string,
    operation: RuntimeProviderOperation,
    outcome: RuntimeProviderAuditOutcome,
    scope: RuntimeScope,
    resourceId: string | undefined,
    durationMs?: number,
    errorCode?: RuntimeAuditErrorCode,
  ): Promise<void> {
    await this.audit.record({
      providerId,
      operation,
      outcome,
      actorId: scope.actorId,
      tenantId: scope.tenantId,
      channelId: scope.channelId,
      correlationId: scope.correlationId,
      ...(resourceId ? { resourceId } : {}),
      ...(durationMs !== undefined ? { durationMs } : {}),
      ...(errorCode ? { errorCode } : {}),
    });
  }
}