swarm-assistent/src/session.js

/**
 * Per-chat generation session — source of truth for Generate params, board, LoRAs.
 * Sparse model deltas merge into the active session; buttons always generate from it.
 */

const MAX_DATA_URL_CHARS = 350_000;

const GEN_KEYS = [
  'prompt', 'negative', 'width', 'height', 'aspect', 'steps', 'cfg', 'sigma_shift',
  'seed', 'sampler', 'scheduler', 'batch', 'checkpoint', 'loras', 'controls',
  'use_init_image', 'clear_init_image', 'init_creativity', 'denoise',
  'use_mask_image', 'clear_mask_image', 'mask_blur', 'mask_grow',
];

export function emptySession() {
  return {
    gen: {
      prompt: '',
      negative: '',
      width: null,
      height: null,
      aspect: null,
      steps: null,
      cfg: null,
      sigma_shift: null,
      seed: null,
      sampler: null,
      scheduler: null,
      batch: null,
      checkpoint: null,
      loras: [],
      controls: {},
      use_init_image: false,
      clear_init_image: false,
      init_creativity: null,
      denoise: null,
      use_mask_image: false,
      clear_mask_image: false,
      mask_blur: null,
      mask_grow: null,
    },
    board: {
      slots: [],
      selectedSlotId: 'ref1',
      genResults: [],
      selectedGenResultId: null,
      refSeq: 1,
    },
    persona: 'neutral',
    pack: 'ordinary',
    context_memory: null,
  };
}

/** Normalize boolean generate + legacy actions:["generate"]. */
export function normalizeDelta(raw) {
  if (!raw || typeof raw !== 'object') {
    return null;
  }
  const delta = { ...raw };
  const acts = Array.isArray(delta.actions) ? delta.actions.map(String) : [];
  if (delta.generate === true || acts.includes('generate')) {
    delta.generate = true;
  }
  if (typeof delta.ask === 'string') {
    delta.ask = [delta.ask];
  }
  if (!Array.isArray(delta.ask)) {
    delete delta.ask;
  } else {
    delta.ask = delta.ask.map(String).filter(Boolean);
  }
  return delta;
}

export function patchWantsGenerate(patch) {
  if (!patch || typeof patch !== 'object') {
    return false;
  }
  if (patch.generate === true) {
    return true;
  }
  const acts = Array.isArray(patch.actions) ? patch.actions.map(String) : [];
  return acts.includes('generate');
}

export function patchAskList(patch) {
  const n = normalizeDelta(patch);
  return Array.isArray(n?.ask) ? n.ask : [];
}

/**
 * Merge sparse model delta into session.gen (and pack/persona if present).
 * Does not run Generate — caller decides via patchWantsGenerate + veto.
 */
export function mergeDelta(session, rawDelta) {
  const base = session && typeof session === 'object' ? structuredCloneSession(session) : emptySession();
  const delta = normalizeDelta(rawDelta);
  if (!delta) {
    return base;
  }
  if (!base.gen) {
    base.gen = emptySession().gen;
  }
  for (const key of GEN_KEYS) {
    if (delta[key] === undefined || delta[key] === null) {
      continue;
    }
    if (key === 'loras' && Array.isArray(delta.loras)) {
      base.gen.loras = delta.loras.map((l) => ({
        name: l?.name || l,
        weight: l?.weight != null ? Number(l.weight) : 1,
        triggers: Array.isArray(l?.triggers) ? l.triggers : undefined,
        trigger_phrase: l?.trigger_phrase || undefined,
      })).filter((l) => l.name);
      continue;
    }
    if (key === 'controls' && typeof delta.controls === 'object') {
      base.gen.controls = { ...(base.gen.controls || {}), ...delta.controls };
      continue;
    }
    if (key === 'checkpoint') {
      base.gen.checkpoint = typeof delta.checkpoint === 'object'
        ? { ...delta.checkpoint }
        : { name: String(delta.checkpoint) };
      continue;
    }
    base.gen[key] = delta[key];
  }
  if (delta.images != null && delta.batch == null) {
    base.gen.batch = delta.images;
  }
  if (delta.pack) {
    base.pack = String(delta.pack);
  }
  if (delta.persona) {
    base.persona = String(delta.persona);
  }
  return base;
}

function structuredCloneSession(session) {
  try {
    return JSON.parse(JSON.stringify(session));
  } catch {
    return emptySession();
  }
}

function slimSrc(src) {
  if (!src || typeof src !== 'string') {
    return null;
  }
  const s = src.trim();
  if (!s || s.startsWith('#')) {
    return null;
  }
  if (s.startsWith('data:') && s.length > MAX_DATA_URL_CHARS) {
    return null;
  }
  return s;
}

