@mcoda/mswarm 0.1.133 → 0.1.135-codali.2

package/dist/runtime.js

@@ -5,9 +5,10 @@ import { spawn } from "node:child_process";
 import { createHash, createHmac, randomBytes, randomUUID } from "node:crypto";
 import { isIP } from "node:net";
 import { Agent, fetch as undiciFetch } from "undici";
+import { nativeReasoningRequest } from "./native-reasoning.js";
 import { GlobalRepository } from "@mcoda/db";
 import { MswarmCodaliExecutor } from "./codali-executor.js";
-import { GENERATIVE_OPERATION_DEFAULT_PATHS, GENERATIVE_REQUEST_PARAMETER_ALLOWLISTS, MSWARM_CAPABILITY_SCHEMA_VERSION, CryptoHelper, assertMswarmSafeRelativePath, validateMswarmArchiveEntry, buildMswarmCapabilityNames, buildMswarmPrivateCapabilityCatalogEntry, buildMswarmLocalArtifactUri, buildMswarmSandboxProfile, defaultMswarmArtifactAccessPolicy, defaultMswarmArtifactRetentionPolicy, normalizeGenerativeModalities, normalizeSelfHostedGenerativeOperation, projectMswarmPublicCapabilities, validateMswarmGenericJobRequest } from "@mcoda/shared";
+import { GENERATIVE_OPERATION_DEFAULT_PATHS, GENERATIVE_REQUEST_PARAMETER_ALLOWLISTS, validateEmbeddingRequest, validateEmbeddingResponse, validateEmbeddingPreprocessing, MSWARM_CAPABILITY_SCHEMA_VERSION, CryptoHelper, assertMswarmSafeRelativePath, validateMswarmArchiveEntry, buildMswarmCapabilityNames, buildMswarmPrivateCapabilityCatalogEntry, buildMswarmLocalArtifactUri, buildMswarmSandboxProfile, defaultMswarmArtifactAccessPolicy, defaultMswarmArtifactRetentionPolicy, normalizeGenerativeModalities, normalizeSelfHostedGenerativeOperation, projectMswarmPublicCapabilities, validateMswarmGenericJobRequest } from "@mcoda/shared";
 const REGISTRATION_POLL_INTERVAL_MS = 5000;
 const REGISTRATION_POLL_TIMEOUT_MS = 30 * 60000;
 /**
@@ -202,12 +203,16 @@ function resolvePersistentNodePath(nodePath) {
 const DEFAULT_EXPOSE_ALL_MODELS = true;
 const LOCAL_OPENAI_COMPATIBLE_ADAPTERS = new Set([
     "openai-compatible-local",
+    "strata-local",
     "vllm-local",
     "llama-cpp-local",
     "llamacpp-local"
 ]);
 const SECRET_LOCAL_RUNNER_HEADER_KEYS = new Set(["authorization", "proxy-authorization", "x-api-key", "api-key"]);
 const RESERVED_LOCAL_RUNNER_EXTRA_BODY_KEYS = new Set([
+    "input",
+    "dimensions",
+    "encoding_format",
     "model",
     "messages",
     "prompt",
@@ -229,6 +234,7 @@ const RESERVED_LOCAL_RUNNER_EXTRA_BODY_KEYS = new Set([
     "temperature"
 ]);
 const REQUIRED_GENERATIVE_OPERATION_REQUEST_PARAMETERS = {
+    embeddings: ["model", "input", "encoding_format"],
     "chat.completions": ["model", "messages"],
     "images.generations": ["model", "prompt", "response_format"],
     "audio.generations": ["model", "prompt", "response_format"],
@@ -244,6 +250,7 @@ const DEFAULT_VIDEO_GENERATIVE_OUTPUT_BYTES = 256 * 1024 * 1024;
 const GENERATIVE_RESPONSE_JSON_OVERHEAD_BYTES = 64 * 1024;
 const LOCAL_RUNNER_KIND_ALIASES = {
     vllm: "vllm",
+    strata: "strata",
     "llama-cpp": "llama-cpp",
     "llama.cpp": "llama-cpp",
     llamacpp: "llama-cpp",
@@ -328,7 +335,7 @@ function isRecord(value) {
 }
 export function normalizeSelfHostedGenerativeOperationName(value) {
     const normalized = optionalText(value);
-    return normalized === "chat.completions" ||
+    return normalized === "embeddings" || normalized === "chat.completions" ||
         normalized === "images.generations" ||
         normalized === "audio.generations" ||
         normalized === "videos.generations"
@@ -369,7 +376,7 @@ function normalizeGenerativeOperationDescriptor(value) {
     ]));
     const responseFormats = normalized.responseFormats?.length
         ? normalized.responseFormats
-        : operation === "chat.completions"
+        : operation === "embeddings" ? ["float"] : operation === "chat.completions"
             ? undefined
             : operation === "images.generations"
                 ? ["b64_json"]
@@ -420,6 +427,8 @@ function normalizeGenerativeOperations(value) {
 function inferGenerativeModalities(operations, direction) {
     const modalities = new Set();
     for (const descriptor of operations) {
+        if (direction === "output" && descriptor.operation === "embeddings")
+            modalities.add("embedding");
         if (direction === "input" || descriptor.operation === "chat.completions") {
             modalities.add("text");
         }
@@ -438,6 +447,10 @@ function inferGenerativeModalities(operations, direction) {
 function generativeOperationCatalogInput(descriptor) {
     const limits = descriptor.limits
         ? {
+            ...(descriptor.limits.maxInputItems !== undefined ? { max_input_items: descriptor.limits.maxInputItems } : {}),
+            ...(descriptor.limits.maxInputChars !== undefined ? { max_input_chars: descriptor.limits.maxInputChars } : {}),
+            ...(descriptor.limits.maxTotalInputChars !== undefined ? { max_total_input_chars: descriptor.limits.maxTotalInputChars } : {}),
+            ...(descriptor.limits.embeddingDimensions !== undefined ? { embedding_dimensions: descriptor.limits.embeddingDimensions } : {}),
             ...(descriptor.limits.maxRequestBytes !== undefined
                 ? { max_request_bytes: descriptor.limits.maxRequestBytes }
                 : {}),
@@ -526,6 +539,8 @@ function defaultLocalRunnerKindForAdapter(adapter) {
     const lookupKey = normalizeLookupKey(normalized);
     if (lookupKey === "vllm-local")
         return "vllm";
+    if (lookupKey === "strata-local")
+        return "strata";
     if (lookupKey === "llama-cpp-local" || lookupKey === "llamacpp-local")
         return "llama-cpp";
     return undefined;
@@ -617,6 +632,7 @@ function normalizeMcodaLocalRunnerConfig(adapter, config) {
         inputModalities: normalizeGenerativeModalities(merged.inputModalities ?? merged.input_modalities),
         outputModalities: normalizeGenerativeModalities(merged.outputModalities ?? merged.output_modalities),
         operations: normalizeGenerativeOperations(merged.operations),
+        embedding: merged.embedding === undefined || merged.embedding === null ? undefined : validateEmbeddingPreprocessing(merged.embedding),
         publicModelId: readConfigString(merged, ["publicModelId", "public_model_id"])
     };
 }
@@ -1228,12 +1244,12 @@ export function agentRunsOnThisNode(agent) {
         return false;
     }
 }
-function isAgentExposed(agentSlug, defaultModel, bestUsage, config) {
+function isAgentExposed(agentSlug, defaultModel, bestUsage, config, supportsEmbeddings = false) {
     const identities = [agentSlug, defaultModel].filter((value) => Boolean(value));
     if (identities.some((identity) => config.modelBlocklist.includes(identity))) {
         return false;
     }
-    if (isEmbeddingUsage(bestUsage) || identities.some((identity) => isEmbeddingModel(identity))) {
+    if (!supportsEmbeddings && (isEmbeddingUsage(bestUsage) || identities.some((identity) => isEmbeddingModel(identity)))) {
         return false;
     }
     if (config.modelAllowlist.length > 0) {
@@ -1500,6 +1516,8 @@ Wants=network-online.target
 Type=simple
 ExecStart=${quoteSystemdValue(input.wrapperPath)}
 Restart=always
+TimeoutStopSec=infinity
+KillMode=mixed
 RestartSec=5
 StandardOutput=append:${input.logPath}
 StandardError=append:${input.errorLogPath}
@@ -2095,6 +2113,8 @@ export function mapMcodaAgentToSelfHostedModel(agent, config) {
         model_id: publicModelId || model,
         public_model_id: publicModelId,
         upstream_model: model,
+        digest: optionalText(legacyConfig.modelDigest),
+        embedding: localRunner?.embedding,
         base_url: localRunner?.baseUrl ?? legacyNativeBase ?? undefined,
         runner_kind: localRunner?.generativeRunnerKind ?? localRunner?.runnerKind ?? (legacyNativeBase ? "llama-cpp" : undefined),
         auth_mode: localRunner?.authMode ?? (legacyNativeBase ? normalizeLocalRunnerAuthMode(legacyConfig.authMode) : undefined),
@@ -2109,10 +2129,13 @@ export function mapMcodaAgentToSelfHostedModel(agent, config) {
         supports_vision: supportsVision,
         supports_json_schema: localRunner?.supportsJsonSchema,
         supports_gbnf: localRunner?.supportsGbnf,
+        // Transport metadata must not change semantic model capability sets: those
+        // sets participate in stable load-balanced pool identifiers on the gateway.
+        native_reasoning_transport_version: localRunner?.runnerKind === "strata" ? 1 : undefined,
         openai_compatible: openaiCompatible,
         exposed: healthStatus !== "blocked" &&
             agentRunsOnThisNode(agent) &&
-            isAgentExposed(slug, defaultModel, bestUsage, config),
+            isAgentExposed(slug, defaultModel, bestUsage, config, operations.some(entry => entry.operation === "embeddings")),
         best_usage: bestUsage,
         capabilities,
         cost_per_million: optionalNumber(agent.costPerMillion, agent.cost_per_million),
@@ -2465,6 +2488,8 @@ function mapMcodaAgentToCodaliAgent(agent, fallbackSlug) {
     const generativeOperations = localRunner?.operations ??
         [defaultGenerativeOperationDescriptor("chat.completions", localRunner ? "/chat/completions" : undefined)];
     return {
+        modelDigest: configText(config, "modelDigest"),
+        orchestrationConfigPath: configText(config, "codaliOrchestrationPath"),
         slug: optionalText(agent.slug) || fallbackSlug,
         adapter,
         provider: resolveCodaliProviderForAgent(agent),
@@ -2491,6 +2516,10 @@ function mapMcodaAgentToCodaliAgent(agent, fallbackSlug) {
         supportsJsonSchema: localRunner?.supportsJsonSchema,
         supportsGbnf: localRunner?.supportsGbnf,
         capabilities: normalizeCapabilities(agent.capabilities),
+        rating: optionalNumber(agent.rating) ?? undefined,
+        reasoningRating: optionalNumber(agent.reasoningRating, agent.reasoning_rating) ?? undefined,
+        maxComplexity: optionalNumber(agent.maxComplexity, agent.max_complexity) ?? undefined,
+        bestUsage: optionalText(agent.bestUsage) ?? optionalText(agent.best_usage) ?? undefined,
         contextWindow: optionalNumber(agent.contextWindow, agent.context_window) ?? undefined,
         maxOutputTokens: optionalNumber(agent.maxOutputTokens, agent.max_output_tokens) ?? undefined,
         generativeOperations,
@@ -4467,11 +4496,18 @@ export class MswarmSelfHostedNodeClient {
     }
 }
 export class SelfHostedNodeRuntime {
+    get isDraining() {
+        return this.admissionClosed || this.config.drainMode === true;
+    }
+    beginDrain() {
+        this.admissionClosed = true;
+    }
     constructor(config, deps) {
         this.activeChatJobs = 0;
         this.activeRequestIds = new Set();
         this.lastModels = [];
         this.loadRuntimeToken = null;
+        this.loadReportQueue = Promise.resolve();
         this.fastLoadUnsupported = false;
         this.nativeActiveJobs = 0;
         this.nativeProbeKeys = new Map();
@@ -4484,6 +4520,9 @@ export class SelfHostedNodeRuntime {
         this.latencySamplesMs = [];
         this.recentFailures = [];
         this.lifecyclePollingDisabled = false;
+        this.admissionClosed = false;
+        this.admittedJobs = 0;
+        this.manualRelayOperations = 0;
         /**
          * The capability snapshot a heartbeat carries, without making the heartbeat
          * wait for it.
@@ -4677,7 +4716,7 @@ export class SelfHostedNodeRuntime {
         return Math.round(total / this.latencySamplesMs.length);
     }
     buildLoadTelemetry(input) {
-        const drainMode = this.config.drainMode === true;
+        const drainMode = this.isDraining;
         const llmMaxConcurrency = Math.max(1, Math.floor(this.config.maxConcurrentLlmJobs || this.config.maxConcurrentJobs || 1));
         const genericMaxConcurrency = Math.max(1, Math.floor(this.config.genericJobMaxConcurrency || 1));
         const maxConcurrency = Math.max(1, Math.floor(this.config.maxConcurrentJobs || 1), llmMaxConcurrency, this.config.genericJobsEnabled ? genericMaxConcurrency : 1);
@@ -5034,7 +5073,7 @@ export class SelfHostedNodeRuntime {
             generic_jobs_enabled: this.config.genericJobsEnabled,
             generic_job_timeout_ms: this.config.genericJobTimeoutMs,
             generic_job_max_concurrency: this.config.genericJobMaxConcurrency,
-            drain_mode: this.config.drainMode === true,
+            drain_mode: this.isDraining,
             load_reporting_enabled: this.config.loadReportingEnabled !== false,
             hardware_telemetry_enabled: this.config.hardwareTelemetryEnabled === true,
             expose_all_models: this.config.exposeAllModels,
@@ -5100,6 +5139,9 @@ export class SelfHostedNodeRuntime {
             body.max_tokens = job.openai_request.max_tokens;
         if (job.openai_request.stop !== undefined)
             body.stop = job.openai_request.stop;
+        if (job.openai_request.response_format !== undefined)
+            body.response_format = job.openai_request.response_format;
+        Object.assign(body, nativeReasoningRequest(job.openai_request));
         const url = `${baseUrl.replace(/\/+$/, "")}/chat/completions`;
         const response = await this.fetchImpl(url, { method: "POST", headers, body: JSON.stringify(body) });
         const text = await response.text();
@@ -5358,6 +5400,61 @@ export class SelfHostedNodeRuntime {
             body
         };
     }
+    prepareEmbeddingPassthrough(job, agent) {
+        if (job.policy?.expected_model_digest && agent.modelDigest !== job.policy.expected_model_digest) {
+            throw new SelfHostedPreStartJobError("selected_agent_unavailable", "Embedding model revision changed after admission");
+        }
+        const descriptor = this.generativeOperationForAgent(agent, "embeddings");
+        const baseUrl = optionalText(agent.baseUrl);
+        let endpoint;
+        try {
+            endpoint = new URL(baseUrl || "");
+        }
+        catch {
+            throw new SelfHostedPreStartJobError("selected_agent_unavailable", "embedding runner URL is missing or invalid");
+        }
+        if (!["http:", "https:"].includes(endpoint.protocol) || endpoint.username || endpoint.password || endpoint.search || endpoint.hash) {
+            throw new SelfHostedPreStartJobError("selected_agent_unavailable", "embedding runner URL is invalid");
+        }
+        let body;
+        try {
+            body = validateEmbeddingRequest(job.openai_request, optionalText(agent.model) || job.openai_request.model, descriptor.limits);
+        }
+        catch (error) {
+            throw new SelfHostedPreStartJobError("validation_failed", error instanceof Error ? error.message : "invalid embedding request");
+        }
+        if (Buffer.byteLength(JSON.stringify(body)) > Math.min(descriptor.limits?.maxRequestBytes || 1048576, 1048576)) {
+            throw new SelfHostedPreStartJobError("validation_failed", "embedding request exceeds byte limit");
+        }
+        const headers = {};
+        const bearer = agent.authMode === "bearer" ? optionalText(agent.apiKey)
+            : agent.authMode === "dummy-bearer" ? optionalText(agent.dummyBearerToken) || "local" : undefined;
+        if (bearer)
+            headers.authorization = `Bearer ${bearer}`;
+        for (const [key, value] of Object.entries(agent.headers ?? {}))
+            if (typeof value === "string")
+                headers[key.toLowerCase()] = value;
+        headers["content-type"] = "application/json";
+        return { descriptor, body, headers, url: `${baseUrl.replace(/\/+$/, "")}${descriptor.path}` };
+    }
+    async executeEmbeddingPassthrough(prepared) {
+        const controller = new AbortController();
+        const timeout = setTimeout(() => controller.abort(), Math.min(this.config.jobTimeoutMs, 120000));
+        try {
+            const response = await this.fetchImpl(prepared.url, {
+                method: "POST", headers: prepared.headers, body: JSON.stringify(prepared.body), signal: controller.signal,
+            });
+            if (!response.ok) {
+                await response.body?.cancel().catch(() => undefined);
+                throw new Error(`local embeddings runner request failed with status ${response.status}`);
+            }
+            const raw = await readBoundedResponseText(response, Math.min(prepared.descriptor.limits?.maxOutputBytes || 16777216, 16777216), "embedding response exceeds byte limit");
+            return validateEmbeddingResponse(JSON.parse(raw), prepared.body.input.length, prepared.descriptor.limits?.embeddingDimensions ?? prepared.body.dimensions, prepared.body.model);
+        }
+        finally {
+            clearTimeout(timeout);
+        }
+    }
     async executeOpenAiGenerativePassthrough(prepared, operation) {
         const outputByteLimit = generativeOutputByteLimit(operation, prepared.descriptor.limits?.maxOutputBytes);
         const jsonResponseByteLimit = generativeJsonResponseByteLimit(outputByteLimit);
@@ -5505,6 +5602,21 @@ export class SelfHostedNodeRuntime {
         return codaliAgent;
     }
     async executeGenericJob(envelope, options = {}) {
+        if (this.isDraining)
+            throw new SelfHostedPreStartJobError("node_draining", "Node is draining; no new work was admitted.");
+        return this.executePreviouslyAdmittedGenericJob(envelope, options);
+    }
+    /** Host scheduler only: finish work accepted before admission was closed. */
+    async executePreviouslyAdmittedGenericJob(envelope, options = {}) {
+        this.admittedJobs += 1;
+        try {
+            return await this.executeAdmittedGenericJob(envelope, options);
+        }
+        finally {
+            this.admittedJobs -= 1;
+        }
+    }
+    async executeAdmittedGenericJob(envelope, options) {
         const startedAt = Date.now();
         this.beginExecutionTelemetry("generic_job");
         const events = [];
@@ -5678,12 +5790,26 @@ export class SelfHostedNodeRuntime {
         }
     }
     async executeJob(job, options = {}) {
+        if (this.isDraining) {
+            return {
+                job_id: job.job_id, request_id: job.request_id, status: "failed",
+                error: { code: "node_draining", message: "Node is draining; no new work was admitted." }
+            };
+        }
+        return this.executeAdmittedJob(job, options);
+    }
+    async executeAdmittedJob(job, options) {
+        this.admittedJobs += 1;
         this.activeRequestIds.add(job.request_id);
         try {
             return await this.executeTrackedJob(job, options);
         }
         finally {
             this.activeRequestIds.delete(job.request_id);
+            // A product often starts its next stage immediately after this result.
+            // Publish the newly free native slot before the gateway delivers it.
+            await this.reportLoad().catch(() => undefined);
+            this.admittedJobs -= 1;
         }
     }
     async executeTrackedJob(job, options) {
@@ -5749,6 +5875,9 @@ export class SelfHostedNodeRuntime {
                 throw new SelfHostedPreStartJobError("unsupported_operation", `unsupported self-hosted operation ${optionalText(job.operation) || String(job.operation)}`);
             }
             if (job.provider === "ollama") {
+                if (Object.keys(nativeReasoningRequest(job.openai_request)).length) {
+                    throw new SelfHostedPreStartJobError("unsupported_operation", "Native reasoning options require an OpenAI-compatible runner");
+                }
                 if (operation !== "chat.completions") {
                     throw new SelfHostedPreStartJobError("unsupported_operation", `Ollama relay jobs do not support ${operation}`);
                 }
@@ -5822,6 +5951,17 @@ export class SelfHostedNodeRuntime {
             const agent = await this.resolveMcodaAgentForJob(job);
             selectedAgent = agent;
             this.generativeOperationForAgent(agent, operation);
+            if (operation === "embeddings") {
+                const prepared = this.prepareEmbeddingPassthrough(job, agent);
+                await acknowledgeStarted(agent);
+                await recordProgress({ type: "agent_selected", operation, job_id: job.job_id,
+                    request_id: job.request_id, agent_slug: agent.slug, adapter: agent.adapter, local_model: prepared.body.model });
+                const result = await this.executeEmbeddingPassthrough(prepared);
+                this.finishExecutionTelemetry({ executionClass, startedAt, ok: true });
+                return { job_id: job.job_id, request_id: job.request_id, operation, status: "success",
+                    openai_response: { ...result, ...(agent.modelDigest ? { embedding_revision: agent.modelDigest } : {}) }, ...(result.usage ? { usage: { ...result.usage, completion_tokens: 0 } } : {}),
+                    ...(progressEvents.length ? { progress_events: progressEvents } : {}), timing: { local_latency_ms: Date.now() - startedAt } };
+            }
             if (operation === "images.generations" ||
                 operation === "audio.generations" ||
                 operation === "videos.generations") {
@@ -5899,6 +6039,7 @@ export class SelfHostedNodeRuntime {
                 attachedMswarmApiKey,
                 policy: buildCodaliPolicy(job),
                 temperature: job.openai_request.temperature,
+                nativeReasoning: nativeReasoningRequest(job.openai_request),
                 responseFormat: isRecord(job.openai_request.response_format)
                     ? job.openai_request.response_format
                     : null,
@@ -5976,6 +6117,7 @@ export class SelfHostedNodeRuntime {
                         agent_slug: response.metadata.agent_slug,
                         runtime: response.metadata.runtime,
                         codali_run_id: response.metadata.run_id,
+                        orchestration: response.metadata.orchestration,
                         tool_calls_executed: response.metadata.tool_calls_executed,
                         called_tools: response.metadata.called_tools,
                         dynamic_tools_considered: response.metadata.dynamic_tools_considered,
@@ -6057,14 +6199,15 @@ export class SelfHostedNodeRuntime {
     async refreshNativeLoad() {
         const runners = new Map();
         for (const model of this.lastModels) {
+            const isStrata = model.runner_kind === "strata" || model.adapter === "strata-local";
             if (model.exposed === false || !model.base_url ||
-                (model.runner_kind !== "llama-cpp" && model.adapter !== "llama-cpp-local"))
+                (!isStrata && model.runner_kind !== "llama-cpp" && model.adapter !== "llama-cpp-local"))
                 continue;
             try {
                 const url = new URL(model.base_url);
                 if (!["127.0.0.1", "localhost", "[::1]"].includes(url.hostname))
                     continue;
-                const endpoint = new URL("/slots", url).toString();
+                const endpoint = new URL(isStrata ? "/v1/status" : "/slots", url).toString();
                 const models = runners.get(endpoint) || [];
                 models.push(model);
                 runners.set(endpoint, models);
@@ -6089,12 +6232,30 @@ export class SelfHostedNodeRuntime {
                 }
                 const response = await this.fetchImpl(endpoint, { signal: AbortSignal.timeout(1500),
                     ...(apiKey ? { headers: { authorization: `Bearer ${apiKey}` } } : {}) });
-                // Slot introspection is optional and may require different local credentials.
-                if ([401, 403, 404].includes(response.status))
+                const isStrata = new URL(endpoint).pathname === "/v1/status";
+                // llama.cpp slot introspection is optional. Strata status also proves the model is loaded.
+                if (!isStrata && [401, 403, 404].includes(response.status))
                     return;
                 if (!response.ok)
                     throw new Error("runner_load_unavailable");
                 const slots = await response.json();
+                if (isStrata) {
+                    const concurrency = isRecord(slots) && isRecord(slots.concurrency) ? slots.concurrency.serving : undefined;
+                    const inFlight = isRecord(slots) && isRecord(slots.activity) ? slots.activity.in_flight : undefined;
+                    if (!isRecord(slots) || slots.loaded !== true ||
+                        typeof concurrency !== "number" || !Number.isSafeInteger(concurrency) || concurrency < 1 ||
+                        typeof inFlight !== "number" || !Number.isSafeInteger(inFlight) || inFlight < 0) {
+                        throw new Error("runner_load_unavailable");
+                    }
+                    const processing = Math.min(concurrency, inFlight);
+                    active += processing;
+                    const capacity = { max_concurrency: concurrency, active_jobs: processing,
+                        queued_jobs: Math.max(0, inFlight - concurrency), free_slots: Math.max(0, concurrency - inFlight),
+                        pool: `runner:${createHash("sha256").update(endpoint).digest("hex").slice(0, 16)}` };
+                    for (const model of models)
+                        modelCapacity[model.model_id || model.model || model.name] = capacity;
+                    return;
+                }
                 if (!Array.isArray(slots))
                     return;
                 const processing = slots.filter((slot) => isRecord(slot) && slot.is_processing === true).length;
@@ -6114,7 +6275,14 @@ export class SelfHostedNodeRuntime {
         this.nativeModelCapacity = modelCapacity;
         this.unavailableModelIds = [...new Set(unavailable)];
     }
-    async reportLoad() {
+    reportLoad() {
+        // Serialize probes and writes so an older busy snapshot cannot overwrite
+        // the completion snapshot when the periodic reporter is already running.
+        const pending = this.loadReportQueue.then(() => this.reportLoadSnapshot());
+        this.loadReportQueue = pending.catch(() => undefined);
+        return pending;
+    }
+    async reportLoadSnapshot() {
         if (this.config.loadReportingEnabled === false || this.fastLoadUnsupported || !this.loadRuntimeToken)
             return;
         await this.refreshNativeLoad();
@@ -6210,7 +6378,7 @@ export class SelfHostedNodeRuntime {
                 relay_mode: this.config.relayMode || "outbound",
                 load_reporting_enabled: loadReportingEnabled,
                 hardware_telemetry_enabled: this.config.hardwareTelemetryEnabled === true,
-                drain_mode: this.config.drainMode === true,
+                drain_mode: this.isDraining,
                 ...(lifecycleProtocolDegraded
                     ? {
                         lifecycle: {
@@ -6385,7 +6553,9 @@ export class SelfHostedNodeRuntime {
                 streamEventForwardingFailed = true;
             }
         };
-        let result = await this.executeJob(job, {
+        // A claim returned by an outstanding poll already belongs to this node,
+        // including a poll that finishes after drain began. Never discard it.
+        let result = await this.executeAdmittedJob(job, {
             attachedMswarmApiKey: claim.attachedMswarmApiKey,
             onStarted: async (event) => {
                 try {
@@ -6473,11 +6643,19 @@ export class SelfHostedNodeRuntime {
         return { executed: true, job_id: job.job_id, status: result.status };
     }
     async pollAndExecuteJob(waitMs = DEFAULT_JOB_POLL_WAIT_MS) {
-        const polled = await this.pollRelayJob(waitMs);
-        if (!polled.claim) {
-            return { executed: false, ...(polled.status ? { status: polled.status } : {}) };
+        if (this.isDraining)
+            return { executed: false };
+        this.manualRelayOperations += 1;
+        try {
+            const polled = await this.pollRelayJob(waitMs);
+            if (!polled.claim) {
+                return { executed: false, ...(polled.status ? { status: polled.status } : {}) };
+            }
+            return await this.executeRelayJobClaim(polled.claim);
+        }
+        finally {
+            this.manualRelayOperations -= 1;
         }
-        return this.executeRelayJobClaim(polled.claim);
     }
     async doctor() {
         const checks = [];
@@ -6591,6 +6769,7 @@ export class SelfHostedNodeRuntime {
         const currentRelayRetryDelayMs = () => Math.max(pollRetryDelayMs(relayPollFailureStreak), pollRetryDelayMs(relayExecutionFailureStreak));
         const schedulePoll = (delayMs) => {
             if (stopped ||
+                this.isDraining ||
                 polling ||
                 revoked ||
                 this.lifecyclePollingDisabled ||
@@ -6612,6 +6791,7 @@ export class SelfHostedNodeRuntime {
         };
         const poll = () => {
             if (stopped ||
+                this.isDraining ||
                 polling ||
                 revoked ||
                 this.lifecyclePollingDisabled ||
@@ -6710,7 +6890,27 @@ export class SelfHostedNodeRuntime {
             schedule();
             schedulePoll(0);
         });
-        return { stop: halt };
+        return {
+            stop: halt,
+            drain: async ({ timeoutMs } = {}) => {
+                if (timeoutMs !== undefined && (!Number.isFinite(timeoutMs) || timeoutMs < 0)) {
+                    throw new Error("drain timeoutMs must be a finite non-negative number");
+                }
+                this.beginDrain();
+                clearPollTimer();
+                // Keep regular heartbeats alive during a drain and publish zero free
+                // slots promptly. Failure to report cannot reopen local admission.
+                void this.reportLoad().catch(() => undefined);
+                const deadline = timeoutMs === undefined ? Infinity : Date.now() + timeoutMs;
+                while (polling || activeRelayExecutions > 0 || this.admittedJobs > 0 || this.queuedGenericJobs > 0 || this.manualRelayOperations > 0) {
+                    if (Date.now() >= deadline)
+                        return false;
+                    await sleep(Math.min(25, Math.max(1, deadline - Date.now())));
+                }
+                await this.reportLoad().catch(() => undefined);
+                return true;
+            }
+        };
     }
 }
 //# sourceMappingURL=runtime.js.map