stack/packages/agent/src/durable-session.ts

export interface DurableSessionIdentity {
  /** Provisioned roster identity; isolates durable state between named agents. */
  agentName: string;
  sessionId: string;
  tenantId: string;
  ownerId: string;
  providerId: string;
  runtimeSessionId: string;
}

export type DurableInboxStatus = 'pending' | 'processing' | 'processed';
export type DurableOutboxStatus = 'pending' | 'processing' | 'delivered';

export interface DurableInboxInput {
  sessionId: string;
  idempotencyKey: string;
  correlationId: string;
  content: string;
}

export interface DurableInboxEntry extends DurableInboxInput {
  status: DurableInboxStatus;
}

export interface DurableOutboxInput {
  sessionId: string;
  idempotencyKey: string;
  correlationId: string;
  channelId: string;
  kind: string;
  content: string;
}

export interface DurableOutboxEntry extends DurableOutboxInput {
  status: DurableOutboxStatus;
}

export interface DurableCheckpointInput {
  sessionId: string;
  checkpointId: string;
  cursor: string;
  summary: string;
  compactionEpoch: number;
}

export interface DurableCheckpoint extends DurableCheckpointInput {}

export interface DurableHandoffInput {
  sessionId: string;
  handoffId: string;
  destination: string;
  correlationId: string;
  checkpointId: string;
  status: 'pending' | 'accepted';
}

export interface DurableHandoff extends DurableHandoffInput {}

export interface DurableSessionSnapshot {
  identity: DurableSessionIdentity;
  inbox: DurableInboxEntry[];
  outbox: DurableOutboxEntry[];
  checkpoint?: DurableCheckpoint;
  handoffs: DurableHandoff[];
}

export interface DurableHandoffRecovery {
  identity: DurableSessionIdentity;
  checkpoint: DurableCheckpoint;
  handoff: DurableHandoff;
}

export interface DurableEnqueueResult<TStatus extends string> {
  accepted: boolean;
  status: TStatus;
}

/**
 * A durable-state port. Implementations must atomically claim and complete work
 * records. Recovery may requeue interrupted inbox work, but never an
 * externally visible outbox effect: an ambiguous provider result remains
 * claimed until a separately authorized reconciliation proves it safe.
 */
export interface DurableSessionStore {
  create(identity: DurableSessionIdentity): Promise<void>;
  snapshot(sessionId: string): Promise<DurableSessionSnapshot | null>;
  enqueueInbox(input: DurableInboxInput): Promise<DurableEnqueueResult<DurableInboxStatus>>;
  claimInbox(sessionId: string): Promise<DurableInboxEntry | null>;
  completeInbox(sessionId: string, idempotencyKey: string): Promise<void>;
  releaseInbox(sessionId: string, idempotencyKey: string): Promise<void>;
  enqueueOutbox(input: DurableOutboxInput): Promise<DurableEnqueueResult<DurableOutboxStatus>>;
  claimOutbox(sessionId: string): Promise<DurableOutboxEntry | null>;
  claimOutboxByKey(sessionId: string, idempotencyKey: string): Promise<DurableOutboxEntry | null>;
  completeOutbox(sessionId: string, idempotencyKey: string): Promise<void>;
  releaseOutbox(sessionId: string, idempotencyKey: string): Promise<void>;
  checkpoint(input: DurableCheckpointInput): Promise<void>;
  findCheckpoint(sessionId: string, checkpointId: string): Promise<DurableCheckpoint | null>;
  handoff(input: DurableHandoffInput): Promise<void>;
  findHandoff(handoffId: string): Promise<DurableHandoff | null>;
  requeueInFlight(sessionId: string): Promise<void>;
}

export class DurableSessionNotFoundError extends Error {
  constructor(sessionId: string) {
    super(`Durable session not found: ${sessionId}`);
    this.name = 'DurableSessionNotFoundError';
  }
}

export class DurableSessionCoordinator {
  constructor(private readonly store: DurableSessionStore) {}

  async create(identity: DurableSessionIdentity): Promise<void> {
    this.assertIdentity(identity);
    await this.store.create(identity);
  }

  async receive(input: DurableInboxInput): Promise<DurableEnqueueResult<DurableInboxStatus>> {
    this.assertRecord(input.sessionId, input.idempotencyKey, input.correlationId, input.content);
    return this.store.enqueueInbox(input);
  }

  async enqueueOutbox(
    input: DurableOutboxInput,
  ): Promise<DurableEnqueueResult<DurableOutboxStatus>> {
    this.assertRecord(
      input.sessionId,
      input.idempotencyKey,
      input.correlationId,
      input.channelId,
      input.content,
    );
    if (input.kind.trim().length === 0) throw new Error('Durable outbox kind is required');
    return this.store.enqueueOutbox(input);
  }

