packages/runtime-host/src/server/connection-session.ts

/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License.  You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing,
 * software distributed under the License is distributed on an
 * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
 * KIND, either express or implied.  See the License for the
 * specific language governing permissions and limitations
 * under the License.
 */

import {
  decodeClientFrame,
  isClientCapabilityClientFrameKind,
  RUNTIME_HOST_MAX_IN_FLIGHT_DOMAIN_REQUESTS,
  type ClientCapabilityClientFrame,
  type HostOperationErrorCode,
  type RequestFrame,
} from '../protocol/index.js';
import { runtimeHostLogBuffer } from '../process-diagnostics.js';
import type { RuntimeHostMessageTransport } from '../transport/message-transport.js';
import {
  dispatchOperation,
  operationFailureResponse,
  type ConnectionContext,
  type OperationHandlerMap,
  type OperationResidency,
} from './operation-dispatcher.js';
import { BoundedSerialOutboundWriter } from './serial-outbound-writer.js';
import { RuntimeHostTransportError } from '../transport/framed-transport.js';
import type {
  SessionContinuityConnection,
  SessionContinuityService,
} from './session-continuity-service.js';
import type {
  ClientCapabilityConnection,
  ClientCapabilityService,
} from './client-capability-service.js';
import type {
  HostChangeFeed,
  HostChangeSubscription,
  HostChangeSubscriptionMask,
} from './host-change-feed.js';
import type { RuntimeHostConnectionAuthority } from './connection-authority.js';
import {
  authorizeClientCapabilityFrame,
  authorizeRuntimeHostOperation,
  hasRuntimeHostOperationGrant,
} from './connection-authority.js';
import { boundedFailureDiagnostic } from './failure-diagnostic.js';

type AcceptedConnectionContext = Omit<
  ConnectionContext,
  'acquireResidency' | 'principal' | 'inputClosedSignal'
> & {
  readonly clientInstanceId: string;
  readonly authority: RuntimeHostConnectionAuthority;
};

export interface ConnectionOperationLease {
  acquireResidency(): OperationResidency;
  seal(): void;
  finish(): void;
}

export interface RuntimeHostConnectionSessionOptions {
  transport: RuntimeHostMessageTransport;
  connection: AcceptedConnectionContext;
  resolveHandlers(): OperationHandlerMap;
  resolveContinuity(): SessionContinuityService | undefined;
  resolveClientCapabilities?(): ClientCapabilityService | undefined;
  resolveHostChanges?(): HostChangeFeed | undefined;
  resolveSharedSessionId?(): string | undefined;
  beginOperation(frame: RequestFrame): Promise<ConnectionOperationLease | HostOperationErrorCode>;
  onDiagnostic?(diagnostic: string): void;
  onTeardown(): void;
}

export class RuntimeHostConnectionSession {
  readonly #options: RuntimeHostConnectionSessionOptions;
  readonly #writer: BoundedSerialOutboundWriter;
  readonly #onDiagnostic: (diagnostic: string) => void;
  readonly #requests = new Map<string, Promise<void>>();
  #transcriptPageTail: Promise<void> = Promise.resolve();
  #inFlightStatusRequests = 0;
  #continuityService: SessionContinuityService | undefined;
  #continuity: SessionContinuityConnection | undefined;
  #clientCapabilityService: ClientCapabilityService | undefined;
  #clientCapabilities: ClientCapabilityConnection | undefined;
  #clientCapabilityCloseTask: Promise<void> | undefined;
  #hostChanges: HostChangeSubscription | undefined;
  readonly #inputClosedAbort = new AbortController();
  #closed = false;

  constructor(options: RuntimeHostConnectionSessionOptions) {
    this.#options = options;
    this.#onDiagnostic =
      options.onDiagnostic ?? ((diagnostic) => runtimeHostLogBuffer.append('error', diagnostic));
    this.#writer = new BoundedSerialOutboundWriter(options.transport, (error) => this.#fail(error));
  }

