@mongodb-js/agent-engine-runner-shared 0.11.3

package/dist/server/drain.js

@@ -0,0 +1,455 @@
+/**
+ * Execution-scoped drain receiver for Runner SDK servers.
+ *
+ * Mirrors Python's `agent_engine_runner_shared.server.drain`. Both runtimes are parallel
+ * implementations of the same server architecture, so this file and its
+ * Python twin must be changed together.
+ *
+ * When the OE durably cancels an execution it POSTs `/drain` to the Tool/AER
+ * runtime owning that execution. The runtime blocks new work for the
+ * execution, aborts tracked in-flight work where a signal channel exists,
+ * and records an idempotent final outcome the caller retrieves by re-sending
+ * the same request.
+ *
+ * Contract (pinned by the shared fixture in
+ * `client-libraries/test-fixtures/drain/contract.json`):
+ *
+ * - `POST /drain {request_id, execution_id, reason, deadline_at_ms, workspace_id}`
+ *   — `workspace_id` is mandatory on workspace-scoped runtimes (`APP_ID`
+ *   set) and must match it exactly.
+ * - `202 {"outcome": "accepted"}` while draining; `200` with the final
+ *   outcome (`completed` / `timed_out` / `delivery_failed`) afterwards.
+ * - `deadline_at_ms` is absolute. The runtime never extends it.
+ * - `accepted` is HTTP-level acceptance only; `completed` is the only
+ *   quiescence signal.
+ * - An execution this runtime never served is `delivery_failed` with
+ *   `reason_code=execution_not_found` — exact OE targeting does not prove a
+ *   false `completed` safe, since one execution can span AER and Tool
+ *   runtimes independently.
+ *
+ * TS-vs-Python divergence, matching the existing boundary tests: Python can
+ * cancel coroutine tasks outright; a promise has no cancellation, so "notify
+ * active work" here means aborting an AbortController. That reaches work
+ * only where a signal channel already exists (AER execution context, LLM
+ * streams). Plain tool functions receive no signal — they are tracked and
+ * admission-blocked, and outliving the deadline is an honest `timed_out`.
+ *
+ * State is process-local by design: a restarted runtime has lost its active
+ * work and answers `delivery_failed`; reconciliation across restarts belongs
+ * to the caller's durable record, not to this registry.
+ */
+import { z } from "zod";
+import { getLogger } from "../logger.js";
+const logger = getLogger("agent_engine_runner_shared.server.drain");
+/** How long a finalized drain record (and never-drained tombstone) survives
+ * for idempotent retries. Must be >= the caller-side retry window. */
+const DEFAULT_RECORD_TTL_S = 900;
+/** Upper bound on the caller-chosen drain window; the runtime never extends
+ * a deadline, and rejects ones set unreasonably far out. */
+const DEFAULT_MAX_DEADLINE_MS = 60_000;
+export const DrainRequestSchema = z.object({
+    request_id: z
+        .string()
+        .min(1)
+        .max(128)
+        .regex(/^[A-Za-z0-9._:-]+$/),
+    execution_id: z.string().min(1).max(256),
+    reason: z.enum(["execution_cancelled", "rolling_restart"]),
+    deadline_at_ms: z.number().int().positive(),
+    // Null is an accepted spelling of omission, matching the Python twin and
+    // the published OpenAPI contract (string|null). The route's truthiness
+    // checks treat it as a missing scope claim either way.
+    workspace_id: z.string().max(256).nullish(),
+});
+/**
+ * `AbortController.abort()` reason used by the per-call abort, distinct from
+ * DRAIN_ABORT_REASON so error paths can tell the two apart.
+ */
+export const CALL_INTERRUPT_REASON = "call_interrupted";
+/**
+ * `AbortController.abort()` reason used by the drain, so error paths can tell
+ * "stopped by cancellation" apart from the execution timeout, which shares
+ * the same controller in the AER.
+ */
+export const DRAIN_ABORT_REASON = "execution_drained";
+/**
+ * How long an abort that preceded its call's registration stays claimable by
+ * `beginWork`. The window covers the OE-claim → in-graph-dispatch handoff; the
+ * platform's dispatch claim makes a same-step collision inside it impossible.
+ * Mirrors Python's `PRE_ABORT_RETENTION_S`.
+ */
+const PRE_ABORT_RETENTION_MS = 60_000;
+function newIdle() {
+    let resolve;
+    const promise = new Promise((r) => {
+        resolve = r;
+    });
+    return { promise, resolve };
+}
+function drainingError() {
+    return Object.assign(new Error("execution_draining"), {
+        statusCode: 409,
+    });
+}
+/**
+ * Reject malformed drain tunables rather than silently collapsing retention
+ * or the deadline ceiling — `Number("bogus")` is NaN, which would evict
+ * records immediately or disable the deadline check entirely. Mirrors
+ * Python's fail-fast `_positive`.
+ */
+function parsePositiveNumber(raw, fallback, name) {
+    if (raw === undefined || raw.trim() === "")
+        return fallback;
+    const value = Number(raw);
+    if (!Number.isFinite(value) || value <= 0) {
+        throw new Error(`${name} must be a positive finite number, got ${raw}`);
+    }
+    return value;
+}
+/**
+ * Execution-scoped admission + active-work registry for one runtime process.
+ *
+ * Mutating methods are synchronous, so check-and-set is atomic on Node's
+ * single thread. The registry is created per server; per-process state only.
+ */
+export class DrainRegistry {
+    recordTtlMs;
+    maxDeadlineMsValue;
+    entries = new Map();
+    /**
+     * Aborts that arrived before their call registered, keyed by execution and
+     * step, valued by expiry. The window is the OE-claim → in-graph-dispatch
+     * handoff; the retention covers graph scheduling delays, not call lifetimes.
+     */
+    preAborted = new Map();
+    constructor(opts = {}) {
+        const env = opts.env ?? process.env;
+        this.recordTtlMs =
+            opts.recordTtlMs ??
+                parsePositiveNumber(env["RUNNER_DRAIN_RECORD_TTL_S"], DEFAULT_RECORD_TTL_S, "RUNNER_DRAIN_RECORD_TTL_S") * 1000;
+        this.maxDeadlineMsValue =
+            opts.maxDeadlineMs ??
+                parsePositiveNumber(env["RUNNER_DRAIN_MAX_DEADLINE_MS"], DEFAULT_MAX_DEADLINE_MS, "RUNNER_DRAIN_MAX_DEADLINE_MS");
+    }
+    get maxDeadlineMs() {
+        return this.maxDeadlineMsValue;
+    }
+    /** Reject work for a draining execution without registering it. */
+    checkAdmission(executionId) {
+        const entry = this.entries.get(executionId);
+        if (entry?.draining)
+            throw drainingError();
+    }
+    /**
+     * Register in-flight work for an execution, rejecting drained ones.
+     *
+     * `controller` is the abort handle the drain fires; omit it for work the
+     * runtime can only wait out (e.g. a tool function with no signal channel).
+     * `stepNumber` makes the work addressable by the per-call abort
+     * (`abortCall`); work registered without it is only ever drained
+     * execution-wide.
+     *
+     * Returns true when a per-call abort preceded this registration (the Stop
+     * landed in the OE-claim → in-graph-dispatch handoff): the caller must not
+     * start the body and must settle the call interrupted.
+     */
+    beginWork(executionId, controller, stepNumber) {
+        let entry = this.entries.get(executionId);
+        if (entry === undefined) {
+            entry = {
+                activeCount: 0,
+                controllers: new Set(),
+                draining: false,
+                drain: null,
+                idle: newIdle(),
+                byStep: new Map(),
+            };
+            this.entries.set(executionId, entry);
+        }
+        else {
+            if (entry.draining)
+                throw drainingError();
+            this.cancelEviction(entry);
+            if (entry.activeCount === 0) {
+                // Reviving a tombstone (e.g. a HITL resume reusing the id).
+                entry.idle = newIdle();
+            }
+        }
+        entry.activeCount += 1;
+        if (controller !== undefined)
+            entry.controllers.add(controller);
+        let preAborted = false;
+        if (stepNumber !== undefined) {
+            preAborted = this.popPreAborted(executionId, stepNumber);
+            // A step maps to one in-flight call by the platform's dispatch claim,
+            // so this can only replace a settled entry.
+            entry.byStep.set(stepNumber, {
+                controller,
+                settled: false,
+                aborted: preAborted,
+            });
+            if (preAborted && controller !== undefined) {
+                controller.abort(CALL_INTERRUPT_REASON);
+            }
+        }
+        return preAborted;
+    }
+    preAbortKey(executionId, stepNumber) {
+        return `${executionId}\n${stepNumber}`;
+    }
+    popPreAborted(executionId, stepNumber) {
+        const key = this.preAbortKey(executionId, stepNumber);
+        const expiry = this.preAborted.get(key);
+        this.preAborted.delete(key);
+        this.sweepPreAborted();
+        return expiry !== undefined && expiry > Date.now();
+    }
+    sweepPreAborted() {
+        if (this.preAborted.size === 0)
+            return;
+        const now = Date.now();
+        for (const [key, expiry] of this.preAborted) {
+            if (expiry <= now)
+                this.preAborted.delete(key);
+        }
+    }
+    /**
+     * Signal exactly one in-flight call, leaving the execution open — the
+     * surgical sibling of a drain: no admission latch, no record, no deadline,
+     * and later steps of the same execution proceed. Idempotent: a repeat
+     * abort of the same step re-reads the same outcome.
+     */
+    abortCall(executionId, stepNumber) {
+        const entry = this.entries.get(executionId);
+        const work = entry?.byStep.get(stepNumber);
+        if (work === undefined) {
+            // The abort can race the call's registration: OE claims the call, then
+            // the runtime registers only when the in-graph body starts. Retain it
+            // briefly so beginWork settles the late registration as interrupted
+            // instead of running a call the caller stopped.
+            this.sweepPreAborted();
+            this.preAborted.set(this.preAbortKey(executionId, stepNumber), Date.now() + PRE_ABORT_RETENTION_MS);
+            return "not_found";
+        }
+        if (work.settled)
+            return "already_settled";
+        // The missing-controller check precedes the repeat check: a repeat abort
+        // of track-only work must keep answering not_cancellable, matching the
+        // Python receiver.
+        if (work.controller === undefined)
+            return "not_cancellable";
+        if (work.aborted)
+            return "interrupted";
+        work.aborted = true;
+        work.controller.abort(CALL_INTERRUPT_REASON);
+        return "interrupted";
+    }
+    /**
+     * Close a call's interruptibility as its result report begins: the body
+     * already produced its outcome, so a late abort must read already_settled
+     * rather than claim a stop the durable record will contradict.
+     */
+    claimSettlement(executionId, stepNumber) {
+        const work = this.entries.get(executionId)?.byStep.get(stepNumber);
+        if (work !== undefined)
+            work.settled = true;
+    }
+    /**
+     * Attach an abort handle to already-registered work. The AER builds its
+     * execution-wide controller inside the handler, after admission — a drain
+     * that landed in between aborts it here, at attach time.
+     */
+    attachController(executionId, controller) {
+        const entry = this.entries.get(executionId);
+        if (entry === undefined)
+            return;
+        entry.controllers.add(controller);
+        // Abort with the drain reason: the AER shares this controller with its
+        // execution timeout and tells them apart by signal.reason.
+        if (entry.draining)
+            controller.abort(DRAIN_ABORT_REASON);
+    }
+    endWork(executionId, controller, stepNumber) {
+        const entry = this.entries.get(executionId);
+        if (entry === undefined)
+            return;
+        if (controller !== undefined)
+            entry.controllers.delete(controller);
+        if (stepNumber !== undefined) {
+            // Keep the settled step mapped: a late abort must read already_settled,
+            // not not_found. The entry's TTL eviction reclaims it.
+            const work = entry.byStep.get(stepNumber);
+            if (work !== undefined)
+                work.settled = true;
+        }
+        entry.activeCount -= 1;
+        if (entry.activeCount > 0)
+            return;
+        // No work remains: drop every tracked handle (drained or not) so a
+        // tombstone never retains controllers and their listeners for its TTL.
+        entry.controllers.clear();
+        entry.idle.resolve();
+        // A pending drain's finalizer owns eviction from here; otherwise the
+        // tombstone (and any finalized drain record) expires on the TTL.
+        if (entry.drain === null || entry.drain.outcome !== null) {
+            this.scheduleEviction(executionId, entry);
+        }
+    }
+    /** Apply one drain request; idempotent by request_id, coalescing per execution. */
+    apply(request) {
+        const entry = this.entries.get(request.execution_id);
+        if (entry?.drain != null) {
+            // Same request_id or a different one: one drain per execution.
+            entry.drain.requestIds.add(request.request_id);
+            return this.current(entry.drain);
+        }
+        if (entry === undefined) {
+            // Unknown execution: not applied here, but still record the outcome
+            // and latch admission. A dispatch that slips past the OE gate and lands
+            // after this drain must not start, and a retry of the same request must
+            // see the same answer for the retention window.
+            const record = {
+                requestIds: new Set([request.request_id]),
+                reason: request.reason,
+                deadlineAtMs: request.deadline_at_ms,
+                outcome: "delivery_failed",
+                reasonCode: "execution_not_found",
+                finalizedAtMs: Date.now(),
+            };
+            const latched = {
+                activeCount: 0,
+                controllers: new Set(),
+                draining: true,
+                drain: record,
+                idle: newIdle(),
+                byStep: new Map(),
+            };
+            this.entries.set(request.execution_id, latched);
+            this.scheduleEviction(request.execution_id, latched);
+            return this.final("delivery_failed", "execution_not_found");
+        }
+        if (entry.activeCount === 0) {
+            // Known, already-ended execution: nothing to drain here. Latch
+            // admission anyway — a drained execution never accepts work again — and
+            // restart retention so the record lives a full TTL from now.
+            entry.draining = true;
+            const record = {
+                requestIds: new Set([request.request_id]),
+                reason: request.reason,
+                deadlineAtMs: request.deadline_at_ms,
+                outcome: "completed",
+                finalizedAtMs: Date.now(),
+            };
+            entry.drain = record;
+            this.scheduleEviction(request.execution_id, entry);
+            return this.final("completed");
+        }
+        entry.draining = true;
+        const record = {
+            requestIds: new Set([request.request_id]),
+            reason: request.reason,
+            deadlineAtMs: request.deadline_at_ms,
+            outcome: null,
+        };
+        entry.drain = record;
+        for (const controller of entry.controllers) {
+            controller.abort(DRAIN_ABORT_REASON);
+        }
+        void this.finalize(request.execution_id, entry, record);
+        logger.info(`Drain accepted for execution ${request.execution_id} (reason=${request.reason}, active=${entry.activeCount})`);
+        return { status: 202, body: { outcome: "accepted" } };
+    }
+    async finalize(executionId, entry, record) {
+        const remainingMs = record.deadlineAtMs - Date.now();
+        let timer;
+        const timedOut = await Promise.race([
+            entry.idle.promise.then(() => false),
+            new Promise((r) => {
+                timer = setTimeout(() => r(true), Math.max(0, remainingMs));
+            }),
+        ]);
+        clearTimeout(timer);
+        if (timedOut) {
+            // Late completion must not overwrite the recorded timeout.
+            record.outcome = "timed_out";
+            record.reasonCode = "deadline_exceeded";
+        }
+        else {
+            record.outcome = "completed";
+        }
+        record.finalizedAtMs = Date.now();
+        logger.info(`Drain ${record.outcome} for execution ${executionId}`);
+        if (entry.activeCount === 0) {
+            this.scheduleEviction(executionId, entry);
+        }
+    }
+    scheduleEviction(executionId, entry) {
+        this.cancelEviction(entry);
+        const timer = setTimeout(() => this.evict(executionId, entry), this.recordTtlMs);
+        // The timer is memory hygiene, not work the process must stay alive for.
+        timer.unref();
+        entry.eviction = timer;
+    }
+    cancelEviction(entry) {
+        if (entry.eviction !== undefined) {
+            clearTimeout(entry.eviction);
+            entry.eviction = undefined;
+        }
+    }
+    evict(executionId, entry) {
+        // A retried drain after eviction finds no entry and gets
+        // execution_not_found; it never re-runs side effects.
+        if (this.entries.get(executionId) === entry && entry.activeCount === 0) {
+            this.entries.delete(executionId);
+        }
+    }
+    current(record) {
+        if (record.outcome === null) {
+            return { status: 202, body: { outcome: "accepted" } };
+        }
+        return this.final(record.outcome, record.reasonCode);
+    }
+    final(outcome, reasonCode) {
+        return {
+            status: 200,
+            body: reasonCode === undefined
+                ? { outcome }
+                : { outcome, reason_code: reasonCode },
+        };
+    }
+}
+/**
+ * Register `POST /drain` on a runner server.
+ *
+ * The bearer-auth hook already gates the path; this adds the
+ * execution/workspace-scoped validation on top. Callers of this endpoint
+ * must not follow redirects — the bearer token must never be forwarded to
+ * another host.
+ *
+ * Scope check: the bearer token authenticates the caller but does not
+ * establish that the named drain belongs to this runtime, so when the
+ * platform scopes this process to a workspace (`APP_ID` set — every managed
+ * runtime), the request must name that workspace exactly; a missing
+ * `workspace_id` is a 400, a mismatch a 403. When `APP_ID` is unset (local
+ * development against a single unscoped runtime) no workspace check is
+ * possible and none is enforced.
+ */
+export function registerDrainRoute(app, registry, env = process.env) {
+    const workspaceId = (env["APP_ID"] ?? "").trim();
+    app.post("/drain", async (request, reply) => {
+        const body = DrainRequestSchema.parse(request.body);
+        if (workspaceId) {
+            if (!body.workspace_id) {
+                return reply.code(400).send({ detail: "workspace_id required" });
+            }
+            if (body.workspace_id !== workspaceId) {
+                return reply.code(403).send({ detail: "workspace mismatch" });
+            }
+        }
+        if (body.deadline_at_ms > Date.now() + registry.maxDeadlineMs) {
+            return reply.code(400).send({ detail: "deadline exceeds maximum" });
+        }
+        const verdict = registry.apply(body);
+        return reply.code(verdict.status).send(verdict.body);
+    });
+}