  async checkpoint(input: DurableCheckpointInput): Promise<void> {
    this.assertRecord(input.sessionId, input.checkpointId, input.cursor, input.summary);
    if (!Number.isInteger(input.compactionEpoch) || input.compactionEpoch < 0) {
      throw new Error('Durable checkpoint compaction epoch must be a non-negative integer');
    }
    await this.store.checkpoint(input);
  }

  async handoff(input: DurableHandoffInput): Promise<void> {
    this.assertRecord(input.sessionId, input.handoffId, input.destination, input.correlationId);
    if (input.checkpointId.trim().length === 0)
      throw new Error('Durable handoff checkpoint is required');
    await this.store.handoff(input);
  }

  /** Read durable state without changing claim status; safe during normal operation. */
  async snapshot(sessionId: string): Promise<DurableSessionSnapshot> {
    const snapshot = await this.store.snapshot(sessionId);
    if (!snapshot) throw new DurableSessionNotFoundError(sessionId);
    return snapshot;
  }

  /** Requeue interrupted inbox work during recovery; preserve ambiguous outbox claims. */
  async recover(sessionId: string): Promise<DurableSessionSnapshot> {
    await this.store.requeueInFlight(sessionId);
    return this.snapshot(sessionId);
  }

  async drainInbox(
    sessionId: string,
    handler: (entry: DurableInboxEntry) => Promise<void>,
  ): Promise<void> {
    for (;;) {
      const entry = await this.store.claimInbox(sessionId);
      if (!entry) return;
      try {
        await handler(entry);
      } catch (error: unknown) {
        await this.store.releaseInbox(sessionId, entry.idempotencyKey);
        throw error;
      }
      // If this write fails after the handler succeeded, leave the record
      // processing. A recovery path can retry it with its stable idempotency key.
      await this.store.completeInbox(sessionId, entry.idempotencyKey);
    }
  }

  async dispatchOutbox(
    sessionId: string,
    dispatcher: (entry: DurableOutboxEntry) => Promise<void>,
  ): Promise<void> {
    for (;;) {
      const entry = await this.store.claimOutbox(sessionId);
      if (!entry) return;
      // A provider failure can be ambiguous: it may occur after the receiver
      // accepted the idempotency key. Preserve the claim for reconciliation.
      await dispatcher(entry);
      // Do not requeue an effect after it has been applied but before its
      // terminal state could be persisted; recovery preserves the claim.
      await this.store.completeOutbox(sessionId, entry.idempotencyKey);
    }
  }

  async dispatchOutboxEntry(
    sessionId: string,
    idempotencyKey: string,
    dispatcher: (entry: DurableOutboxEntry) => Promise<void>,
  ): Promise<void> {
    const entry = await this.store.claimOutboxByKey(sessionId, idempotencyKey);
    if (!entry) return;
    // A provider failure can be ambiguous, so this stays processing until
    // separately authorized reconciliation proves it safe to resolve.
    await dispatcher(entry);
    await this.store.completeOutbox(sessionId, entry.idempotencyKey);
  }

  async resumeHandoff(handoffId: string): Promise<DurableHandoffRecovery> {
    const handoff = await this.store.findHandoff(handoffId);
    if (!handoff) throw new Error(`Durable handoff not found: ${handoffId}`);
    const snapshot = await this.snapshot(handoff.sessionId);
    const checkpoint = await this.store.findCheckpoint(handoff.sessionId, handoff.checkpointId);
    if (!checkpoint) {
      throw new Error(`Durable handoff checkpoint is unavailable: ${handoff.checkpointId}`);
    }
    return { identity: snapshot.identity, checkpoint, handoff };
  }

  private assertIdentity(identity: DurableSessionIdentity): void {
    this.assertRecord(
      identity.agentName,
      identity.sessionId,
      identity.tenantId,
      identity.ownerId,
      identity.providerId,
    );
    if (identity.runtimeSessionId.trim().length === 0) {
      throw new Error('Durable runtime session identity is required');
    }
  }

  private assertRecord(...values: string[]): void {
    if (values.some((value: string): boolean => value.trim().length === 0)) {
      throw new Error('Durable session records require non-empty fields');
    }
  }
}

interface InMemorySessionState {
  identity: DurableSessionIdentity;
  inbox: Map<string, DurableInboxEntry>;
  outbox: Map<string, DurableOutboxEntry>;
  checkpoints: Map<string, DurableCheckpoint>;
  handoffs: Map<string, DurableHandoff>;
}