/** Snapshot live UI + board into a persistable session object. */
export function snapshotFromLive({
  genFields,
  board,
  persona,
  pack,
  context_memory,
}) {
  const session = emptySession();
  if (genFields && typeof genFields === 'object') {
    for (const key of GEN_KEYS) {
      if (genFields[key] !== undefined) {
        session.gen[key] = genFields[key];
      }
    }
  }
  session.persona = persona || 'neutral';
  session.pack = pack || 'ordinary';
  if (context_memory && typeof context_memory === 'object') {
    session.context_memory = context_memory;
  }
  if (board && typeof board === 'object') {
    session.board = {
      slots: (board.slots || []).map((s) => ({
        id: s.id,
        type: s.type,
        label: s.label,
        src: slimSrc(s.src),
        attach: !!s.attach,
        note: s.note || null,
      })),
      selectedSlotId: board.selectedSlotId || 'ref1',
      genResults: (board.genResults || []).map((r) => ({
        id: r.id,
        label: r.label,
        src: slimSrc(r.src),
        patch: r.patch || null,
      })),
      selectedGenResultId: board.selectedGenResultId || null,
      refSeq: board.refSeq || 1,
    };
  }
  return session;
}

/** Convert legacy flat chat.params into session shape. */
export function sessionFromLegacyParams(params) {
  if (!params || typeof params !== 'object') {
    return emptySession();
  }
  if (params.gen && typeof params.gen === 'object') {
    const s = emptySession();
    s.gen = { ...s.gen, ...params.gen };
    if (params.board && typeof params.board === 'object') {
      s.board = { ...s.board, ...params.board };
    }
    s.persona = params.persona || s.persona;
    s.pack = params.pack || s.pack;
    if (params.context_memory && typeof params.context_memory === 'object') {
      s.context_memory = params.context_memory;
    }
    return s;
  }
  const s = emptySession();
  for (const key of GEN_KEYS) {
    if (params[key] !== undefined && params[key] !== null) {
      s.gen[key] = params[key];
    }
  }
  if (Array.isArray(params.loras)) {
    s.gen.loras = params.loras;
  }
  if (params.checkpoint) {
    s.gen.checkpoint = typeof params.checkpoint === 'object'
      ? params.checkpoint
      : { name: String(params.checkpoint) };
  }
  s.persona = params.persona || 'neutral';
  s.pack = params.pack || 'ordinary';
  if (params.context_memory && typeof params.context_memory === 'object') {
    s.context_memory = params.context_memory;
  }
  s.board.genResults = Array.isArray(params.genResults) ? params.genResults : [];
  s.board.selectedGenResultId = params.selectedGenResultId || null;
  if (Array.isArray(params.slots)) {
    s.board.slots = params.slots;
  }
  if (params.selectedSlotId) {
    s.board.selectedSlotId = params.selectedSlotId;
  }
  if (params.refSeq) {
    s.board.refSeq = params.refSeq;
  }
  return s;
}

/** Persist blob for AssistentSaveChat.params */
export function toPersistParams(session) {
  const s = session && typeof session === 'object' ? session : emptySession();
  const out = {
    gen: s.gen || emptySession().gen,
    board: s.board || emptySession().board,
    persona: s.persona || 'neutral',
    pack: s.pack || 'ordinary',
  };
  if (s.context_memory && typeof s.context_memory === 'object') {
    out.context_memory = s.context_memory;
  }
  return out;
}

function slimText(t, max) {
  const s = String(t || '');
  if (s.length <= max) {
    return s;
  }
  return `${s.slice(0, max)}…`;
}

/** Compact context for every LLM turn. */
export function compactContext(session, extras = {}) {
  const s = session && typeof session === 'object' ? session : emptySession();
  const g = s.gen || {};
  const board = s.board || {};
  const slots = board.slots || [];
  const genSlot = slots.find((x) => x.type === 'generate' || x.id === 'generate');
  const refs = slots.filter((x) => x.type === 'ref' || String(x.id || '').startsWith('ref'));
  return {
    session: true,
    prompt: slimText(g.prompt, extras.promptMax || 2000),
    negative: slimText(g.negative, 500),
    aspect: g.aspect || null,
    width: g.width ?? null,
    height: g.height ?? null,
    steps: g.steps ?? null,
    cfg: g.cfg ?? null,
    seed: g.seed ?? null,
    sigma_shift: g.sigma_shift ?? null,
    sampler: g.sampler || null,
    scheduler: g.scheduler || null,
    batch: g.batch ?? null,
    checkpoint: g.checkpoint?.name || g.checkpoint || null,
    selected_loras: (g.loras || []).map((l) => ({
      name: l.name || l,
      weight: l.weight != null ? l.weight : 1,
    })),
    persona: s.persona || 'neutral',
    pack: s.pack || 'ordinary',
    board: {
      has_generate: !!(genSlot?.src || (board.genResults || []).some((r) => r.src)),
      refs: refs.map((r) => ({ id: r.id, has_image: !!r.src, attach: !!r.attach })),
      gen_results: (board.genResults || []).map((r) => ({
        id: r.id,
        label: r.label,
        has_image: !!r.src,
        selected: r.id === board.selectedGenResultId,
      })),
      selected_slot: board.selectedSlotId || null,
    },
    architecture_ok: extras.architecture_ok !== false,
    ...extras.extra,
  };
}

