rivet-design 0.21.9 → 0.22.0

package/dist/agent-variants/SessionStore.js

@@ -1,6 +1,7 @@
 "use strict";
 Object.defineProperty(exports, "__esModule", { value: true });
 exports.SessionStore = exports.buildImageRefs = exports.SOURCE_PLAN_ACTION_OPTIONS = exports.SOURCE_PLAN_OUTPUT_SCHEMA = exports.SCAFFOLD_LEASE_TTL_MS = exports.DEFAULT_LEASE_TTL_MS = void 0;
+const preview_routes_1 = require("../types/preview-routes");
 const variant_reference_1 = require("../types/variant-reference");
 const crypto_1 = require("crypto");
 const errors_1 = require("./errors");
@@ -8,6 +9,7 @@ const contracts_1 = require("./contracts");
 const briefContent_1 = require("./briefContent");
 const runLabel_1 = require("./runLabel");
 const runPlan_1 = require("./runPlan");
+const workItemInput_1 = require("./workItemInput");
 const MAX_SUBTITLE_WORDS = 10;
 const MAX_RELEVANT_FILES = 12;
 const capSubtitleWords = (text) => {
@@ -21,11 +23,10 @@ exports.DEFAULT_LEASE_TTL_MS = 20 * 60 * 1000;
 exports.SCAFFOLD_LEASE_TTL_MS = 30 * 60 * 1000;
 const SYSTEM_CLOCK = { now: () => Date.now() };
 function isAdoptedVariant(item) {
-    return Boolean(item.input.adoptedFrom);
+    return (0, workItemInput_1.isVariantWorkItem)(item) && Boolean(item.input.adoptedFrom);
 }
 function isForkedStaticPreviewVariant(item) {
-    return (item.kind === 'static_preview' &&
-        Boolean(item.input.sourceVariantId));
+    return item.kind === 'static_preview' && Boolean(item.input.sourceVariantId);
 }
 function fingerprint(parts) {
     return parts.map((p) => (p === undefined ? '∅' : String(p))).join('|');
@@ -761,12 +762,13 @@ class SessionStore {
         item.error = args.error;
         if (args.runtime)
             item.runtime = args.runtime;
-        if (args.status === 'succeeded' && (args.title || args.description)) {
-            const input = item.input;
+        if (args.status === 'succeeded' &&
+            (args.title || args.description) &&
+            (0, workItemInput_1.isVariantWorkItem)(item)) {
             if (args.title)
-                input.briefLabel = args.title.trim().slice(0, 80);
+                item.input.briefLabel = args.title.trim().slice(0, 80);
             if (args.description)
-                input.briefBody = args.description.trim();
+                item.input.briefBody = args.description.trim();
         }
         item.acceptedReports.push(fp);
         item.lease = undefined;
@@ -796,8 +798,8 @@ class SessionStore {
         item.lease = undefined;
         item.output = undefined;
         item.error = undefined;
-        const input = (item.input ?? {});
-        item.input = { ...input, qaRetry: args.qaRetry };
+        if ((0, workItemInput_1.isVariantWorkItem)(item))
+            item.input = { ...item.input, qaRetry: args.qaRetry };
         const result = {
             stage: session.stage,
             progress: this.progressOf(session),
@@ -817,8 +819,7 @@ class SessionStore {
         const activeRefinement = [...session.workItems.values()].find((item) => {
             if (item.kind !== 'refine_variant')
                 return false;
-            const input = item.input;
-            return input.variantId === args.variantId && item.status === 'pending';
+            return (item.input.variantId === args.variantId && item.status === 'pending');
         });
         if (activeRefinement) {
             throw new errors_1.AgentVariantsError('PENDING_CHANGE_CONFLICT', 'A refinement is already queued for this variant');
@@ -846,7 +847,7 @@ class SessionStore {
                 ...(args.currentHtml ? { currentHtml: args.currentHtml } : {}),
                 ...(args.targetElement ? { targetElement: args.targetElement } : {}),
                 ...(args.requestId ? { requestId: args.requestId } : {}),
-                projectContext: targetInput.projectContext ?? session.projectContext,
+                projectContext: targetInput.projectContext,
                 ...(targetInput.context ? { context: targetInput.context } : {}),
                 ...(relevantFiles.length > 0 ? { relevantFiles } : {}),
             },
@@ -869,8 +870,7 @@ class SessionStore {
         const activeRefinement = [...session.workItems.values()].find((item) => {
             if (item.kind !== 'refine_variant')
                 return false;
-            const input = item.input;
-            return (input.variantId === args.sourceVariantId &&
+            return (item.input.variantId === args.sourceVariantId &&
                 (item.status === 'pending' || item.status === 'running'));
         });
         if (activeRefinement) {
@@ -878,6 +878,7 @@ class SessionStore {
         }
         const sourceInput = source.input;
         const relevantFiles = mergeRelevantFiles(sourceInput.relevantFiles, args.relevantFiles);
+        const previewRoutes = (0, preview_routes_1.mergePreviewRoutes)(args.previewRoutes, sourceInput.previewRoutes);
         const seedHtml = args.currentHtml ??
             source.output?.html;
         const briefs = args.briefs?.map(briefContent_1.normalizeBriefContent);
@@ -925,6 +926,7 @@ class SessionStore {
                     briefId: `${sourceInput.briefId ?? 'v'}-vary-${(0, crypto_1.randomUUID)().slice(0, 8)}`,
                     briefLabel: brief?.label ?? `${labelBase}${suffix}`,
                     briefBody: brief?.body ?? prettyInstruction,
+                    ...(previewRoutes.length ? { previewRoutes } : {}),
                     references: (0, variant_reference_1.mergeVariantReferences)(sourceInput.references, context?.references),
                     ...(args.runLabel ? { runLabel: args.runLabel } : {}),
                     ...(brief?.visualReferenceUrl
@@ -943,7 +945,7 @@ class SessionStore {
                     ...(args.targetElement ? { targetElement: args.targetElement } : {}),
                     ...(args.requestId ? { requestId: args.requestId } : {}),
                     ...(args.fidelity ? { fidelity: args.fidelity } : {}),
-                    projectContext: sourceInput.projectContext ?? session.projectContext,
+                    projectContext: sourceInput.projectContext,
                     ...(context ? { context } : {}),
                     ...(relevantFiles.length > 0 ? { relevantFiles } : {}),
                 },
@@ -981,16 +983,15 @@ class SessionStore {
             const activeRefinement = [...session.workItems.values()].find((item) => {
                 if (item.kind !== 'refine_variant')
                     return false;
-                const input = item.input;
-                return (input.variantId === args.sourceVariantId &&
+                return (item.input.variantId === args.sourceVariantId &&
                     (item.status === 'pending' || item.status === 'running'));
             });
             if (activeRefinement) {
                 throw new errors_1.AgentVariantsError('PENDING_CHANGE_CONFLICT', 'Cannot vary a variant while an in-place refinement is pending or running');
             }
         }
-        const sourceInput = (source?.input ?? {});
-        const relevantFiles = mergeRelevantFiles(sourceInput.relevantFiles, args.relevantFiles);
+        const sourceInput = source && (0, workItemInput_1.isVariantWorkItem)(source) ? source.input : undefined;
+        const relevantFiles = mergeRelevantFiles(sourceInput?.relevantFiles, args.relevantFiles);
         const createdWorkItemIds = [];
         for (const [batchSlotIndex, brief] of args.briefs.entries()) {
             const context = resolveVariantWorkItemContext({
@@ -1012,7 +1013,7 @@ class SessionStore {
                     briefId: brief.briefId,
                     briefLabel: brief.label,
                     briefBody: brief.body,
-                    references: (0, variant_reference_1.mergeVariantReferences)(sourceInput.references, context?.references),
+                    references: (0, variant_reference_1.mergeVariantReferences)(sourceInput?.references, context?.references),
                     instruction: args.instruction,
                     guidance: args.guidance ?? [],
                     ...(brief.notes ? { briefNotes: brief.notes } : {}),
@@ -1021,10 +1022,10 @@ class SessionStore {
                     ...(args.sourceVariantId
                         ? { sourceVariantId: args.sourceVariantId }
                         : {}),
-                    ...(source && sourceInput.briefLabel
+                    ...(sourceInput?.briefLabel
                         ? { sourceLabel: sourceInput.briefLabel }
                         : {}),
-                    ...(source && sourceInput.briefBody
+                    ...(sourceInput?.briefBody
                         ? { sourceBody: sourceInput.briefBody }
                         : {}),
                     target: session.target,
@@ -1050,10 +1051,9 @@ class SessionStore {
         }
         const existing = this.findAdoptedVariant(session, args.source);
         if (existing) {
-            const input = existing.input;
             return {
                 workItemId: existing.id,
-                briefId: input.briefId ?? '',
+                briefId: existing.input.briefId ?? '',
                 created: false,
             };
         }
@@ -1074,6 +1074,7 @@ class SessionStore {
                 ...(args.runLabel ? { runLabel: args.runLabel } : {}),
                 adoptedFrom: { ...args.source },
                 references: (0, variant_reference_1.mergeVariantReferences)(args.references),
+                previewRoutes: args.previewRoutes,
                 projectContext: session.projectContext,
             },
             acceptedReports: [],
@@ -1106,6 +1107,7 @@ class SessionStore {
             label: args.label,
             description: args.description,
             references: args.references,
+            previewRoutes: args.previewRoutes,
             runLabel: args.runLabel,
         });
         return {
@@ -1116,10 +1118,9 @@ class SessionStore {
     }
     findAdoptedVariant(session, source) {
         for (const item of session.workItems.values()) {
-            if (item.removed)
+            if (item.removed || !(0, workItemInput_1.isVariantWorkItem)(item))
                 continue;
-            const adoptedFrom = item.input
-                .adoptedFrom;
+            const { adoptedFrom } = item.input;
             if (adoptedFrom &&
                 adoptedFrom.sessionId === source.sessionId &&
                 adoptedFrom.variantId === source.variantId) {
@@ -1128,29 +1129,41 @@ class SessionStore {
         }
         return undefined;
     }
+    assignVariantPreviewRoutes(args) {
+        const session = this.requireSession(args.sessionId);
+        const source = args.sourceVariantId
+            ? this.requireWorkItem(session, args.sourceVariantId)
+            : undefined;
+        const merged = (0, preview_routes_1.mergePreviewRoutes)(args.previewRoutes, source && (0, workItemInput_1.isVariantWorkItem)(source)
+            ? source.input.previewRoutes
+            : undefined);
+        if (!merged.length)
+            return;
+        const previewRoutes = preview_routes_1.previewRoutesSchema.parse(merged);
+        for (const workItemId of args.workItemIds) {
+            const item = this.requireWorkItem(session, workItemId);
+            if ((0, workItemInput_1.isVariantWorkItem)(item))
+                item.input = { ...item.input, previewRoutes: [...previewRoutes] };
+        }
+    }
     assignVariantRequestId(sessionId, workItemIds, requestId) {
         const session = this.requireSession(sessionId);
         for (const id of workItemIds) {
             const item = session.workItems.get(id);
-            if (!item)
-                continue;
-            const input = (item.input ?? {});
-            item.input = { ...input, requestId };
+            if (item && (0, workItemInput_1.isVariantWorkItem)(item))
+                item.input = { ...item.input, requestId };
         }
     }
     getPendingRefinement(sessionId, variantId) {
-        return [...this.requireSession(sessionId).workItems.values()].find((item) => {
-            if (item.kind !== 'refine_variant')
-                return false;
-            const input = item.input;
-            return input.variantId === variantId && item.status === 'pending';
-        });
+        return [...this.requireSession(sessionId).workItems.values()].find((item) => item.kind === 'refine_variant' &&
+            item.input.variantId === variantId &&
+            item.status === 'pending');
     }
     appendPendingRefineVariantInstruction(args) {
         const refineItem = this.getPendingRefinement(args.sessionId, args.variantId);
         if (!refineItem)
             return false;
-        const input = refineItem.input;
+        const { input } = refineItem;
         input.guidance = args.guidance ?? [];
         input.requestId = args.requestId;
         input.instruction =
@@ -1170,7 +1183,7 @@ class SessionStore {
         const runningRefinement = [...session.workItems.values()].find((item) => {
             if (item.kind !== 'refine_variant')
                 return false;
-            const input = item.input;
+            const { input } = item;
             return input.variantId === args.variantId && item.status === 'running';
         });
         if (runningRefinement) {
@@ -1179,7 +1192,7 @@ class SessionStore {
         const refineItem = [...session.workItems.values()].find((item) => {
             if (item.kind !== 'refine_variant')
                 return false;
-            const input = item.input;
+            const { input } = item;
             return input.variantId === args.variantId && item.status === 'pending';
         });
         if (!refineItem) {
@@ -1277,7 +1290,7 @@ class SessionStore {
         for (const other of session.workItems.values()) {
             if (other.kind !== 'refine_variant')
                 continue;
-            const input = other.input;
+            const { input } = other;
             if (input.variantId !== item.id)
                 continue;
             if (other.status === 'pending' || other.status === 'running') {
@@ -1296,8 +1309,8 @@ class SessionStore {
     renameWorkItem(args) {
         const session = this.requireSession(args.sessionId);
         const item = this.requireWorkItem(session, args.workItemId);
-        const input = item.input;
-        input.briefLabel = args.label;
+        if ((0, workItemInput_1.isVariantWorkItem)(item))
+            item.input.briefLabel = args.label;
         return {
             workItemId: item.id,
             label: args.label,
@@ -1310,9 +1323,9 @@ class SessionStore {
         const workItemIds = [];
         for (const id of session.workItemOrder) {
             const item = session.workItems.get(id);
-            if (!item || item.removed)
+            if (!item || item.removed || !(0, workItemInput_1.isVariantWorkItem)(item))
                 continue;
-            const input = item.input;
+            const { input } = item;
             if (input.requestId !== args.requestId)
                 continue;
             const effective = input.runLabel ?? sessionFallback;
@@ -1379,12 +1392,7 @@ class SessionStore {
             const { lease } = item;
             if (!lease || item.status !== 'running')
                 return [];
-            const input = item.input && typeof item.input === 'object'
-                ? item.input
-                : {};
-            const variantId = item.kind === 'refine_variant' && typeof input.variantId === 'string'
-                ? input.variantId
-                : item.id;
+            const variantId = (item.kind === 'refine_variant' && item.input.variantId) || item.id;
             return {
                 workId: item.id,
                 variantId,
@@ -1410,14 +1418,25 @@ class SessionStore {
     hasSession(sessionId) {
         return this.sessions.has(sessionId);
     }
-    getWorkItemInput(sessionId, workItemId) {
+    getVariantWorkItemInput(sessionId, workItemId) {
         const session = this.requireSession(sessionId);
         const item = this.requireWorkItem(session, workItemId);
-        return item.input;
+        return (0, workItemInput_1.isVariantWorkItem)(item)
+            ? (0, workItemInput_1.parseVariantWorkItemInput)(item)
+            : undefined;
+    }
+    requireVariantWorkItemInput(args) {
+        const session = this.requireSession(args.sessionId);
+        const item = this.requireWorkItem(session, args.workItemId);
+        const kind = args.kinds.find((candidate) => candidate === item.kind);
+        if (!kind) {
+            throw new errors_1.AgentVariantsError('WORK_ITEM_NOT_FOUND', `Work item ${item.id} is a ${item.kind} work item, not ${args.kinds.join(' or ')}`);
+        }
+        return (0, workItemInput_1.parseVariantWorkItemInput)({ id: item.id, kind, input: item.input });
     }
     getWorkItemIdsForRequest(sessionId, requestId) {
         return [...this.requireSession(sessionId).workItems.values()]
-            .filter((item) => item.input.requestId === requestId)
+            .filter((item) => (0, workItemInput_1.isVariantWorkItem)(item) && item.input.requestId === requestId)
             .map((item) => item.id);
     }
     getWorkItemOutput(sessionId, workItemId) {
@@ -1441,7 +1460,7 @@ class SessionStore {
     }
     isVaryForkWorkItem(sessionId, workItemId) {
         const item = this.sessions.get(sessionId)?.workItems.get(workItemId);
-        return Boolean(item && item.input.sourceVariantId);
+        return Boolean(item && (0, workItemInput_1.isVariantWorkItem)(item) && item.input.sourceVariantId);
     }
     isWorkItemAborted(sessionId, workItemId) {
         const session = this.sessions.get(sessionId);
@@ -1529,8 +1548,7 @@ class SessionStore {
         return out;
     }
     hasLeasableCodeGenForStaticPreview(session, staticPreview) {
-        const staticInput = staticPreview.input;
-        if (!staticInput.briefId)
+        if (!(0, workItemInput_1.isVariantWorkItem)(staticPreview) || !staticPreview.input.briefId)
             return false;
         for (const candidate of session.workItems.values()) {
             if (candidate.kind !== 'code_gen')
@@ -1539,8 +1557,7 @@ class SessionStore {
                 continue;
             if (!this.dependenciesMet(session, candidate))
                 continue;
-            const codeInput = candidate.input;
-            if (codeInput.briefId === staticInput.briefId)
+            if (candidate.input.briefId === staticPreview.input.briefId)
                 return true;
         }
         return false;
@@ -1639,14 +1656,11 @@ class SessionStore {
             successCount: items.filter((i) => i.status === 'succeeded').length,
             failureCount: items.filter((i) => i.status === 'failed').length,
             cancelledCount: items.filter((i) => i.status === 'cancelled').length,
-            variants: items.map((i) => {
-                const input = i.input;
-                return {
-                    workItemId: i.id,
-                    status: i.status,
-                    label: input.briefLabel,
-                };
-            }),
+            variants: items.map((i) => ({
+                workItemId: i.id,
+                status: i.status,
+                label: (0, workItemInput_1.isVariantWorkItem)(i) ? i.input.briefLabel : undefined,
+            })),
         };
     }
     refinementSnapshotPriority(item) {
@@ -1662,7 +1676,7 @@ class SessionStore {
         for (const item of session.workItems.values()) {
             if (item.kind !== 'refine_variant')
                 continue;
-            const input = item.input;
+            const { input } = item;
             if (!input.variantId)
                 continue;
             const existing = refinementsByVariantId.get(input.variantId);
@@ -1676,16 +1690,19 @@ class SessionStore {
             .map((id) => session.workItems.get(id))
             .filter((i) => !!i &&
             !i.removed &&
+            (0, workItemInput_1.isVariantWorkItem)(i) &&
             (i.kind === kind ||
                 isAdoptedVariant(i) ||
                 isForkedStaticPreviewVariant(i)))
             .map((item) => {
-            const input = item.input;
+            const { input } = item;
             const refineItem = refinementsByVariantId.get(item.id);
-            const refineRequestId = refineItem?.input
-                ?.requestId;
+            const refineRequestId = refineItem?.input.requestId;
             return {
                 workItemId: item.id,
+                ...(input.previewRoutes
+                    ? { previewRoutes: [...input.previewRoutes] }
+                    : {}),
                 briefId: input.briefId ?? '',
                 label: input.briefLabel ?? 'Variant',
                 references: input.references,