/** Reference store for deterministic domain tests; production uses the gateway DB adapter. */
export class InMemoryDurableSessionStore implements DurableSessionStore {
  private readonly sessions = new Map<string, InMemorySessionState>();

  async create(identity: DurableSessionIdentity): Promise<void> {
    const existing = this.sessions.get(identity.sessionId);
    if (existing) {
      if (!sameEnrollmentScope(existing.identity, identity)) {
        throw new Error(`Durable session identity conflict: ${identity.sessionId}`);
      }
      existing.identity.providerId = identity.providerId;
      existing.identity.runtimeSessionId = identity.runtimeSessionId;
      return;
    }
    this.sessions.set(identity.sessionId, {
      identity: copyIdentity(identity),
      inbox: new Map(),
      outbox: new Map(),
      checkpoints: new Map(),
      handoffs: new Map(),
    });
  }

  async snapshot(sessionId: string): Promise<DurableSessionSnapshot | null> {
    const state = this.sessions.get(sessionId);
    if (!state) return null;
    const checkpoint = latestCheckpoint(state.checkpoints);
    return {
      identity: copyIdentity(state.identity),
      inbox: [...state.inbox.values()].map(copyInbox),
      outbox: [...state.outbox.values()].map(copyOutbox),
      ...(checkpoint ? { checkpoint: copyCheckpoint(checkpoint) } : {}),
      handoffs: [...state.handoffs.values()].map(copyHandoff),
    };
  }

  async enqueueInbox(input: DurableInboxInput): Promise<DurableEnqueueResult<DurableInboxStatus>> {
    const state = this.require(input.sessionId);
    const existing = state.inbox.get(input.idempotencyKey);
    if (existing) {
      if (!sameInbox(existing, input)) {
        throw new Error(`Durable inbox idempotency conflict: ${input.idempotencyKey}`);
      }
      return { accepted: false, status: existing.status };
    }
    state.inbox.set(input.idempotencyKey, { ...input, status: 'pending' });
    return { accepted: true, status: 'pending' };
  }

  async claimInbox(sessionId: string): Promise<DurableInboxEntry | null> {
    const state = this.require(sessionId);
    const entry = [...state.inbox.values()].find(
      (candidate: DurableInboxEntry): boolean => candidate.status === 'pending',
    );
    if (!entry) return null;
    entry.status = 'processing';
    return copyInbox(entry);
  }

  async completeInbox(sessionId: string, idempotencyKey: string): Promise<void> {
    this.requireEntry(this.require(sessionId).inbox, idempotencyKey, 'inbox').status = 'processed';
  }

  async releaseInbox(sessionId: string, idempotencyKey: string): Promise<void> {
    this.requireEntry(this.require(sessionId).inbox, idempotencyKey, 'inbox').status = 'pending';
  }

  async enqueueOutbox(
    input: DurableOutboxInput,
  ): Promise<DurableEnqueueResult<DurableOutboxStatus>> {
    const state = this.require(input.sessionId);
    const existing = state.outbox.get(input.idempotencyKey);
    if (existing) {
      if (!sameOutbox(existing, input)) {
        throw new Error(`Durable outbox idempotency conflict: ${input.idempotencyKey}`);
      }
      return { accepted: false, status: existing.status };
    }
    state.outbox.set(input.idempotencyKey, { ...input, status: 'pending' });
    return { accepted: true, status: 'pending' };
  }

  async claimOutbox(sessionId: string): Promise<DurableOutboxEntry | null> {
    const state = this.require(sessionId);
    const entry = [...state.outbox.values()].find(
      (candidate: DurableOutboxEntry): boolean => candidate.status === 'pending',
    );
    if (!entry) return null;
    entry.status = 'processing';
    return copyOutbox(entry);
  }

  async claimOutboxByKey(
    sessionId: string,
    idempotencyKey: string,
  ): Promise<DurableOutboxEntry | null> {
    const entry = this.require(sessionId).outbox.get(idempotencyKey);
    if (!entry || entry.status !== 'pending') return null;
    entry.status = 'processing';
    return copyOutbox(entry);
  }

  async completeOutbox(sessionId: string, idempotencyKey: string): Promise<void> {
    this.requireEntry(this.require(sessionId).outbox, idempotencyKey, 'outbox').status =
      'delivered';
  }

  async releaseOutbox(sessionId: string, idempotencyKey: string): Promise<void> {
    this.requireEntry(this.require(sessionId).outbox, idempotencyKey, 'outbox').status = 'pending';
  }

