src/model-picker/apply-session-model-selection.ts

import { resolveAgentDir, type AgentModelPrimaryWriteTarget } from "../agents/agent-scope.js";
import type { ModelCatalogEntry } from "../agents/model-catalog.js";
import { modelKey, resolveDefaultModelForAgent } from "../agents/model-selection.js";
import {
  createModelVisibilityPolicy,
  type ModelVisibilityPolicy,
} from "../agents/model-visibility-policy.js";
import { resolveContextConfigProviderForRuntime } from "../agents/openai-routing.js";
import {
  persistStickyModelSelectionBestEffort,
  type StickyModelSelectionDispatchOutcome,
} from "../agents/sticky-model-selection.js";
import { resolveEffectiveAgentRuntime } from "../agents/thinking-runtime.js";
import { applyModelRuntimeDirective } from "../auto-reply/reply/directive-handling.model-runtime.js";
import {
  prepareModelSelectionRuntime,
  findSelectedCatalogEntry,
} from "../auto-reply/reply/model-runtime-normalization.js";
import { resolveContextTokens } from "../auto-reply/reply/model-selection-context.js";
import { refreshQueuedFollowupSession } from "../auto-reply/reply/queue.js";
import { persistReplySessionEntry } from "../auto-reply/reply/session-entry-persistence.js";
import { resolveSupportedThinkingLevel } from "../auto-reply/thinking.js";
import type { ThinkLevel } from "../auto-reply/thinking.shared.js";
import { resolveCollapsedSessionAuthPinSource } from "../config/sessions/auth-profile-override-provenance.js";
import {
  adoptPersistedSessionSnapshot,
  SESSION_MODEL_OVERRIDE_TRANSACTION_FIELDS,
  sessionModelOverrideChangesApplied,
} from "../config/sessions/session-snapshot-merge.js";
import type { InternalSessionEntry as SessionEntry } from "../config/sessions/types.js";
import type { OpenClawConfig } from "../config/types.openclaw.js";
import { triggerSessionPatchHook } from "../gateway/session-patch-hooks.js";
import { resolveSessionWorkerPlacementContext } from "../gateway/session-worker-placement-context.js";
import { resolveWorkerPlacementSessionRuntimeCapabilities } from "../gateway/worker-environments/placement-session-runtime.js";
import { enqueueSystemEvent } from "../infra/system-events.js";
import { applyModelOverrideWithAuthProfileCompatibility } from "../sessions/auth-profile-preservation.js";
import {
  isModelSelectionLocked,
  MODEL_SELECTION_LOCKED_MESSAGE,
} from "../sessions/model-overrides.js";
import { emitSessionLifecycleEvent } from "../sessions/session-lifecycle-events.js";

export type SessionModelSelectionRequest = {
  provider: string;
  model: string;
  isDefault: boolean;
  resetToDefault?: true;
  alias?: string;
  profileOverride?: string;
  runtime: { kind: "unchanged" } | { kind: "clear" } | { kind: "set"; runtime: string };
};

export type ApplySessionModelSelectionParams = {
  cfg: OpenClawConfig;
  agentId: string;
  sessionKey: string;
  storePath?: string;
  sessionEntry: SessionEntry;
  sessionStore: Record<string, SessionEntry>;
  allowCreate?: boolean;
  defaultProvider: string;
  defaultModel: string;
  currentProvider: string;
  currentModel: string;
  modelPolicy?: Omit<ModelVisibilityPolicy, "catalog">;
  modelCatalog: readonly ModelCatalogEntry[];
  thinkingCatalog?: readonly ModelCatalogEntry[];
  canPersistStickyModelSelection?: boolean;
  stickyModelSelectionTarget?: AgentModelPrimaryWriteTarget;
  validateAuthProfileSelection?: () => string | undefined;
  request: SessionModelSelectionRequest;
  /** Raw directive text used only by the existing session patch hook. */
  patchModel?: string;
  markLiveSwitchPending: true;
};

export type ApplySessionModelSelectionResult =
  | {
      status: "applied";
      provider: string;
      model: string;
      effectiveModelRef: string;
      agentRuntime: string;
      changed: boolean;
      contextTokens: number;
      configuredDefaultUpdate?: StickyModelSelectionDispatchOutcome;
      runtimeChange?: { kind: "clear" } | { kind: "set"; runtime: string };
      thinkingRemap?: {
        from: ThinkLevel;
        to: ThinkLevel;
        provider: string;
        model: string;
      };
    }
  | {
      status: "rejected";
      reason: "locked" | "not-allowed" | "invalid-runtime" | "unknown-provider";
      message: string;
    }
  | { status: "conflict"; message: string };

