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 type {  McpExecutionOwnerDirectory,  McpExecutionOwnerRecord,  McpExecutionOwnerRoute,} from "./execution-owner-directory";import {  MAX_PAUSED_SESSION_IDLE_MS,  SESSION_TIMEOUT_MS,  decideSessionAlarm,  pausedLeaseExtensionLog,  runningLeaseExtensionLog,} from "./session-alarm-policy"; export type IncomingTraceHeaders = IncomingPropagationHeaders; export interface McpSessionInit {  readonly organizationId: string;  readonly userId: string;  readonly elicitationMode: McpElicitationMode;  /** 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?: "browser" | "model" | "native";  /** The MCP resource the session serves (carried from {@link McpSessionInit});   *  `buildMcpServer` scopes the tool catalog to it. */  readonly resource: McpResource;  readonly webOrigin?: string;} 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";/** 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;  private onStartPromise: Promise<void> | null = null;  private lastActivityMs = 0;  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>;   protected abstract resolveSessionMeta(token: McpSessionInit): Effect.Effect<SessionMeta>;   protected abstract buildMcpServer(    sessionMeta: SessionMeta,    dbHandle: TDbHandle,  ): Effect.Effect<BuiltMcpServer>;   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;  }   protected flushTelemetry(): Promise<void> {    return Promise.resolve();  }   protected get sessionId(): string {    return this.getSessionId();  }   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 {    return null;  }   protected executionOwnerRoute(): McpExecutionOwnerRoute {    return { sessionId: this.sessionId };  }   protected sameExecutionOwnerRoute(a: McpExecutionOwnerRoute, b: McpExecutionOwnerRoute): boolean {    return a.sessionId === b.sessionId;  }   protected forwardModelResumeToOwner(    _owner: McpExecutionOwnerRoute,    _identity: McpApprovalOwner,    _executionId: string,    _response: ResumeResponse,  ): Effect.Effect<McpSessionModelResumeResult, unknown> {    return Effect.succeed({      status: "execution_expired",      ttlMs: PAUSED_APPROVAL_TIMEOUT_MS,    });  }   protected readonly browserApprovalStore: BrowserApprovalStore = {    takeResponse: (executionId) => this.takeApprovalResponse(executionId),    waitForResponse: (executionId) => this.waitForApprovalResponse(executionId),  };   protected readonly modelResumeFallback = (    executionId: string,    response: ResumeResponse,  ): Effect.Effect<ResumeFallbackOutcome | null> =>    this.resumeFromExecutionOwnerDirectory(executionId, response);   protected readonly pausedExecutionHooks: PausedExecutionHooks = {    onExecutionPaused: (executionId, deadline) =>      Effect.sync(() => {        this.queuePendingApprovalLeaseStart(executionId, deadline);      }),    onResumeStarted: (executionId) => this.beginPendingApprovalResume(executionId),    onResumeSettled: (executionId) => this.finishPendingApprovalResume(executionId),  };   override async onConnect(conn: Connection, context: ConnectionContext): Promise<void> {    const requestIds = readActivePostRequestIds(context.request);    if (requestIds.length === 0) {      await super.onConnect(conn, context);      return;    }     await this.keepAliveWhile(async () => {      await this.setStreamRequestIds(conn.id, [...requestIds]);      await super.onConnect(conn, context);    });  }   private openSessionDbHandle(): Effect.Effect<TDbHandle> {    return Effect.promise(() => Promise.resolve(this.openSessionDb()));  }   private loadSessionMeta(): Effect.Effect<SessionMeta | null> {    return Effect.promise(async () => {      if (this.sessionMeta) return this.sessionMeta;      const stored = await this.ctx.storage.get<SessionMeta>(SESSION_META_KEY);      // Backfill `resource` for sessions persisted before scoped toolkits added      // the field. Their stored meta has no `resource`, and every such session      // was minted against the default `/mcp` endpoint, so default it here      // rather than let owner validation read `.kind` off undefined.      this.sessionMeta = stored        ? { ...stored, resource: stored.resource ?? defaultMcpResource }        : null;      return this.sessionMeta;    }).pipe(Effect.withSpan("mcp.session.load_meta"));  }   private async saveSessionMeta(sessionMeta: SessionMeta): Promise<void> {    this.sessionMeta = sessionMeta;    await this.ctx.storage.put(SESSION_META_KEY, sessionMeta);  }   private async markActivity(now = Date.now()): Promise<void> {    this.lastActivityMs = now;    await Promise.all([      this.ctx.storage.put(LAST_ACTIVITY_KEY, now),      this.ctx.storage.setAlarm(now + this.sessionTimeoutMs()),    ]);  }   private async loadLastActivity(): Promise<number> {    if (this.lastActivityMs > 0) return this.lastActivityMs;    const stored = await this.ctx.storage.get<number>(LAST_ACTIVITY_KEY);    this.lastActivityMs = stored ?? 0;    return this.lastActivityMs;  }   private async hasPartyServerName(): Promise<boolean> {    if (this.ctx.id.name) return true;    const stored = await this.ctx.storage.get<string>(PARTYSERVER_NAME_KEY);    return !!stored;  }   private activeStreamCount(): number {    return this.connectionsOrNone().length;  }   private async runningExecutionCount(): Promise<number> {    // Only requests still awaiting a result count as running work. Undelivered    // response markers (the transport's __mcp_undelivered_stream__: keys)    // deliberately do NOT extend the lease: the response is persisted in storage,    // which survives disposeIdleRuntime, so a later reconnect GET re-inits the    // DO and replays it. Counting them would make every delivered-but-unacked    // POST response pin the runtime alive indefinitely.    const rows = await this.ctx.storage.list<readonly JsonRpcRequestId[]>({      prefix: MCP_STREAM_REQS_KEY_PREFIX,      limit: 1_000,    });    let count = 0;    for (const requestIds of rows.values()) {      if (Array.isArray(requestIds)) count += requestIds.length;    }    return count;  }   private closeActiveStreams(): void {    for (const connection of this.connectionsOrNone()) {      // oxlint-disable-next-line executor/no-try-catch-or-throw -- boundary: best-effort WebSocket close during runtime disposal.      try {        connection.close(1000, "Session closed");      } catch {}    }  }   /**   * partyserver's `getConnections` dereferences a `#connectionManager`   * private field that is only initialized once the DO has accepted a   * websocket (never in unit harnesses), and partyserver exposes no   * non-throwing probe for that state, so "it throws" IS the signal for   * "no connections yet". Treating that as an empty set is safe for both   * callers: `closeActiveStreams` then has nothing to close, and   * `activeStreamCount` feeds the idle-lease decision where zero at worst   * disposes an idle-looking runtime whose undelivered responses are   * persisted in durable storage and replayed by the next reconnect GET.   * Before this guard the alarm crashed and retried instead, which kept   * the session pinned without ever making progress.   */  private connectionsOrNone(): ReadonlyArray<Connection> {    const getConnections = (this as { getConnections?: () => Iterable<Connection> }).getConnections;    if (!getConnections) return [];    // oxlint-disable-next-line executor/no-try-catch-or-throw -- boundary: see doc comment; partyserver offers no non-throwing way to ask whether the connection manager exists.    try {      return Array.from(getConnections.call(this));    } catch {      return [];    }  }   private async cleanupUnaddressableSessionAlarm(): Promise<void> {    await Effect.runPromise(this.closeRuntime());    await Effect.runPromise(      Effect.all([        Effect.ignore(Effect.tryPromise(() => this.ctx.storage.deleteAlarm())),        Effect.ignore(Effect.tryPromise(() => this.ctx.storage.delete(LAST_ACTIVITY_KEY))),      ]),    );  }   private async disposeIdleRuntime(input: {    readonly idleMs: number;    readonly pausedExecutionCount: number;  }): Promise<void> {    console.info(      JSON.stringify({        event: "mcp_session_idle_runtime_dispose",        sessionId: this.sessionId,        idleMs: input.idleMs,        pausedExecutionCount: input.pausedExecutionCount,      }),    );    await Effect.runPromise(this.closeRuntime());    await Effect.runPromise(      Effect.all([        Effect.ignore(Effect.tryPromise(() => this.ctx.storage.deleteAlarm())),        Effect.ignore(Effect.tryPromise(() => this.ctx.storage.delete(LAST_ACTIVITY_KEY))),      ]),    );  }   private resolveAndStoreSessionMeta(token: McpSessionInit) {    const self = this;    return Effect.gen(function* () {      const resolved = yield* self.resolveSessionMeta(token);      const sessionMeta: SessionMeta = {        ...resolved,        ...(token.webOrigin ? { webOrigin: token.webOrigin } : {}),      };      yield* Effect.promise(() => self.saveSessionMeta(sessionMeta)).pipe(        Effect.withSpan("mcp.session.save_meta"),      );      return sessionMeta;    }).pipe(Effect.withSpan("mcp.session.resolve_and_store_meta"));  }   private recordCauseOnSpan(cause: Cause.Cause<unknown>): Effect.Effect<void> {    const errors = Cause.prettyErrors(cause);    if (errors.length === 0) return Effect.void;    const first = errors[0];    return Effect.annotateCurrentSpan({      "exception.type": first?.name ?? "Error",      "exception.message": first?.message ?? "unknown",      "exception.stacktrace": Cause.pretty(cause),    });  }   private logExecutionOwnerDirectoryFailure(input: {    readonly operation: "put" | "get" | "delete";    readonly executionId: string;    readonly cause: Cause.Cause<unknown>;  }): Effect.Effect<void> {    const self = this;    return Effect.gen(function* () {      const first = Cause.prettyErrors(input.cause)[0];      console.error(        JSON.stringify({          event: "mcp_execution_owner_directory_error",          operation: input.operation,          executionId: input.executionId,          sessionId: self.sessionId,          exceptionType: first?.name ?? "Error",          exceptionMessage: first?.message ?? "unknown",          cause: Cause.pretty(input.cause),

Lines 1–500 of 1309. Load more