package/dist/server/function.d.ts

@@ -0,0 +1,77 @@
+/**
+ * Function mode for the Runner SDK.
+ *
+ * Function mode runs a tool CLI-style: the process boots, runs a single named
+ * tool from the registry, reports the result to the OE, and exits. It is the
+ * per-call counterpart to the long-running tool server (mode `tool`).
+ *
+ * Mirrors Python's `agent_engine_runner_shared/server/function.py`.
+ *
+ * How the invocation arrives: fctr seeds `RunRequest.metadata` into the guest
+ * metadata directory before the workload starts. The OE sets one key,
+ * `request`, to a JSON `ToolFunctionRequest` envelope carrying both the
+ * `ToolPodExecuteRequest` and the tool-call `step` number (required by
+ * `ToolResultRequest`, owned by OE dispatch). The result is POSTed to
+ * `{oe_url}/tool/result`, the same endpoint the AER already uses.
+ */
+import { z } from "zod";
+import { type ToolResultRequest } from "../models.js";
+import type { ITenantRuntime } from "./base.js";
+/** The function-mode invocation envelope delivered at /run/meta/request. */
+export declare const ToolFunctionRequestSchema: z.ZodObject<{
+    request: z.ZodObject<{
+        platform_trace_id: z.ZodOptional<z.ZodNullable<z.ZodString>>;
+        execution_id: z.ZodString;
+        tool_name: z.ZodString;
+        arguments: z.ZodRecord<z.ZodString, z.ZodType<import("../models.js").JsonValue, unknown, z.core.$ZodTypeInternals<import("../models.js").JsonValue, unknown>>>;
+        tool_call_id: z.ZodOptional<z.ZodString>;
+        step_number: z.ZodOptional<z.ZodNumber>;
+        session_id: z.ZodString;
+        user_id: z.ZodOptional<z.ZodString>;
+        oe_url: z.ZodOptional<z.ZodString>;
+        oe_owner_url: z.ZodOptional<z.ZodNullable<z.ZodString>>;
+        authorization: z.ZodOptional<z.ZodObject<{
+            token: z.ZodString;
+            expires_at: z.ZodOptional<z.ZodNullable<z.ZodNumber>>;
+        }, z.core.$strip>>;
+        custom_headers: z.ZodOptional<z.ZodRecord<z.ZodString, z.ZodString>>;
+        payload: z.ZodOptional<z.ZodRecord<z.ZodString, z.ZodUnknown>>;
+        metadata: z.ZodDefault<z.ZodRecord<z.ZodString, z.ZodUnknown>>;
+    }, z.core.$strip>;
+    step: z.ZodNumber;
+}, z.core.$strip>;
+export type ToolFunctionRequest = z.infer<typeof ToolFunctionRequestSchema>;
+/**
+ * Encode the result to a JSON-safe object for the /tool/result POST.
+ *
+ * A tool may return a value the JSON encoder cannot handle (e.g. a circular
+ * reference). Downgrade any encode failure to an error report so the OE still
+ * receives a terminal signal. Mirrors Python's _result_payload.
+ */
+export declare function buildResultPayload(result: ToolResultRequest, executionId: string): Record<string, unknown>;
+/**
+ * Runs one metadata-delivered tool call and exits. Not a server.
+ *
+ * Reads the ToolFunctionRequest envelope from the guest metadata directory,
+ * invokes the named tool from the registry, reports the ToolResultRequest
+ * to {oe_url}/tool/result, and returns. Mirrors Python's ToolFunctionRunner.
+ *
+ * `metadataDir` is injectable for tests.
+ */
+export declare class ToolFunctionRunner {
+    private readonly runtime;
+    private readonly metadataDir;
+    constructor(runtime: ITenantRuntime, { metadataDir }?: {
+        metadataDir?: string;
+    });
+    /**
+     * Run one tool from the metadata-delivered request and report the result.
+     *
+     * Re-raises if the result POST cannot be delivered after retries so fctr
+     * records the execution as failed.
+     */
+    run(): Promise<void>;
+    private _prepare;
+    private _invokeToolFn;
+}
+//# sourceMappingURL=function.d.ts.map