type AppliedRuntimeDirective = Exclude<
  Parameters<typeof applyModelRuntimeDirective>[1],
  { kind: "invalid" }
>;

type ApplySessionModelSelectionToEntryResult = {
  changed: boolean;
  runtimeChange?: { kind: "clear" } | { kind: "set"; runtime: string };
};

/** Applies the model transaction field family to one caller-owned snapshot. */
function applySessionModelSelectionToEntry(params: {
  cfg: OpenClawConfig;
  agentDir: string;
  entry: SessionEntry;
  currentProvider: string;
  request: SessionModelSelectionRequest;
  runtime: AppliedRuntimeDirective;
  markLiveSwitchPending?: boolean;
}): ApplySessionModelSelectionToEntryResult {
  const modelChange = applyModelOverrideWithAuthProfileCompatibility({
    cfg: params.cfg,
    agentDir: params.agentDir,
    entry: params.entry,
    currentProvider: params.currentProvider,
    selection: params.request,
    explicitDefaultSelection: params.request.isDefault,
    profileOverride: params.request.profileOverride,
    markLiveSwitchPending: params.markLiveSwitchPending,
  });
  const runtimeChange = applyModelRuntimeDirective(params.entry, params.runtime);
  return {
    changed: modelChange.updated || runtimeChange.updated,
    ...(params.runtime.kind === "clear" || params.runtime.kind === "set"
      ? { runtimeChange: params.runtime }
      : {}),
  };
}

function formatModelSwitchEvent(provider: string, model: string, alias?: string): string {
  const label = `${provider}/${model}`;
  return alias ? `Model switched to ${alias} (${label}).` : `Model switched to ${label}.`;
}

function rejectNotAllowed(provider: string, model: string): ApplySessionModelSelectionResult {
  return {
    status: "rejected",
    reason: "not-allowed",
    message: `Model ${provider}/${model} is not available for this agent.`,
  };
}

/**
 * Rejects a model selection when the candidate runtime is incompatible with an
 * active cloud-worker placement. Mirrors the sessions.patch guard so directive
 * model changes are validated before they persist.
 */
function resolveActivePlacementModelSelectionError(params: {
  cfg: OpenClawConfig;
  agentId: string;
  sessionKey: string;
  entry: SessionEntry;
}): string | undefined {
  const sessionId = params.entry.sessionId;
  if (!sessionId) {
    return undefined;
  }
  const placementService = resolveSessionWorkerPlacementContext().workerSessionPlacementService;
  const placement = placementService?.getMany([sessionId]).get(sessionId);
  if (!placement || placement.state === "local") {
    return undefined;
  }
  const { executionMode } = resolveWorkerPlacementSessionRuntimeCapabilities({
    cfg: params.cfg,
    entry: params.entry,
    agentId: params.agentId,
    sessionKey: params.sessionKey,
  });
  if (executionMode === placement.executionMode) {
    return undefined;
  }
  return executionMode
    ? `Session cannot change cloud placement execution mode while placement is ${placement.state}.`
    : `Session cannot select a runtime without cloud placement support while cloud worker placement is ${placement.state}.`;
}

