@opengeni/core 2.6.4 → 2.7.2-canary.0

package/src/domain/scheduled-tasks.ts

@@ -54,6 +54,7 @@ import {
 } from "../session-authorization";
 import type { SessionWorkflowClient } from "../dependencies";
 import type { ObjectStorageDependency } from "../dependencies";
+import { lockActiveCustomModelForAdmission, workspaceCustomModelReference } from "../model-catalog";
 import { settingsWithEnabledCapabilityMcpServers } from "./capabilities";
 import { validateVariableSetAttachment } from "./environments";
 import {
@@ -99,6 +100,28 @@ export function scheduledTaskToolsProvided(rawPayload: unknown): boolean {
   );
 }
 
+function workspaceCustomModelCommitGuard(input: {
+  settings: Settings;
+  accountId: string;
+  workspaceId: string;
+  modelId: string;
+}): ((tx: Database) => Promise<void>) | undefined {
+  const reference = workspaceCustomModelReference(input.settings, input.modelId);
+  if (!reference) return undefined;
+  return async (tx: Database): Promise<void> => {
+    const active = await lockActiveCustomModelForAdmission(tx, {
+      accountId: input.accountId,
+      workspaceId: input.workspaceId,
+      reference,
+    });
+    if (!active) {
+      throw new HTTPException(422, {
+        message: `model is not available: ${input.modelId}`,
+      });
+    }
+  };
+}
+
 export function scheduledConnectionSurfaceEligibility(
   settings: Settings,
   target: Pick<Session, "firstPartyMcpTools" | "firstPartyMcpPermissions"> | null,
@@ -146,7 +169,7 @@ export async function createValidatedScheduledTask(input: {
       action: knowledgeAction,
     });
   }
-  const agentConfig = knowledgeAction
+  const agentConfig: ScheduledTaskAgentConfig = knowledgeAction
     ? input.payload.agentConfig
     : await validateScheduledTaskAgentConfig({
         ...input,
@@ -235,6 +258,15 @@ export async function createValidatedScheduledTask(input: {
           workspaceId: input.grant.workspaceId,
           subjectId: input.grant.subjectId,
         });
+  const beforeCreateCommit =
+    !knowledgeAction && input.payload.runMode !== "existing_session"
+      ? workspaceCustomModelCommitGuard({
+          settings: input.settings,
+          accountId: input.grant.accountId,
+          workspaceId: input.grant.workspaceId,
+          modelId: agentConfig.model ?? input.settings.openaiModel,
+        })
+      : undefined;
   return await withScheduledTaskAuthorityWriteErrors(() =>
     createScheduledTask(input.db, {
       id,
@@ -257,6 +289,7 @@ export async function createValidatedScheduledTask(input: {
       variableSetId: input.payload.variableSetId ?? null,
       rigId: input.payload.rigId ?? null,
       metadata: input.payload.metadata,
+      ...(beforeCreateCommit ? { beforeCreateCommit } : {}),
     }),
   );
 }
@@ -734,6 +767,15 @@ export async function validatedScheduledTaskUpdate(input: {
     (input.payload.metadata !== undefined &&
       !isDeepStrictEqual(input.payload.metadata, input.existing.metadata)) ||
     (input.existing.status === "paused" && input.payload.status === "active");
+  if (materialExecutionChange && nextRunMode !== "existing_session") {
+    const beforeUpdateCommit = workspaceCustomModelCommitGuard({
+      settings: input.settings,
+      accountId: input.existing.accountId,
+      workspaceId: input.existing.workspaceId,
+      modelId: nextAgentConfig.model ?? input.settings.openaiModel,
+    });
+    if (beforeUpdateCommit) update.beforeUpdateCommit = beforeUpdateCommit;
+  }
   const existingXaiAuthority = await getScheduledTaskXaiProviderAccountAuthoritySnapshot(
     input.db,
     input.existing.workspaceId,