@nodetool-ai/websocket 0.8.1 → 0.8.2

package/dist/session/job-execution.js

@@ -1,3 +1,6 @@
+import { ROOT_CONTEXT, propagation, trace, isSpanContextValid } from "@opentelemetry/api";
+import { getRegisteredRunTrace } from "@nodetool-ai/models";
+import { ensureRunTraceTelemetry, openAppRunTrace, registerWorkflowRunTrace } from "@nodetool-ai/execution";
 /**
  * Everything between "run_job arrived" and "terminal status persisted": the
  * admission gates, the concurrency queue, run start-up, the message stream,
@@ -9,7 +12,9 @@
  * {@link JobExecutionDeps}.
  */
 import { randomUUID } from "node:crypto";
+import { operationTarget, parseInputStateKey } from "@nodetool-ai/app-runtime";
 import { ApiErrorCode } from "../error-codes.js";
+import { captureError } from "../error-traces.js";
 import { admitSpend, releaseSpend, reserveSpend } from "../credit-gate.js";
 import { JobConcurrencyQueue } from "../job-queue.js";
 import { isBoolean, isFiniteNumber, isFunctionValue, isNonEmptyString, isNumber, isObjectLike, isRecord, isString } from "../lib/wire-values.js";
@@ -18,17 +23,19 @@ import { ExecutionSession, isExecutionPreflightError, isUnitBilledCapability, pr
 import { persistableOutputs, withoutInlineAssetBytes } from "@nodetool-ai/execution/service";
 import { createRunSupervisor } from "../run-supervisor.js";
 import { jobRunRegistry } from "../job-run-registry.js";
-import { Application, Asset, Job, listApplicationVersions, invocationIdInUse, releasedApplicationRelease, releasedApplicationVersion, reserveInvocation, settleInvocation, Workflow } from "@nodetool-ai/models";
+import { Application, AppRunError, claimAppRun, ensureDefaultAppInstance, findAppRunByInvocation, getAppInstance, getAppRun, Asset, Job, listApplicationVersions, invocationIdInUse, releasedApplicationRelease, releasedApplicationVersion, reconcileAppRunCost, reserveInvocation, reserveAppRun, setAppRunInputs, settleAppRun, updateAppRunEstimate, settleInvocation, Workflow } from "@nodetool-ai/models";
 import { getInstanceId } from "../lib/instance-id.js";
 import { requestRemoteJobCancel } from "../job-control.js";
 import { estimateWorkflowCost, nodeExpectedQuantity } from "@nodetool-ai/node-sdk/cost-estimate";
 import { extractPricingParams } from "@nodetool-ai/node-sdk/pricing-params";
 import { getModelUnitPrice } from "@nodetool-ai/model-pricing";
 import { applicationReleaseResponse } from "@nodetool-ai/protocol/api-schemas/applications.js";
+import { appRunSnapshot } from "@nodetool-ai/protocol/api-schemas/app-runs.js";
 import { confineRunRequest, isRunRefusal } from "../lib/app-session-scope.js";
 import { releaseBlockedReason } from "../lib/app-deployment-service.js";
 import { autoSaveAssets, primaryTextOutputName } from "./asset-autosave.js";
 import { createRuntimeContext } from "./model-interfaces.js";
+import { createAppRunFold } from "./app-run-fold.js";
 import { formatSanitizedError, sanitizeLargeText } from "./sanitize.js";
 import { createLogger } from "@nodetool-ai/config";
 const log = createLogger("nodetool.websocket.runner");
@@ -374,7 +381,19 @@ export class JobExecutionManager {
      */
     async emitBeforeRunFailure(jobId, workflowId, err, persistJob) {
         const errorMessage = err instanceof Error ? err.message : String(err);
+        const appRun = await findAppRunByInvocation(this.session.requireUserId(), jobId);
+        if (appRun) {
+            await settleAppRun(this.session.requireUserId(), appRun.id, {
+                status: "failed",
+                error: errorMessage
+            });
+        }
         this.session.logError("beforeRunJob failed", err);
+        captureError(err, {
+            source: "job",
+            userId: this.session.userId,
+            context: { job_id: jobId, workflow_id: workflowId ?? "" }
+        });
         await this.session.send({
             type: "job_update",
             status: "failed",
@@ -589,6 +608,9 @@ export class JobExecutionManager {
      * telemetry, which is the same ledger.
      */
     async admitApplicationRun(req) {
+        if (req.app_run_id || this.session.appSession) {
+            return this.admitDurableAppRun(req);
+        }
         const applicationId = req.application_id;
         if (!applicationId)
             return true;
@@ -712,6 +734,144 @@ export class JobExecutionManager {
         }
         return true;
     }
+    async admitDurableAppRun(req) {
+        const userId = this.session.requireUserId();
+        const jobId = req.job_id ?? randomUUID();
+        req.job_id = jobId;
+        let claimedRunId = null;
+        try {
+            if (this.session.appSession && req.application_id) {
+                const release = await releasedApplicationRelease(req.application_id, userId);
+                if (!release) {
+                    throw new Error("Application release not found");
+                }
+                const instance = await ensureDefaultAppInstance({
+                    userId,
+                    applicationId: req.application_id,
+                    sourceId: `public:${req.application_id}:${release.version}`,
+                    version: release.version,
+                    name: "Visitor runs",
+                    snapshot: appRunSnapshot.parse({
+                        document: release.document,
+                        workflow_graphs: Object.fromEntries(release.workflows.flatMap((item) => item.graph ? [[item.workflowId, item.graph]] : [])),
+                        script_documents: {}
+                    })
+                });
+                const decision = await reserveAppRun({
+                    userId,
+                    instanceId: instance.id,
+                    invocationId: jobId,
+                    operationId: req.operation_id ?? "",
+                    origin: "public",
+                    estimatedUsd: this.estimateRunCost(req),
+                    requireFiniteBudget: true
+                });
+                if (!decision.allowed) {
+                    return this.refuseRun(req, jobId, ApiErrorCode.BUDGET_EXCEEDED, decision.reason);
+                }
+                req.app_run_id = decision.run.id;
+                req.instance_id = instance.id;
+            }
+            const run = req.app_run_id
+                ? await getAppRun(userId, req.app_run_id)
+                : null;
+            const instance = req.instance_id
+                ? await getAppInstance(userId, req.instance_id)
+                : null;
+            if (!run ||
+                !instance ||
+                run.instance_id !== instance.id ||
+                run.invocation_id !== jobId) {
+                return this.refuseRun(req, jobId, ApiErrorCode.NOT_FOUND, "App run not found");
+            }
+            if (run.status !== "running") {
+                return this.refuseRun(req, jobId, ApiErrorCode.ALREADY_EXISTS, "App run has already finished");
+            }
+            const operation = instance.snapshot.document.operations.find((item) => item.id === run.operation_id);
+            if (!operation ||
+                operationTarget(operation).kind !== "workflow" ||
+                (req.operation_id && req.operation_id !== run.operation_id)) {
+                return this.refuseRun(req, jobId, ApiErrorCode.INVALID_INPUT, "Operation does not match this run");
+            }
+            const target = operationTarget(operation);
+            if (target.kind !== "workflow" ||
+                !target.workflowId ||
+                (req.workflow_id && req.workflow_id !== target.workflowId)) {
+                return this.refuseRun(req, jobId, ApiErrorCode.INVALID_INPUT, "Workflow does not match this run");
+            }
+            if (run.origin !== "public") {
+                const frozen = run.snapshot?.workflow_graphs[target.workflowId];
+                if (!frozen) {
+                    return this.refuseRun(req, jobId, ApiErrorCode.INVALID_INPUT, "Run has no frozen workflow");
+                }
+                const overlays = new Map();
+                const inputs = instance.variables["__app_inputs"];
+                for (const [key, value] of Object.entries(isRecord(inputs) ? inputs : {})) {
+                    const binding = parseInputStateKey(key);
+                    if (!binding?.property || binding.operationId !== run.operation_id) {
+                        continue;
+                    }
+                    const properties = overlays.get(binding.nodeId) ?? {};
+                    properties[binding.property] = value;
+                    overlays.set(binding.nodeId, properties);
+                }
+                req.graph = {
+                    ...frozen,
+                    nodes: frozen.nodes.map((node) => {
+                        const overlay = overlays.get(node.id);
+                        if (!overlay) {
+                            return node;
+                        }
+                        const properties = isRecord(node.properties)
+                            ? node.properties
+                            : isRecord(node.data)
+                                ? node.data
+                                : {};
+                        return { ...node, properties: { ...properties, ...overlay } };
+                    })
+                };
+                await updateAppRunEstimate(userId, run.id, this.estimateRunCost(req));
+            }
+            if (!(await claimAppRun(userId, run.id, getInstanceId()))) {
+                return this.refuseRun(req, jobId, ApiErrorCode.ALREADY_EXISTS, "App run execution already started");
+            }
+            claimedRunId = run.id;
+            await setAppRunInputs(userId, run.id, req.params ?? {});
+            req.app_run_id = run.id;
+            req.instance_id = instance.id;
+            req.user_id = userId;
+            req.workflow_id = target.workflowId;
+            req.application_id = run.application_id;
+            req.application_version = run.version;
+            req.operation_id = run.operation_id;
+            return true;
+        }
+        catch (error) {
+            if (claimedRunId) {
+                try {
+                    await settleAppRun(userId, claimedRunId, {
+                        status: "failed",
+                        actualUsd: 0,
+                        knownLlmUsd: 0,
+                        error: error instanceof Error
+                            ? error.message
+                            : "App run admission failed"
+                    });
+                }
+                catch (settlementError) {
+                    if (!(settlementError instanceof AppRunError) ||
+                        settlementError.code !== "not_found") {
+                        this.session.logError("app run admission settlement failed", settlementError);
+                    }
+                }
+            }
+            this.session.logError("app run admission failed", error);
+            const code = error instanceof AppRunError && error.code === "budget_exceeded"
+                ? ApiErrorCode.BUDGET_EXCEEDED
+                : ApiErrorCode.INVALID_INPUT;
+            return this.refuseRun(req, jobId, code, error instanceof Error ? error.message : "App run admission failed");
+        }
+    }
     /**
      * Gate a run on the user's credit balance — but only the part of the run
      * that spends through NodeTool's managed provider. A graph with no
@@ -808,8 +968,15 @@ export class JobExecutionManager {
         req._accepted_at_ms ??= performance.now();
         if (!(await this.admitApplicationRun(req)))
             return;
-        if (!(await this.admitCreditRun(req)))
+        if (!(await this.admitCreditRun(req))) {
+            if (req.app_run_id) {
+                await settleAppRun(this.session.requireUserId(), req.app_run_id, {
+                    status: "failed",
+                    error: "Run refused by spending limit"
+                });
+            }
             return;
+        }
         const max = await this.deps.getMaxConcurrentJobs();
         const perWorkflowMax = await this.perWorkflowLimitFor(req);
         // Queue the run when over the global cap, or when this workflow already has
@@ -949,6 +1116,13 @@ export class JobExecutionManager {
         try {
             await this.startJobInner(req, releaseSlot);
         }
+        catch (error) {
+            if (req.app_run_id) {
+                await this.emitBeforeRunFailure(req.job_id ?? "", req.workflow_id ?? null, error, true);
+                return;
+            }
+            throw error;
+        }
         finally {
             // Safety net: if startJobInner returned/threw without registering, the
             // slot is freed so it doesn't leak and permanently shrink the cap.
@@ -1011,11 +1185,23 @@ export class JobExecutionManager {
         const workspace = this.session.workspaceResolver
             ? await this.session.workspaceResolver(workflowId ?? null, userId)
             : null;
+        const appRun = req.app_run_id
+            ? await getAppRun(userId, req.app_run_id)
+            : null;
         const context = createRuntimeContext({
             jobId,
             workflowId,
             projectId,
             userId,
+            ...(appRun && {
+                appRunContext: {
+                    userId,
+                    appRunId: appRun.id,
+                    instanceId: appRun.instance_id,
+                    traceId: appRun.trace_id,
+                    origin: appRun.origin
+                }
+            }),
             workspace,
             assetOutputMode: this.session.mode === "text" ? "data_uri" : "temp_url",
             persistOutputAssets: executionOptions.assetPersistence === "auto"
@@ -1044,6 +1230,9 @@ export class JobExecutionManager {
                 const existing = await Job.get(jobId);
                 if (existing) {
                     if (existing.status === "cancelled") {
+                        if (appRun) {
+                            await settleAppRun(userId, appRun.id, { status: "cancelled" });
+                        }
                         log.info("Skipping start of cancelled job", { jobId });
                         // Nothing was registered in activeJobs yet — free the reserved
                         // slot and promote any queued run, matching every other slot
@@ -1134,11 +1323,27 @@ export class JobExecutionManager {
         // uses: a bare throw reaches handleCommand as a generic `invalid_command`
         // the UI never associates with the job, so the run appears to spin
         // forever instead of failing with the reason.
+        let finishAppTrace;
+        try {
+            await ensureRunTraceTelemetry();
+            const incomingParent = req.traceparent ? trace.getSpanContext(propagation.extract(ROOT_CONTEXT, { traceparent: req.traceparent })) : undefined;
+            const acceptedParent = incomingParent && isSpanContextValid(incomingParent) && !this.session.appSession && appRun && incomingParent.traceId === appRun.trace_id && await getRegisteredRunTrace(userId, incomingParent.traceId, appRun.id) ? incomingParent : undefined;
+            finishAppTrace = appRun ? await openAppRunTrace(context, appRun, acceptedParent) : undefined;
+            if (executionOptions.persistence === "job") {
+                await registerWorkflowRunTrace(context, { jobId, workflowId: workflowId ?? null, inlineGraph: Boolean(req.graph), origin: this.session.appSession ? "public" : "ui" });
+            }
+        }
+        catch (error) {
+            await finishAppTrace?.("failed", error instanceof Error ? error.message : String(error));
+            await this.emitBeforeRunFailure(jobId, workflowId, error, executionOptions.persistence === "job");
+            return;
+        }
         let session;
         try {
             session = await ExecutionSession.create(sessionOptions);
         }
         catch (err) {
+            await finishAppTrace?.("failed", err instanceof Error ? err.message : String(err));
             if (!isExecutionPreflightError(err))
                 throw err;
             await this.emitBeforeRunFailure(jobId, workflowId, err, executionOptions.persistence === "job");
@@ -1163,7 +1368,9 @@ export class JobExecutionManager {
                 persistenceMs,
                 kernelStartedAt: performance.now()
             },
-            applicationId: req.application_id ?? null
+            applicationId: req.application_id ?? null,
+            ...(appRun && { appRunFold: createAppRunFold(appRun) }),
+            ...(finishAppTrace && { finishAppTrace })
         };
         // Decouple the run from this socket: from here on every frame carrying
         // this job_id is stamped with `job_seq` and buffered, so a client that
@@ -1304,6 +1511,16 @@ export class JobExecutionManager {
                     }
                     else if (active.status === "failed") {
                         job.markFailed(active.error ?? "Unknown error");
+                        // The run's error is a string by now. A stack built here would
+                        // point at this method, not at the failure, so none is sent.
+                        captureError(active.error ?? "Unknown error", {
+                            source: "job",
+                            userId: job.user_id,
+                            context: {
+                                job_id: active.jobId,
+                                workflow_id: active.workflowId ?? ""
+                            }
+                        });
                     }
                     else if (active.status === "cancelled") {
                         job.markCancelled();
@@ -1327,6 +1544,38 @@ export class JobExecutionManager {
      * the estimate standing rather than handing the spend back. Never throws.
      */
     async settleApplicationInvocation(active) {
+        await active.finishAppTrace?.(active.status === "failed" ? "failed" : active.status === "cancelled" ? "cancelled" : "completed", active.error, this.runMeasuredCost(active));
+        const identity = active.context.appRunContext;
+        if (identity) {
+            try {
+                const settlement = {
+                    status: active.status === "failed"
+                        ? "failed"
+                        : active.status === "cancelled"
+                            ? "cancelled"
+                            : "completed",
+                    error: active.context.runTraceContext?.policy.contentSuppressed && active.error ? "Execution failed" : active.error ?? null,
+                    actualUsd: this.runMeasuredCost(active),
+                    outputs: active.appRunFold?.outputs() ?? {},
+                    documents: active.context.getAppRunDocuments(),
+                    contentSuppressed: active.context.runTraceContext?.policy.contentSuppressed,
+                    secretValues: [...active.context.getResolvedSecretValues()]
+                };
+                const llmCost = active.context.getAppRunLlmCost();
+                if (llmCost !== null) {
+                    settlement.knownLlmUsd = llmCost;
+                }
+                await settleAppRun(identity.userId, identity.appRunId, settlement);
+                await reconcileAppRunCost(identity.userId, identity.appRunId);
+                return;
+            }
+            catch (error) {
+                if (!(error instanceof AppRunError) || error.code !== "not_found") {
+                    this.session.logError("app run settlement failed", error);
+                }
+            }
+            return;
+        }
         if (!active.applicationId)
             return;
         try {
@@ -1587,18 +1836,25 @@ export class JobExecutionManager {
                     continue;
                 }
                 const status = outbound.type === "job_update" ? String(outbound.status ?? "") : "";
-                const suppressProvisionalCompletion = active.requireTerminalResult &&
-                    status === "completed" &&
-                    outbound.result === undefined;
+                const suppressProvisionalCompletion = (Boolean(active.context.appRunContext) &&
+                    TERMINAL_JOB_STATUSES.includes(status)) ||
+                    (active.requireTerminalResult &&
+                        status === "completed" &&
+                        outbound.result === undefined);
                 if (!suppressProvisionalCompletion) {
                     if (outbound.type === "job_update") {
                         if (status === "running") {
                             runningSeen = true;
                         }
                         else if (TERMINAL_JOB_STATUSES.includes(status)) {
+                            if (active.context.appRunContext) {
+                                await executionSettled;
+                                await this.settleApplicationInvocation(active);
+                            }
                             await ensureRunningFrame();
                         }
                     }
+                    active.appRunFold?.fold(outbound);
                     await this.session.send(outbound);
                 }
                 if (outbound.type === "job_update" && !suppressProvisionalCompletion) {
@@ -1620,12 +1876,30 @@ export class JobExecutionManager {
         if (Object.keys(finalOutputs).length > 0) {
             finalOutputs = await this.normalizeFinalOutputs(active, finalOutputs);
         }
+        if (active.appRunFold) {
+            for (const [key, values] of Object.entries(finalOutputs)) {
+                const nodeId = this.resolveOutputNodeForKey(active, key)?.id ?? key;
+                for (const [index, value] of values.entries()) {
+                    active.appRunFold.fold({
+                        type: "output_update",
+                        job_id: active.jobId,
+                        node_id: nodeId,
+                        value,
+                        disposition: index === 0 ? "replace" : "append",
+                        done: true
+                    });
+                }
+            }
+        }
         if (active.executionOptions.eventDetail !== "terminal" &&
             !outputUpdateSeen &&
             Object.keys(finalOutputs).length > 0) {
             await this.sendOutputUpdates(active, finalOutputs);
         }
         const relayCompletedAt = performance.now();
+        if (active.context.appRunContext) {
+            await this.settleApplicationInvocation(active);
+        }
         if (!terminalSeen ||
             (!terminalWithResultSeen && Object.keys(finalOutputs).length > 0)) {
             await ensureRunningFrame();
@@ -1706,7 +1980,7 @@ export class JobExecutionManager {
             // indistinguishable from a missing one — it must be neither reported
             // nor settled below.
             const row = (await Job.get(jobId));
-            const job = row && row.user_id === (this.session.requireUserId()) ? row : null;
+            const job = row && row.user_id === this.session.requireUserId() ? row : null;
             const settled = job != null &&
                 (job.status === "completed" ||
                     job.status === "failed" ||
@@ -1781,6 +2055,11 @@ export class JobExecutionManager {
         // and tell the client it's cancelled before it ever starts.
         const queued = this.jobQueue.remove(jobId);
         if (queued) {
+            if (queued.app_run_id) {
+                await settleAppRun(this.session.requireUserId(), queued.app_run_id, {
+                    status: "cancelled"
+                });
+            }
             releaseSpend(this.session.requireUserId(), jobId);
             const cancelledWorkflowId = queued.workflow_id ?? workflowId ?? null;
             // Mark the persisted queued row cancelled so it leaves the queue in
@@ -2062,6 +2341,11 @@ export class JobExecutionManager {
             if (!queuedId)
                 continue;
             try {
+                if (queued.app_run_id) {
+                    await settleAppRun(this.session.requireUserId(), queued.app_run_id, {
+                        status: "cancelled"
+                    });
+                }
                 const job = await Job.get(queuedId);
                 if (job) {
                     job.markCancelled();
@@ -2074,6 +2358,12 @@ export class JobExecutionManager {
         }
         for (const dequeuedId of this.dequeuedJobs) {
             try {
+                const run = await findAppRunByInvocation(this.session.requireUserId(), dequeuedId);
+                if (run) {
+                    await settleAppRun(this.session.requireUserId(), run.id, {
+                        status: "cancelled"
+                    });
+                }
                 const job = await Job.get(dequeuedId);
                 if (job) {
                     job.markCancelled();