  async run(): Promise<void> {
    this.attachGlobalChanges();
    try {
      try {
        await this.#pumpInbound();
      } catch (error) {
        if (!isReadEof(error)) throw error;
        await this.#closeAfterDispatchedReplies();
      }
    } catch (error) {
      this.#fail(error);
    } finally {
      this.#teardown();
      await Promise.allSettled(this.#requests.values());
      await Promise.all([
        this.#writer.settled(),
        this.#options.transport.closed,
        this.#clientCapabilityCloseTask?.catch(() => undefined),
      ]);
    }
  }

  async #closeAfterDispatchedReplies(): Promise<void> {
    this.#inputClosedAbort.abort();
    this.#detachContinuity();
    this.#detachClientCapabilities();
    this.#detachHostChanges();
    const outcome = await Promise.race([
      Promise.allSettled([...this.#requests.values()]).then(() => 'drained' as const),
      this.#options.transport.closed.then(() => 'closed' as const),
    ]);
    if (outcome === 'closed') {
      this.#teardown();
      return;
    }
    if (this.#closed) return;
    await this.#writer.settled();
    if (this.#closed) return;
    this.#closed = true;
    this.#writer.close();
    this.#options.transport.closeAfterFlush();
    this.#options.onTeardown();
  }

  async #pumpInbound(): Promise<void> {
    while (!this.#closed) {
      const frame = decodeClientFrame(await this.#options.transport.read(0));
      if ('kind' in frame) {
        if (isClientCapabilityClientFrameKind(frame.kind)) {
          const capabilityFrame = frame as ClientCapabilityClientFrame;
          if (
            !authorizeClientCapabilityFrame(this.#options.connection.authority, capabilityFrame)
          ) {
            this.#fail(new Error('Runtime Host Client Capability frame is not authorized'));
            return;
          }
          this.#ensureClientCapabilities()?.accept(capabilityFrame);
          continue;
        }
        throw new Error('Unexpected handshake frame after acceptance');
      }
      const usesLivenessReserve =
        this.#requests.size === RUNTIME_HOST_MAX_IN_FLIGHT_DOMAIN_REQUESTS &&
        (frame.operation === 'host.status' || this.#inFlightStatusRequests > 0);
      if (this.#requests.has(frame.requestId)) {
        this.#fail(new Error('Runtime Host Client reused an active request id'));
        return;
      }
      if (
        this.#requests.size >= RUNTIME_HOST_MAX_IN_FLIGHT_DOMAIN_REQUESTS &&
        !usesLivenessReserve
      ) {
        this.#fail(new Error('Runtime Host Client exceeded the in-flight request limit'));
        return;
      }
      this.#dispatch(frame);
    }
  }

  #dispatch(frame: RequestFrame): void {
    if (frame.operation === 'host.status') this.#inFlightStatusRequests += 1;
    const handling =
      frame.operation === 'session.transcript.page'
        ? this.#transcriptPageTail.then(() => this.#handleRequest(frame))
        : this.#handleRequest(frame);
    const task = handling
      .catch((error: unknown) => this.#fail(error))
      .finally(() => {
        if (this.#requests.get(frame.requestId) === task) {
          this.#requests.delete(frame.requestId);
          if (frame.operation === 'host.status') this.#inFlightStatusRequests -= 1;
        }
      });
    this.#requests.set(frame.requestId, task);
    if (frame.operation === 'session.transcript.page') {
      this.#transcriptPageTail = task.catch(() => undefined);
    }
  }

  async #handleRequest(frame: RequestFrame): Promise<void> {
    if (this.#closed) return;
    if (!authorizeRuntimeHostOperation(this.#options.connection.authority, frame)) {
      if (this.#closed) return;
      await this.#writer.enqueue(
        operationFailureResponse(frame, 'unauthorized', 'Runtime Host operation is not authorized'),
      ).flushed;
      return;
    }
    const admission = await this.#options.beginOperation(frame);
    if (typeof admission === 'string') {
      if (this.#closed) return;
      await this.#writer.enqueue(
        operationFailureResponse(
          frame,
          admission,
          admission === 'host_draining' ? 'Runtime Host is draining' : 'Runtime Host is not ready',
        ),
      ).flushed;
      return;
    }

    try {
      if (this.#closed) return;
      this.#ensureClientCapabilities();
      const continuity =
        frame.operation === 'subscription.open' ||
        frame.operation === 'subscription.close' ||
        frame.operation === 'session.transcript.page'
          ? this.#ensureContinuity()
          : undefined;
      const response = await dispatchOperation(frame, this.#options.resolveHandlers(), {
        ...this.#options.connection,
        inputClosedSignal: this.#inputClosedAbort.signal,
        principal: this.#options.connection.authority.principalId,
        principalKind: this.#options.connection.authority.principalKind,
        ...(this.#options.connection.authority.credentialId
          ? { credentialId: this.#options.connection.authority.credentialId }
          : {}),
        ...(this.#options.connection.authority.clientInstanceId
          ? { credentialClientInstanceId: this.#options.connection.authority.clientInstanceId }
          : {}),
        acquireResidency: () => admission.acquireResidency(),
      });
      admission.seal();
      const receipt = this.#writer.enqueue(response);
      const openedSubscriptionId =
        response.ok && response.operation === 'subscription.open'
          ? response.result.subscriptionId
          : undefined;
      try {
        await receipt.flushed;
        // Subscriber-local queues retain pre-activation events. Expose them
        // only after the open result leaves the connection-wide writer, or a
        // restore fan-out can make legal responses and first frames overflow it.
        if (openedSubscriptionId) continuity?.activate(openedSubscriptionId);
      } catch (error) {
        if (openedSubscriptionId) continuity?.abort(openedSubscriptionId);
        throw error;
      }
    } finally {
      admission.finish();
    }
  }

  #ensureContinuity(): SessionContinuityConnection | undefined {
    if (this.#closed || this.#inputClosedAbort.signal.aborted) return;
    const service = this.#options.resolveContinuity();
    if (!service) return;
    if (this.#continuityService && this.#continuityService !== service) {
      throw new Error('Runtime Host continuity service changed within one connection');
    }
    if (!this.#continuity) {
      this.#continuityService = service;
      this.#continuity = service.attachConnection(this.#options.connection.connectionId, {
        send: (frame) => {
          try {
            return this.#writer.enqueue(frame).flushed;
          } catch (error) {
            return Promise.reject(error);
          }
        },
      });
    }
    return this.#continuity;
  }

  #detachContinuity(): void {
    this.#continuity?.close();
    this.#continuity = undefined;
    this.#continuityService = undefined;
  }

  #ensureClientCapabilities(): ClientCapabilityConnection | undefined {
    if (this.#closed || this.#inputClosedAbort.signal.aborted) return;
    // Guests observe and submit approval requests; they cannot provide Client
    // Capabilities or directly admit a Turn. In particular, their pending and
    // finalized connections may overlap while the credential becomes bound to
    // the Client. Neither connection owns a capability-provider registration.
    if (this.#options.connection.authority.principalKind === 'session_guest') return;
    const service = this.#options.resolveClientCapabilities?.();
    if (!service) return;
    if (this.#clientCapabilityService && this.#clientCapabilityService !== service) {
      throw new Error('Runtime Host Client Capability service changed within one connection');
    }
    if (!this.#clientCapabilities) {
      this.#clientCapabilityService = service;
      this.#clientCapabilities = service.attachConnection(
        {
          connectionId: this.#options.connection.connectionId,
          principalId: this.#options.connection.authority.principalId,
          clientInstanceId: this.#options.connection.clientInstanceId,
          ...(this.#options.connection.authority.clientInstanceId
            ? {
                credentialBoundClientInstanceId:
                  this.#options.connection.authority.clientInstanceId,
              }
            : {}),
          principalKind: this.#options.connection.authority.principalKind,
          ...(this.#options.connection.authority.capabilityOwner
            ? { capabilityOwner: this.#options.connection.authority.capabilityOwner }
            : {}),
        },
        {
          send: (frame) => {
            try {
              return this.#writer.enqueue(frame).flushed;
            } catch (error) {
              return Promise.reject(error);
            }
          },
        },
      );
    }
    return this.#clientCapabilities;
  }

  #detachClientCapabilities(): void {
    const connection = this.#clientCapabilities;
    this.#clientCapabilities = undefined;
    this.#clientCapabilityService = undefined;
    if (!connection || this.#clientCapabilityCloseTask) return;
    this.#clientCapabilityCloseTask = Promise.resolve().then(() => connection.close());
    void this.#clientCapabilityCloseTask.catch(() => undefined);
  }

  attachGlobalChanges(): void {
    if (this.#closed || this.#inputClosedAbort.signal.aborted) return;
    const service = this.#options.resolveHostChanges?.();
    if (!service || this.#hostChanges) return;
    const sharedSessionId =
      this.#options.connection.authority.principalKind === 'session_guest' &&
      hasRuntimeHostOperationGrant(this.#options.connection.authority, 'session.shared.query')
        ? this.#options.resolveSharedSessionId?.()
        : undefined;
    const sessionCatalog: HostChangeSubscriptionMask['sessionCatalog'] =
      sharedSessionId !== undefined
        ? {
            sessionId: sharedSessionId,
            principalId: this.#options.connection.authority.principalId,
          }
        : hasRuntimeHostOperationGrant(this.#options.connection.authority, 'session.catalog.query')
          ? true
          : undefined;
    this.#hostChanges = service.attachConnection(
      this.#options.connection.connectionId,
      {
        configuration: hasRuntimeHostOperationGrant(
          this.#options.connection.authority,
          'runtime.policy.query',
        ),
        connectionCatalog: hasRuntimeHostOperationGrant(
          this.#options.connection.authority,
          'connection.catalog.query',
        ),
        projectCatalog: hasRuntimeHostOperationGrant(
          this.#options.connection.authority,
          'project.catalog.query',
        ),
        sessionCatalog,
        scheduledTask: hasRuntimeHostOperationGrant(
          this.#options.connection.authority,
          'scheduled-task.query',
        ),
      },
      {
        send: (frame) => {
          try {
            return this.#writer.enqueue(frame).flushed;
          } catch (error) {
            return Promise.reject(error);
          }
        },
      },
    );
  }

  #detachHostChanges(): void {
    this.#hostChanges?.close();
    this.#hostChanges = undefined;
  }

  #fail(error: unknown): void {
    if (this.#closed) return;
    try {
      this.#onDiagnostic(
        `[runtime-host] connection session failed: ${boundedFailureDiagnostic(error)}`,
      );
    } catch {
      // Diagnostics cannot keep an invalid connection alive.
    }
    this.#teardown();
  }

  #teardown(): void {
    if (this.#closed) return;
    this.#closed = true;
    this.#inputClosedAbort.abort();
    this.#detachContinuity();
    this.#detachClientCapabilities();
    this.#detachHostChanges();
    this.#writer.close();
    this.#options.transport.abort();
    this.#options.onTeardown();
  }
}

function isReadEof(error: unknown): boolean {
  return error instanceof RuntimeHostTransportError && error.code === 'read_eof';
}