  async checkpoint(input: DurableCheckpointInput): Promise<void> {
    const checkpoints = this.require(input.sessionId).checkpoints;
    const existing = checkpoints.get(input.checkpointId);
    if (existing && !sameCheckpoint(existing, input)) {
      throw new Error(`Durable checkpoint identity conflict: ${input.checkpointId}`);
    }
    if (!existing) checkpoints.set(input.checkpointId, { ...input });
  }

  async findCheckpoint(sessionId: string, checkpointId: string): Promise<DurableCheckpoint | null> {
    const checkpoint = this.require(sessionId).checkpoints.get(checkpointId);
    return checkpoint ? copyCheckpoint(checkpoint) : null;
  }

  async handoff(input: DurableHandoffInput): Promise<void> {
    const state = this.require(input.sessionId);
    if (!state.checkpoints.has(input.checkpointId)) {
      throw new Error(`Durable handoff checkpoint is unavailable: ${input.checkpointId}`);
    }
    const existing = state.handoffs.get(input.handoffId);
    if (existing && !sameHandoff(existing, input)) {
      throw new Error(`Durable handoff identity conflict: ${input.handoffId}`);
    }
    if (!existing) state.handoffs.set(input.handoffId, { ...input });
  }

  async findHandoff(handoffId: string): Promise<DurableHandoff | null> {
    for (const state of this.sessions.values()) {
      const handoff = state.handoffs.get(handoffId);
      if (handoff) return copyHandoff(handoff);
    }
    return null;
  }

  async requeueInFlight(sessionId: string): Promise<void> {
    const state = this.require(sessionId);
    for (const entry of state.inbox.values()) {
      if (entry.status === 'processing') entry.status = 'pending';
    }
  }

  private require(sessionId: string): InMemorySessionState {
    const state = this.sessions.get(sessionId);
    if (!state) throw new DurableSessionNotFoundError(sessionId);
    return state;
  }

  private requireEntry<T extends { status: string }>(
    records: Map<string, T>,
    idempotencyKey: string,
    kind: string,
  ): T {
    const entry = records.get(idempotencyKey);
    if (!entry) throw new Error(`Durable ${kind} entry not found: ${idempotencyKey}`);
    return entry;
  }
}

function sameEnrollmentScope(left: DurableSessionIdentity, right: DurableSessionIdentity): boolean {
  return (
    left.agentName === right.agentName &&
    left.sessionId === right.sessionId &&
    left.tenantId === right.tenantId &&
    left.ownerId === right.ownerId
  );
}

function sameInbox(left: DurableInboxEntry, right: DurableInboxInput): boolean {
  return (
    left.sessionId === right.sessionId &&
    left.idempotencyKey === right.idempotencyKey &&
    left.correlationId === right.correlationId &&
    left.content === right.content
  );
}

function sameOutbox(left: DurableOutboxEntry, right: DurableOutboxInput): boolean {
  return (
    left.sessionId === right.sessionId &&
    left.idempotencyKey === right.idempotencyKey &&
    left.correlationId === right.correlationId &&
    left.channelId === right.channelId &&
    left.kind === right.kind &&
    left.content === right.content
  );
}

function sameCheckpoint(left: DurableCheckpoint, right: DurableCheckpointInput): boolean {
  return (
    left.sessionId === right.sessionId &&
    left.checkpointId === right.checkpointId &&
    left.cursor === right.cursor &&
    left.summary === right.summary &&
    left.compactionEpoch === right.compactionEpoch
  );
}

function latestCheckpoint(
  checkpoints: Map<string, DurableCheckpoint>,
): DurableCheckpoint | undefined {
  return [...checkpoints.values()].sort(
    (left: DurableCheckpoint, right: DurableCheckpoint): number =>
      right.compactionEpoch - left.compactionEpoch,
  )[0];
}

function sameHandoff(left: DurableHandoff, right: DurableHandoffInput): boolean {
  return (
    left.sessionId === right.sessionId &&
    left.handoffId === right.handoffId &&
    left.destination === right.destination &&
    left.correlationId === right.correlationId &&
    left.checkpointId === right.checkpointId &&
    left.status === right.status
  );
}

function copyIdentity(identity: DurableSessionIdentity): DurableSessionIdentity {
  return { ...identity };
}

function copyInbox(entry: DurableInboxEntry): DurableInboxEntry {
  return { ...entry };
}

function copyOutbox(entry: DurableOutboxEntry): DurableOutboxEntry {
  return { ...entry };
}

function copyCheckpoint(checkpoint: DurableCheckpoint): DurableCheckpoint {
  return { ...checkpoint };
}

function copyHandoff(handoff: DurableHandoff): DurableHandoff {
  return { ...handoff };
}