/** Applies one validated picker selection to the authoritative live session. */
export async function applySessionModelSelection(
  params: ApplySessionModelSelectionParams,
): Promise<ApplySessionModelSelectionResult> {
  const startingStoreEntry = params.sessionStore[params.sessionKey];
  const startingEntry = params.storePath
    ? params.sessionEntry
    : (startingStoreEntry ?? params.sessionEntry);
  const initialEntry = { ...startingEntry };
  if (isModelSelectionLocked(startingEntry)) {
    return { status: "rejected", reason: "locked", message: MODEL_SELECTION_LOCKED_MESSAGE };
  }

  const resetToDefault = params.request.resetToDefault === true;
  const selectedRef = resetToDefault
    ? resolveDefaultModelForAgent({ cfg: params.cfg, agentId: params.agentId })
    : params.request;
  const normalizedModelKey = modelKey(selectedRef.provider, selectedRef.model);
  const request: SessionModelSelectionRequest = {
    ...params.request,
    provider: selectedRef.provider,
    model: selectedRef.model,
    isDefault:
      resetToDefault ||
      normalizedModelKey === modelKey(params.defaultProvider, params.defaultModel),
  };
  const policy =
    params.modelPolicy ??
    createModelVisibilityPolicy({
      cfg: params.cfg,
      catalog: [...params.modelCatalog],
      defaultProvider: params.defaultProvider,
      defaultModel: params.defaultModel,
      agentId: params.agentId,
    });
  if (!resetToDefault && !policy.allows(request)) {
    return rejectNotAllowed(request.provider, request.model);
  }

  const prepared = await prepareModelSelectionRuntime({
    cfg: params.cfg,
    agentId: params.agentId,
    workspaceDir: startingEntry.spawnedWorkspaceDir,
    sessionEntry: request.profileOverride
      ? {
          ...startingEntry,
          providerOverride: request.provider,
          modelProvider: request.provider,
          authProfileOverride: request.profileOverride,
          authProfileOverrideSource: "user",
        }
      : startingEntry,
    provider: request.provider,
    model: request.model,
    catalog: params.thinkingCatalog ?? params.modelCatalog,
    rawRuntime:
      request.runtime.kind === "set"
        ? request.runtime.runtime
        : request.runtime.kind === "clear"
          ? "default"
          : undefined,
  });
  if (prepared.status === "rejected") {
    return prepared;
  }
  const validateSelection = () =>
    params.validateAuthProfileSelection?.() ?? prepared.validateRuntimeSelection?.();
  const authProfileError = validateSelection();
  if (authProfileError) {
    return { status: "rejected", reason: "not-allowed", message: authProfileError };
  }
  // Metadata preparation can yield. Memory-only sessions need the same lock and
  // replacement fence that persisted sessions enforce in their atomic write.
  const currentEntry = params.storePath
    ? startingEntry
    : (params.sessionStore[params.sessionKey] ?? params.sessionEntry);
  if (isModelSelectionLocked(currentEntry)) {
    return { status: "rejected", reason: "locked", message: MODEL_SELECTION_LOCKED_MESSAGE };
  }
  if (
    !params.storePath &&
    (params.sessionStore[params.sessionKey] !== startingStoreEntry ||
      currentEntry.sessionId !== initialEntry.sessionId)
  ) {
    return {
      status: "conflict",
      message: "Model change was not applied because the session changed. Retry.",
    };
  }
  const runtime = prepared.runtime;
  const thinkingCatalog = prepared.catalog;
  const selectedCatalogEntry = findSelectedCatalogEntry({ catalog: thinkingCatalog, ...request });
  const nextEntry = { ...startingEntry };
  const applied = applySessionModelSelectionToEntry({
    cfg: params.cfg,
    agentDir: resolveAgentDir(params.cfg, params.agentId),
    entry: nextEntry,
    currentProvider: params.currentProvider,
    request,
    runtime,
    markLiveSwitchPending: params.markLiveSwitchPending,
  });
  const thinkingRuntime = resolveEffectiveAgentRuntime({
    cfg: params.cfg,
    provider: request.provider,
    modelId: request.model,
    modelApi: selectedCatalogEntry?.api,
    modelBaseUrl: selectedCatalogEntry?.baseUrl,
    agentId: params.agentId,
    sessionKey: params.sessionKey,
    sessionEntry: nextEntry,
  });
  const currentThinkingLevel = nextEntry.thinkingLevel as ThinkLevel | undefined;
  let thinkingRemap: Extract<
    ApplySessionModelSelectionResult,
    { status: "applied" }
  >["thinkingRemap"];
  if (currentThinkingLevel) {
    const remapped = resolveSupportedThinkingLevel({
      provider: request.provider,
      model: request.model,
      level: currentThinkingLevel,
      catalog: [...thinkingCatalog],
      agentRuntime: thinkingRuntime,
    });
    if (remapped !== currentThinkingLevel) {
      nextEntry.thinkingLevel = remapped;
      thinkingRemap = {
        from: currentThinkingLevel,
        to: remapped,
        provider: request.provider,
        model: request.model,
      };
    }
  }
  const placementError = resolveActivePlacementModelSelectionError({
    cfg: params.cfg,
    agentId: params.agentId,
    sessionKey: params.sessionKey,
    entry: nextEntry,
  });
  if (placementError) {
    return { status: "rejected", reason: "invalid-runtime", message: placementError };
  }
  // An explicit selection retains the existing persistence and conflict semantics even when idempotent.
  nextEntry.updatedAt = Date.now();
  let persistedEntry: SessionEntry;
  // The pre-persistence read above can be overtaken by placement activation before the
  // durable write commits. Revalidate placement inside the synchronous commit boundary so an
  // override that became incompatible during that window is rejected without mutating state.
  const validateCommit = () =>
    validateSelection() ??
    resolveActivePlacementModelSelectionError({
      cfg: params.cfg,
      agentId: params.agentId,
      sessionKey: params.sessionKey,
      entry: nextEntry,
    });
  if (params.storePath) {
    const persistence = await persistReplySessionEntry({
      storePath: params.storePath,
      sessionKey: params.sessionKey,
      initialEntry,
      entry: nextEntry,
      allowCreate: params.allowCreate,
      reassertLiveModelSwitchPending: applied.changed && nextEntry.liveModelSwitchPending === true,
      requireModelSelectionUnlocked: true,
      touchedFields: SESSION_MODEL_OVERRIDE_TRANSACTION_FIELDS,
      validateCommit,
    });
    if (persistence.entry) {
      params.sessionStore[params.sessionKey] = persistence.entry;
      adoptPersistedSessionSnapshot(params.sessionEntry, persistence.entry);
    }
    if (persistence.status === "model-selection-locked") {
      return { status: "rejected", reason: "locked", message: MODEL_SELECTION_LOCKED_MESSAGE };
    }
    if (persistence.status === "commit-rejected") {
      return { status: "rejected", reason: "not-allowed", message: persistence.error };
    }
    if (
      persistence.status !== "current" ||
      !sessionModelOverrideChangesApplied({
        initial: initialEntry,
        next: nextEntry,
        current: persistence.entry,
        reassertLiveModelSwitchPending:
          applied.changed && nextEntry.liveModelSwitchPending === true,
      })
    ) {
      return {
        status: "conflict",
        message: "Model change was not applied because the session changed. Retry.",
      };
    }
    persistedEntry = persistence.entry;
  } else {
    adoptPersistedSessionSnapshot(params.sessionEntry, nextEntry);
    params.sessionStore[params.sessionKey] = params.sessionEntry;
    persistedEntry = params.sessionEntry;
  }

  const agentRuntime = resolveEffectiveAgentRuntime({
    cfg: params.cfg,
    provider: request.provider,
    modelId: request.model,
    modelApi: selectedCatalogEntry?.api,
    modelBaseUrl: selectedCatalogEntry?.baseUrl,
    agentId: params.agentId,
    sessionKey: params.sessionKey,
    sessionEntry: persistedEntry,
  });

  const provider = request.provider;
  const model = request.model;
  const effectiveModelRef = `${provider}/${model}`;
  const changed = applied.changed || thinkingRemap !== undefined;
  const configuredDefaultUpdate =
    params.canPersistStickyModelSelection === true &&
    (!request.isDefault || params.stickyModelSelectionTarget)
      ? persistStickyModelSelectionBestEffort({
          agentId: params.agentId,
          model: effectiveModelRef,
          // The shipped SDK opt-in resolves its effective layer inside the config mutation.
          // Ordinary chat callers supply an authorized target or leave persistence disabled.
          target: params.stickyModelSelectionTarget ?? "effective",
        })
      : undefined;
  if (changed) {
    emitSessionLifecycleEvent({
      sessionKey: params.sessionKey,
      agentId: params.agentId,
      reason: "patch",
    });
    triggerSessionPatchHook({
      cfg: params.cfg,
      sessionEntry: persistedEntry,
      sessionKey: params.sessionKey,
      patch: { key: params.sessionKey, model: params.patchModel ?? effectiveModelRef },
    });
    refreshQueuedFollowupSession({
      key: params.sessionKey,
      nextProvider: provider,
      nextModel: model,
      nextRouteResolution: "resolved",
      nextModelOverrideSource: request.isDefault ? undefined : "user",
      nextAuthProfileId: persistedEntry.authProfileOverride,
      nextAuthProfileIdSource: resolveCollapsedSessionAuthPinSource(persistedEntry),
      nextThinking: {
        level: persistedEntry.thinkingLevel,
        catalog: [...thinkingCatalog],
        agentRuntime,
      },
    });
  }

  if (`${params.currentProvider}/${params.currentModel}` !== effectiveModelRef) {
    enqueueSystemEvent(formatModelSwitchEvent(provider, model, request.alias), {
      sessionKey: params.sessionKey,
      contextKey: `model:${effectiveModelRef}`,
    });
  }

  const contextProvider = resolveContextConfigProviderForRuntime({
    provider,
    runtimeId: agentRuntime,
    config: params.cfg,
  });
  return {
    status: "applied",
    provider,
    model,
    effectiveModelRef,
    agentRuntime,
    changed,
    contextTokens: resolveContextTokens({
      cfg: params.cfg,
      provider: contextProvider,
      model,
      modelContextWindow: selectedCatalogEntry?.contextWindow,
      modelContextTokens: selectedCatalogEntry?.contextTokens,
    }),
    ...(configuredDefaultUpdate ? { configuredDefaultUpdate } : {}),
    ...(applied.runtimeChange ? { runtimeChange: applied.runtimeChange } : {}),
    ...(thinkingRemap ? { thinkingRemap } : {}),
  };
}