import { Cause, Deferred, Effect, Exit, Option, Schema } from "effect";import type * as Tracer from "effect/Tracer";import type { Connection, ConnectionContext } from "agents";import { McpAgent } from "agents/mcp";import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js"; import { RequestOrgSlug, RequestWebOrigin } from "@executor-js/api/server";import {  formatPausedExecution,  type ExecutionEngine,  type ExecutionResult,  type PausedExecutionDeadline,  type ResumeResponse,} from "@executor-js/execution";import {  PAUSED_APPROVAL_TIMEOUT_MS,  formatMcpExecutionOutcome,  type PausedExecutionHooks,  type ResumeFallbackOutcome,} from "@executor-js/host-mcp/tool-server";import { defaultMcpResource, type McpResource } from "@executor-js/host-mcp"; import type { IncomingPropagationHeaders, McpElicitationMode } from "./do-headers";import { classifyDurableObjectError, type DurableObjectFailure } from "./durable-object-errors";import type {  McpExecutionOwnerDirectory,  McpExecutionOwnerRecord,  McpExecutionOwnerRoute,} from "./execution-owner-directory";import {  MAX_PAUSED_SESSION_IDLE_MS,  SESSION_TIMEOUT_MS,  decideSessionAlarm,  pausedLeaseExtensionLog,  runningLeaseExtensionLog,} from "./session-alarm-policy";import {  acquireResidentRuntime,  currentInFlightColdBuildCount,  currentResidentRuntimeCount,  markEvictionRequested,  pickEvictionCandidate,  registerResidentSession,  releaseColdBuildSlot,  releaseResidentRuntime,  releaseResidentSession,  reserveColdBuildSlot,  type ResidentSessionEntry,  residencyAttributes,  RESIDENT_RUNTIME_SOFT_CAP,  touchResidentSession,} from "./session-runtime-residency"; export type IncomingTraceHeaders = IncomingPropagationHeaders; export interface McpSessionInit {  readonly organizationId: string;  /** The organization's display name, as the worker resolved it while   *  authorizing this very request. Carried so the session DO never has to   *  re-read a row the request already loaded. Absent when the auth plane could   *  not name the org, in which case the host resolves it itself. */  readonly organizationName?: string;  /** The organization's URL slug, from the same resolved record. */  readonly organizationSlug?: string;  readonly userId: string;  readonly elicitationMode: McpElicitationMode;  /** Whether this session serves artifacts, read off `?artifacts=` at connect   *  time. Absent means the default (enabled). */  readonly artifactsEnabled?: boolean;  /** Whether this session serves the per-integration `search_<integration>`   *  tools, read off `?search_tools=` at connect time. Absent means the   *  default (disabled). */  readonly searchToolsEnabled?: boolean;  /** The MCP resource the session was minted against (`/mcp` default vs a   *  `/mcp/toolkits/<slug>` toolkit), so the tool catalog is scoped to it. */  readonly resource: McpResource;  readonly webOrigin?: string;} export interface McpSessionProps extends Record<string, unknown> {  readonly session: McpSessionInit;  readonly propagation?: IncomingTraceHeaders;} export type McpApprovalOwner = {  readonly accountId: string;  readonly organizationId: string;}; type McpSessionApprovalErrorResult =  | { readonly status: "not_found" }  | { readonly status: "forbidden" }; type PendingApprovalLease = {  readonly disposeKeepAlive: () => void;  timeout: ReturnType<typeof setTimeout> | null;  expiring: boolean;}; export type McpSessionApprovalResult =  | {      readonly status: "ok";      readonly text: string;      readonly structured: Record<string, unknown>;    }  | McpSessionApprovalErrorResult; export type McpSessionResumeApprovalResult =  | {      readonly status: "ok";      readonly executionStatus: "completed" | "paused";      readonly text: string;      readonly structured: Record<string, unknown>;      readonly isError?: boolean;    }  | McpSessionApprovalErrorResult; export type McpSessionModelResumeResult = ResumeFallbackOutcome; export interface SessionDbHandle {  readonly end: () => Promise<void> | void;} export interface SessionMeta {  readonly organizationId: string;  readonly organizationName: string;  /** The org's URL slug, when the host's `resolveSessionMeta` carried one.   * Pins browser-handoff URLs to the right org's console. */  readonly organizationSlug?: string;  readonly userId: string;  readonly elicitationMode?: McpElicitationMode;  /** Whether the session serves artifacts (carried from {@link McpSessionInit}).   *  Absent — including for sessions persisted before the flag existed — means   *  the default (enabled). */  readonly artifactsEnabled?: boolean;  /** Whether the session serves the per-integration search tools (carried from   *  {@link McpSessionInit}). Absent — including for sessions persisted before   *  the flag existed — means the default (disabled). */  readonly searchToolsEnabled?: boolean;  /** The MCP resource the session serves (carried from {@link McpSessionInit});   *  `buildMcpServer` scopes the tool catalog to it. */  readonly resource: McpResource;  readonly webOrigin?: string;  /**   * Whether this session's client advertised MCP-Apps support at `initialize`.   *   * Capabilities are negotiated once, into the server instance's memory. When   * the DO is evicted (deploy, idle) and a later request cold-restores it, the   * rebuilt server never sees an `initialize` — so without persisting this, an   * apps-capable client silently drops to artifact deep links mid-conversation.   * Absent — including for sessions persisted before this field existed — means   * unknown, which behaves as disabled until the next `initialize`.   */  readonly appsEnabled?: boolean;} export interface BuiltMcpServer {  readonly mcpServer: McpServer;  readonly engine: ExecutionEngine<Cause.YieldableError>;} export interface BrowserApprovalStore {  readonly takeResponse: (executionId: string) => Effect.Effect<ResumeResponse | null>;  readonly waitForResponse: (executionId: string) => Effect.Effect<ResumeResponse | null>;} const SESSION_META_KEY = "session-meta";const LAST_ACTIVITY_KEY = "last-activity-ms";const PARTYSERVER_NAME_KEY = "__ps_name";/** * Stand-in session id for a log line written when the DO's name could not be * resolved. A named placeholder keeps the field present and the log shape * stable, and is countable when it happens. */const UNRESOLVED_SESSION_ID = "unresolved";/** The agents SDK's durable "condemned" marker (`_cf_scheduleDestroy`). */const AGENTS_DESTROY_PENDING_KEY = "cf_agents_destroy_pending";const MCP_HTTP_METHOD_HEADER = "cf-mcp-method";const MCP_MESSAGE_HEADER = "cf-mcp-message";const MODEL_RESUME_FORWARD_TIMEOUT_MS = 10_000;const MCP_STREAM_REQS_KEY_PREFIX = "__mcp_stream_reqs__:";const approvalResponseKey = (executionId: string) => `approval-response:${executionId}`; type JsonRpcRequestId = string | number;const JsonRpcRequestWithId = Schema.Struct({  id: Schema.Union([Schema.String, Schema.Number]),  method: Schema.String,});const JsonRpcPostPayload = Schema.fromJsonString(Schema.Unknown);const decodeJsonRpcPostPayload = Schema.decodeUnknownOption(JsonRpcPostPayload);const decodeJsonRpcRequestWithId = Schema.decodeUnknownOption(JsonRpcRequestWithId); const resumeApprovalResult = (  executionId: string,  response: ResumeResponse,): Extract<McpSessionResumeApprovalResult, { readonly status: "ok" }> => {  const textByAction = {    accept: "I've approved it",    decline: "I've denied it",    cancel: "I've canceled it",  } satisfies Record<ResumeResponse["action"], string>;  const statusByAction = {    accept: "approved",    decline: "denied",    cancel: "canceled",  } satisfies Record<ResumeResponse["action"], string>;   return {    status: "ok",    executionStatus: "completed",    text: textByAction[response.action],    structured: { status: statusByAction[response.action], executionId },    isError: false,  };}; const isSessionProps = (props: unknown): props is McpSessionProps =>  typeof props === "object" &&  props !== null &&  "session" in props &&  typeof (props as { readonly session?: unknown }).session === "object" &&  (props as { readonly session?: unknown }).session !== null; const readActivePostRequestIds = (request: Request): readonly JsonRpcRequestId[] => {  if (request.headers.get(MCP_HTTP_METHOD_HEADER) !== "POST") return [];  const encoded = request.headers.get(MCP_MESSAGE_HEADER);  if (!encoded) return [];  const decoded = Effect.runSyncExit(    Effect.try({      try: () => atob(encoded),      catch: () => "invalid_base64" as const,    }),  );  if (Exit.isFailure(decoded)) {    console.warn(      JSON.stringify({        event: "mcp_active_post_response_wait_parse_failed",        reason: "invalid_base64",      }),    );    return [];  }  const parsed = decodeJsonRpcPostPayload(decoded.value);  if (Option.isNone(parsed)) {    console.warn(      JSON.stringify({        event: "mcp_active_post_response_wait_parse_failed",        reason: "invalid_json",      }),    );    return [];  }  const messages = Array.isArray(parsed.value) ? parsed.value : [parsed.value];  const requestIds: JsonRpcRequestId[] = [];  for (const message of messages) {    const decoded = decodeJsonRpcRequestWithId(message);    if (Option.isSome(decoded)) requestIds.push(decoded.value.id);  }  return requestIds;}; export abstract class McpAgentSessionDOBase<  Env extends Cloudflare.Env = Cloudflare.Env,  TDbHandle extends SessionDbHandle = SessionDbHandle,> extends McpAgent<Env, unknown, McpSessionProps> {  server!: McpServer;  private engine: ExecutionEngine<Cause.YieldableError> | null = null;  private dbHandle: TDbHandle | null = null;  private sessionMeta: SessionMeta | null = null;  private initialized = false;  /** Whether this instance is currently counted in the isolate's residency   *  gauge. Tracked separately from `engine` because `closeRuntime` runs on   *  paths where nothing was ever built, and it must not decrement then. */  private countedAsResident = false;  /** Whether THIS init's cold-build admission is currently holding a reserved   *  slot in the isolate-wide in-flight counter. Mirrors `countedAsResident`'s   *  guard discipline: gates {@link releaseColdBuildSlotIfReserved} so it   *  releases exactly once per reservation, from whichever of its two callers   *  (success or failure/interrupt) gets there first. */  private reservedColdBuildSlot = false;  /**   * The `AbortController` backing the currently in-flight `init` call's   * `Effect.runPromise`, so its root fiber can be interrupted from outside   * the Promise it returns. Nothing in production ever calls `.abort()` on   * it today — it exists so a unit test can interrupt `init` deterministically   * at a chosen suspension point (e.g. mid `openSessionDbHandle`) and assert   * the cold-build reservation is still released, rather than only ever being   * exercised through failure/defect paths.   */  private initAbortController: AbortController | null = null;  /** Same purpose as {@link initAbortController}, for `disposeIdleRuntime`'s   *  root effect (`closeRuntime` plus its alarm/activity bookkeeping) — lets   *  a unit test interrupt a disposal deterministically to confirm the   *  uninterruptible teardown in `closeRuntime` actually finishes rather than   *  being cut short. */  private disposeAbortController: AbortController | null = null;  /**   * The in-progress runtime disposal, if one is running right now — from   * whichever of `closeRuntime`'s callers (the idle alarm, a cap eviction   * request, `cleanup`) got there first. Set at the very start of   * `closeRuntime`'s body, before its first async close, and cleared in a   * finally once that disposal settles.   *   * Exists to close two races at once: `init` awaits this before deciding   * whether to rebuild, so a request that lands mid-teardown never races the   * closes it is waiting on; and a second `closeRuntime` call that lands   * while this is set (the idle alarm and a cap eviction request landing   * together) waits on the SAME promise instead of running the teardown a   * second time.   */  private disposingRuntime: Promise<void> | null = null;  private onStartPromise: Promise<void> | null = null;  private lastActivityMs = 0;  private resolvedSessionName: string | undefined = undefined;  private approvalResponses = new Map<string, ResumeResponse>();  private approvalWaiters = new Map<string, Deferred.Deferred<ResumeResponse>>();  private pendingApprovalLeases = new Map<string, PendingApprovalLease>();   protected abstract openSessionDb(): TDbHandle | Promise<TDbHandle>;   /**   * Build the session's {@link SessionMeta} for this init.   *   * `storedMeta` is what this DO already persisted for the SAME organization on   * an earlier init, or `null`. It is offered first so a host never has to   * re-resolve an org identity it is already holding — on cloud that resolution   * is a fresh Postgres connection, and a cold restore used to die on it. Every   * field the CONNECT carries (resource, elicitation mode, capability flags)   * still comes from `token`; only the org identity may be reused.   */  protected abstract resolveSessionMeta(    token: McpSessionInit,    storedMeta: SessionMeta | null,  ): Effect.Effect<SessionMeta>;   protected abstract buildMcpServer(    sessionMeta: SessionMeta,    dbHandle: TDbHandle,  ): Effect.Effect<BuiltMcpServer>;   /**   * Cheap, count-only proxies for what the runtime `buildMcpServer` just built   * holds — e.g. connected-integration counts — attributed alongside   * `residencyAttributes()` on the same `McpSessionDO.init` span so per-session   * memory footprint is queryable next to isolate residency without a   * cross-span join. Empty by default: a host that has nothing free to read   * (or nothing beyond what `buildMcpServer` already returns) need not   * override this. MUST stay O(1)/O(count) over already-materialized state —   * never trigger a new query or serialize a catalog to compute these.   */  protected sessionFootprintAttributes(): Record<string, number> {    return {};  }   protected withTelemetry<A, E>(    effect: Effect.Effect<A, E>,    _incoming?: IncomingTraceHeaders,  ): Effect.Effect<A, E> {    return effect;  }   protected captureCause(_cause: Cause.Cause<unknown>): void {}   protected captureCauseEffect(cause: Cause.Cause<unknown>): Effect.Effect<string | undefined> {    return Effect.sync(() => {      this.captureCause(cause);      return undefined;    });  }   protected prepareErrorCaptureScope(): Effect.Effect<void> {    return Effect.void;  }   /**   * Declare that the Durable Object has finished deciding what to do about this   * cause — either it reported it through `captureCauseEffect`, or it   * recognized it as expected platform behaviour and deliberately did not.   *   * Either way the cause is *owned here*. The host's outer error   * instrumentation wraps the DO's entry points and would otherwise report the   * same rejection a second time as it escapes, producing two issues per   * failure; this is the hook that lets the host drop its own echo. Anything the   * DO never claims (an alarm crash, a transport fault) is untouched and keeps   * being reported by that instrumentation, which is the whole point of   * claiming explicitly instead of disabling it.   */  protected claimCauseHandled(_cause: Cause.Cause<unknown>): Effect.Effect<void> {    return Effect.void;  }   protected flushTelemetry(): Promise<void> {    return Promise.resolve();  }   /**   * The session id as it appears in logs and span attributes.   *   * Purely observational, and therefore total: a log line or a span attribute   * must never be able to abort the work it is describing. The alarm path is   * the one that used to prove this the hard way — it logged its decision   * before doing it, and on an invocation where the name could not be resolved   * the LOG threw and took the whole alarm with it.   */  protected sessionIdForTelemetry(): string {    return this.sessionIdOrUndefined() ?? UNRESOLVED_SESSION_ID;  }   /**   * The session id, or `undefined` when this DO's name cannot be resolved from   * any source. Derived from {@link sessionNameOrUndefined} so it agrees with   * the guard that decided the DO was addressable in the first place.   */  protected sessionIdOrUndefined(): string | undefined {    const name = this.sessionNameOrUndefined();    if (name === undefined) return undefined;    const [, sessionId] = name.split(":");    return sessionId ? sessionId : undefined;  }   /**   * The session id for callers that cannot proceed without one — routing an   * execution back to its owner, addressing an approval URL. Those callers want   * the throw, because a wrong id is worse than a failure.   */  protected get sessionId(): string {    const resolved = this.sessionIdOrUndefined();    if (resolved !== undefined) return resolved;    return this.getSessionId();  }   /**   * The single place that answers "what is this Durable Object's name".   *   * PartyServer has three sources and does not consult them all in the same   * place: `this.name` reads `ctx.id.name` and an in-memory field that is only   * hydrated during initialization, while the durable `__ps_name` record is   * read *only* by that initialization. An entry point that skips   * initialization — the alarm override below, which handles most of its   * decisions itself and only delegates to `super.alarm()` on one branch — can   * therefore find the durable record present, conclude the session is   * addressable, and still have `this.name` throw on the very next read.   *   * Everything in this class goes through this accessor instead, so the guard   * and the reads that follow it can never disagree.   */  private sessionNameOrUndefined(): string | undefined {    return this.ctx.id.name ?? this.resolvedSessionName ?? this.partyServerNameOrUndefined();  }   /**   * Ask every source, including durable storage, and remember the answer.   *   * The remembered value is what makes the synchronous accessor above agree   * with this one for the rest of the invocation: a name recovered from the   * `__ps_name` record stays available to callers that cannot await.   */  private async resolveSessionName(): Promise<string | undefined> {    const inMemory = this.sessionNameOrUndefined();    if (inMemory !== undefined) {      this.resolvedSessionName = inMemory;      return inMemory;    }    const stored = await this.ctx.storage.get<string>(PARTYSERVER_NAME_KEY);    if (!stored) return undefined;    this.resolvedSessionName = stored;    return stored;  }   /**   * Read PartyServer's own name resolution without inheriting its throw.   *   * PartyServer exposes no non-throwing probe for "do you have a name yet", so   * "it throws" IS the signal for "not resolvable from memory" — the same shape   * as {@link connectionsOrNone}. In a unit harness the getter also   * throws on its uninitialized private field, which reads identically here and   * is equally correct.   */  private partyServerNameOrUndefined(): string | undefined {    // oxlint-disable-next-line executor/no-try-catch-or-throw -- boundary: see doc comment; partyserver's `name` getter throws instead of reporting absence.    try {      return (this as { readonly name?: string }).name;    } catch {      return undefined;    }  }   protected currentParentSpan(): Tracer.AnySpan | undefined {    return undefined;  }   protected sessionTimeoutMs(): number {    return SESSION_TIMEOUT_MS;  }   protected maxPausedSessionIdleMs(): number {    return MAX_PAUSED_SESSION_IDLE_MS;  }   protected executionOwnerDirectory(): McpExecutionOwnerDirectory | null {

Lines 1–500 of 2191. Load more