package/dist/server/function.d.ts.map

@@ -0,0 +1 @@
+{"version":3,"file":"function.d.ts","sourceRoot":"","sources":["../../src/server/function.ts"],"names":[],"mappings":"AAAA;;;;;;;;;;;;;;;GAeG;AAKH,OAAO,EAAE,CAAC,EAAE,MAAM,KAAK,CAAC;AASxB,OAAO,EAIL,KAAK,iBAAiB,EACvB,MAAM,cAAc,CAAC;AAgBtB,OAAO,KAAK,EAAE,cAAc,EAAE,MAAM,WAAW,CAAC;AAchD,4EAA4E;AAC5E,eAAO,MAAM,yBAAyB;;;;;;;;;;;;;;;;;;;;;iBAGpC,CAAC;AACH,MAAM,MAAM,mBAAmB,GAAG,CAAC,CAAC,KAAK,CAAC,OAAO,yBAAyB,CAAC,CAAC;AAoE5E;;;;;;GAMG;AACH,wBAAgB,kBAAkB,CAChC,MAAM,EAAE,iBAAiB,EACzB,WAAW,EAAE,MAAM,GAClB,MAAM,CAAC,MAAM,EAAE,OAAO,CAAC,CAqBzB;AAgGD;;;;;;;;GAQG;AACH,qBAAa,kBAAkB;IAI3B,OAAO,CAAC,QAAQ,CAAC,OAAO;IAH1B,OAAO,CAAC,QAAQ,CAAC,WAAW,CAAS;gBAGlB,OAAO,EAAE,cAAc,EACxC,EAAE,WAA0B,EAAE,GAAE;QAAE,WAAW,CAAC,EAAE,MAAM,CAAA;KAAO;IAK/D;;;;;OAKG;IACG,GAAG,IAAI,OAAO,CAAC,IAAI,CAAC;IA+E1B,OAAO,CAAC,QAAQ;YASF,aAAa;CA+E5B"}