/** Full dump for ask:settings hop. */
export function fullSettingsDump(session, extras = {}) {
  const compact = compactContext(session, extras);
  const s = session && typeof session === 'object' ? session : emptySession();
  return {
    ...compact,
    detail: 'settings',
    gen: { ...(s.gen || {}) },
    controls: s.gen?.controls || {},
    exact: extras.exact || null,
    krea_profiles: extras.kreaProfiles || null,
    session_exact: extras.sessionExact || null,
  };
}

export function resolveTurnIntent(patch, userText, { vetoFn } = {}) {
  const delta = normalizeDelta(patch) || {};
  const vetoed = typeof vetoFn === 'function' ? !!vetoFn(userText) : false;
  const generate = !vetoed && patchWantsGenerate(delta);
  const hasLook = delta.look_at != null || delta.vision_from != null || delta.vision_slots != null;
  const look = !!(hasLook && !generate && !vetoed);
  const ask = patchAskList(delta);
  return { generate, look, vetoed, ask };
}

/** Params the client fills from Exact on Generate when the LLM omits them (sparse contract). */
export const EXACT_GENERATE_PARAM_KEYS = ['steps', 'cfg', 'sigma_shift'];

/**
 * Turbo/RAW profile numbers from Exact — no sessionExact overlay.
 * profileName should already reflect the live checkpoint (turbo vs raw).
 */
export function resolveExactProfileDefaults({ exact, profiles, profileName } = {}) {
  const gen = exact?.generation && typeof exact.generation === 'object' ? exact.generation : {};
  const profile = profileName || gen.profile || 'turbo';
  const fromProfile = profiles?.[profile] && typeof profiles[profile] === 'object'
    ? profiles[profile]
    : {};
  return {
    profile,
    steps: fromProfile.steps ?? gen.steps ?? null,
    cfg: fromProfile.cfg ?? gen.cfg ?? null,
    sigma_shift: fromProfile.sigma_shift ?? gen.sigma_shift ?? null,
  };
}

/**
 * On generate, inject Exact profile steps/cfg/sigma when the patch omitted them.
 * Client is authoritative so sparse LLM deltas are safe. Honors explicit patch values
 * and sessionExact when the user asked for params this turn.
 * Returns clearSessionKeys so the UI apply path is not blocked by stale sessionExact.
 */
export function mergeExactParamsForGenerate(patch, {
  exact,
  profiles,
  profileName,
  sessionExact,
  userParamIntent,
} = {}) {
  if (!patchWantsGenerate(patch)) {
    return { patch, clearSessionKeys: [], profile: null };
  }
  const defaults = resolveExactProfileDefaults({ exact, profiles, profileName });
  const out = { ...patch };
  const clearSessionKeys = [];
  for (const key of EXACT_GENERATE_PARAM_KEYS) {
    if (out[key] != null) {
      continue;
    }
    if (userParamIntent && sessionExact?.[key] != null) {
      out[key] = sessionExact[key];
      continue;
    }
    if (defaults[key] != null) {
      out[key] = defaults[key];
      if (sessionExact?.[key] != null && String(sessionExact[key]) !== String(defaults[key])) {
        clearSessionKeys.push(key);
      }
    }
  }
  return { patch: out, clearSessionKeys, profile: defaults.profile };
}

export function attachSession(SA) {
  SA.session = {
    emptySession,
    normalizeDelta,
    mergeDelta,
    patchWantsGenerate,
    patchAskList,
    snapshotFromLive,
    sessionFromLegacyParams,
    toPersistParams,
    compactContext,
    fullSettingsDump,
    resolveTurnIntent,
    resolveExactProfileDefaults,
    mergeExactParamsForGenerate,
    EXACT_GENERATE_PARAM_KEYS,
    GEN_KEYS,
  };
}