dify/web/features/agent-v2/agent-composer/store-modules/env.ts

import type { EnvScope, EnvVariable } from '../form-state'
import type { DraftFieldUpdate } from './utils'
import { atom } from 'jotai'
import { agentComposerDraftAtom } from '../store'
import { resolveDraftFieldUpdate } from './utils'

export const agentComposerEnvVariablesAtom = atom(
  (get) => get(agentComposerDraftAtom).envVariables,
  (get, set, envVariablesUpdate: DraftFieldUpdate<EnvVariable[]>) => {
    const draft = get(agentComposerDraftAtom)

    set(agentComposerDraftAtom, {
      ...draft,
      envVariables: resolveDraftFieldUpdate(draft.envVariables, envVariablesUpdate),
    })
  },
)

const updateEnvVariable = (
  envVariables: EnvVariable[],
  starterVariable: EnvVariable,
  id: string,
  updater: (variable: EnvVariable) => EnvVariable,
) => {
  const existingVariable = envVariables.find((variable) => variable.id === id)

  if (existingVariable) {
    return envVariables.map((variable) => (variable.id === id ? updater(variable) : variable))
  }

  if (id === starterVariable.id) return [updater(starterVariable)]

  return envVariables
}

export const setEnvVariableKeyAtom = atom(
  null,
  (
    _get,
    set,
    {
      id,
      key,
      starterVariable,
    }: {
      id: string
      key: string
      starterVariable: EnvVariable
    },
  ) => {
    set(agentComposerEnvVariablesAtom, (envVariables) =>
      updateEnvVariable(envVariables, starterVariable, id, (variable) => ({ ...variable, key })),
    )
  },
)

export const setEnvVariableScopeAtom = atom(
  null,
  (
    _get,
    set,
    {
      id,
      scope,
      starterVariable,
    }: {
      id: string
      scope: EnvScope
      starterVariable: EnvVariable
    },
  ) => {
    set(agentComposerEnvVariablesAtom, (envVariables) =>
      updateEnvVariable(envVariables, starterVariable, id, (variable) => ({ ...variable, scope })),
    )
  },
)

export const setEnvVariableValueAtom = atom(
  null,
  (
    _get,
    set,
    {
      id,
      starterVariable,
      value,
    }: {
      id: string
      starterVariable: EnvVariable
      value: string
    },
  ) => {
    set(agentComposerEnvVariablesAtom, (envVariables) =>
      updateEnvVariable(envVariables, starterVariable, id, (variable) => ({ ...variable, value })),
    )
  },
)

export const addEnvVariableAtom = atom(
  null,
  (
    _get,
    set,
    {
      starterVariable,
      variable,
    }: {
      starterVariable: EnvVariable
      variable: EnvVariable
    },
  ) => {
    set(agentComposerEnvVariablesAtom, (envVariables) => [
      ...(envVariables.length > 0 ? envVariables : [starterVariable]),
      variable,
    ])
  },
)

export const importEnvVariablesAtom = atom(null, (_get, set, variables: EnvVariable[]) => {
  set(agentComposerEnvVariablesAtom, (envVariables) => [...envVariables, ...variables])
})

export const removeEnvVariableAtom = atom(null, (_get, set, id: string) => {
  set(agentComposerEnvVariablesAtom, (envVariables) =>
    envVariables.filter((variable) => variable.id !== id),
  )
})