windmill/frontend/src/lib/components/flows/linkedAgentDrafts.ts

import {
	ResourceService,
	type FlowModule,
	type FlowValue,
	type InputTransform,
	type Resource
} from '$lib/gen'
import { UserDraft } from '$lib/userDraft.svelte'
import { UserDraftDbSyncer } from '$lib/userDraftDbSyncer.svelte'
import { canWrite } from '$lib/utils'
import type { UserExt } from '$lib/stores'
import { dfs } from './dfs'
import { flowLocalInputs, type AIAgentConfig } from './agentResourceUtils'
import { AGENT_HISTORY_KEYS } from './agentFormFields'
import type { AgentResourceState } from './agentDraft.svelte'
import type { AgentTool } from './agentToolUtils'

/** A step names its agent bare or as `$res:<path>`/`res://<path>`; all three are the same agent,
 *  and a draft index has to answer for a lookup written any of those ways. Same normalization as
 *  `linkedAgentToolsStore`, and as the `trim_start_matches` the worker applies. */
export function normalizeAgentRef(agentRef: string): string {
	return agentRef.replace(/^\$res:/, '').replace(/^res:\/\//, '')
}

/** Every `ai_agent` resource this flow links to, deduped. `dfs` walks agent tool nodes as well as
 *  branches and loops, so a nested linked agent tool is included. */
export function linkedAgentPaths(value: FlowValue | undefined): string[] {
	if (!value?.modules) return []
	const paths = new Set<string>()
	for (const module of dfs(value.modules, (m) => m)) {
		const v = module?.value as { type?: string; agent?: string } | undefined
		if (v?.type === 'aiagent' && v.agent) {
			paths.add(normalizeAgentRef(v.agent))
		}
	}
	return [...paths]
}

/** Point every step of this flow linked to `from` at `to`, for an agent renamed from inside it.
 *  Returns the ids of the steps it moved. */
export function repointLinkedAgent(
	value: FlowValue | undefined,
	from: string,
	to: string
): string[] {
	if (!value?.modules) return []
	const moved: string[] = []
	for (const module of dfs(value.modules, (m) => m)) {
		const v = module?.value as { type?: string; agent?: string } | undefined
		if (v?.type === 'aiagent' && v.agent === from) {
			v.agent = to
			moved.push(module.id)
		}
	}
	return moved
}

/**
 * The unsaved draft for an agent, freshest first: the cell an open agent editor is writing, then
 * what a `get_draft` response carried.
 *
 * Only the live cell is reliably current, and only while an editor holds it: `releaseEntry` drops
 * the cached write at refcount 0 on purpose, so once the agent editor closes the persisted row is
 * the sole answer. Read through `fetchAgentWithDraft`, which settles that row first.
 */
export function agentDraftState(
	response: { draft?: unknown },
	path: string,
	workspace: string | undefined
): AgentResourceState | undefined {
	const live = UserDraft.get<AgentResourceState>('resource', path, { workspace })
	return live ?? (response.draft as AgentResourceState | undefined)
}

/** A refusal that already names the agent and says what is wrong with it, so a caller wrapping it
 *  would only repeat itself. */
export class AgentDraftUnavailable extends Error {}

/**
 * An agent's resource together with the draft a run of it would use.
 *
 * The flush is what makes the answer current. Autosave is debounced by 1.5s (10s ceiling), and
 * closing the agent editor releases the in-memory cell without cancelling that pending POST — so
 * testing or deploying right after closing would otherwise read a row the last edits have not
 * reached yet. `flush` replays the parked save and is a no-op when there is none.
 */
export async function fetchAgentWithDraft(
	path: string,
	workspace: string
): Promise<{ response: Resource; draft: AgentResourceState | undefined }> {
	const query = { workspace, itemKind: 'resource' as const, path }
	await UserDraftDbSyncer.flush(query)
	// `flush` resolves whether or not the save actually landed: `postSave` catches network and
	// server errors into its failure map, and answers a conflicting write by parking a snapshot,
	// returning normally in both cases. The row about to be read is then older than the edit still
	// held in the browser, and nothing downstream could tell. Running that row is a test of the
	// wrong agent; deploying it is worse, because the deploy deletes the draft and takes the newer
	// edit with it. Neither is recoverable from here, so refuse the read.
	const failure = UserDraftDbSyncer.getState(query).failureMessage
	if (failure) {
		throw new AgentDraftUnavailable(`The unsaved changes to ${path} could not be saved: ${failure}`)
	}
	if (UserDraftDbSyncer.getConflict(query).conflict) {
		throw new AgentDraftUnavailable(
			`The unsaved changes to ${path} could not be saved because it was edited elsewhere. Open the agent to resolve it.`
		)
	}
	const response = await ResourceService.getResource({ workspace, path, getDraft: true })
	return { response, draft: agentDraftState(response, path, workspace) }
}

/** One linked agent whose resource the user has an unsaved draft for. */
export interface LinkedAgentDraft {
	path: string
	/** The draft's resource value: what a run of this agent would use. */
	args: AIAgentConfig
	/** The whole draft row, as the resource editors write it — the deploy payload. */
	state: AgentResourceState
	/** No deployed row at this path, so deploying has to create rather than update. */
	noDeployed: boolean
	/** Of the deployed resource, for `agentDraftCanWrite`. */
	extraPerms: Record<string, boolean>
}

/** Whether `user` may write this agent's resource. Split from the load so that resolving the
 *  drafts of a whole flow costs no `whoami` — only the deploy dialog needs the answer, and it
 *  looks the user up once for every agent it lists. */
export function agentDraftCanWrite(draft: LinkedAgentDraft, user: UserExt | undefined): boolean {
	return canWrite(draft.path, draft.extraPerms, user)
}

/** A link that cannot resolve for the user rather than because something went wrong: the agent was
 *  deleted, or sits in a folder they cannot read. Both are ordinary states of a rigid link, and
 *  neither should stop the caller — the flow still tests and deploys, against the deployed agent.
 *  Every other failure is an outage, and answering "no draft" to one would quietly run or deploy
 *  the wrong configuration, which is the whole thing this module exists to prevent. */
export function isExpectedLinkFailure(err: unknown): boolean {
	const status = (err as { status?: number } | null | undefined)?.status
	return status === 401 || status === 403 || status === 404
}

/**
 * The unsaved draft of every given `ai_agent` path, for the paths that have one.
 *
 * Throws when a path fails to load for any reason other than being missing or unreadable, so a
 * caller cannot mistake an outage for an agent with nothing unsaved.
 */
export async function loadLinkedAgentDrafts(
	paths: string[],
	workspace: string | undefined
): Promise<Map<string, LinkedAgentDraft>> {
	const out = new Map<string, LinkedAgentDraft>()
	if (!workspace || paths.length === 0) return out
	await Promise.all(
		paths.map(async (path) => {
			let response: Resource
			let draft: AgentResourceState | undefined
			try {
				;({ response, draft } = await fetchAgentWithDraft(path, workspace))
			} catch (err) {
				if (isExpectedLinkFailure(err)) return
				if (err instanceof AgentDraftUnavailable) throw err
				throw new Error(`Could not load the agent ${path}: ${err}`)
			}
			if (!draft) return
			out.set(path, {
				path,
				args: (draft.args ?? {}) as AIAgentConfig,
				state: draft,
				noDeployed: Boolean((response as { no_deployed?: boolean }).no_deployed),
				extraPerms: response.extra_perms ?? {}
			})
		})
	)
	return out
}

/**
 * Every argument a saved agent carries, as a static input transform. Not only the keys the agent
 * form renders: a run reads them all, and an agent holding its own `user_message` answers with it
 * when nothing overrides it. `tools` is the step's own roster rather than an input, so it rides on
 * the module's `tools` key instead.
 */
export function agentArgsToTransforms(args: AIAgentConfig): Record<string, InputTransform> {
	const it: Record<string, InputTransform> = {}
	for (const [key, value] of Object.entries(args ?? {})) {
		if (key === 'tools' || value === undefined) continue
		it[key] = { type: 'static', value } as InputTransform
	}
	return it
}

type AiAgentValue = Extract<FlowModule['value'], { type: 'aiagent' }>

/**
 * The standalone step a linked step's draft would run as: the draft's brain and tools inlined, with
 * the step's own flow-local inputs kept on top.
 *
 * The overlay order is the worker's (`ai_executor.rs`): its linked branch interpolates the whole
 * resource brain and only then writes the flow-local inputs (`user_message`, `user_attachments`,
 * `enabled_tools`, `memory_id`, `previous_messages`) back from the step's own args. `tool_inputs`
 * stays untouched — the worker overlays it onto the tools in both branches, so an inlined step
 * keeps the host flow's tool bindings. The worker never reads a history input from the resource,
 * so one a draft happens to carry is left out here too.
 */
export function inlineAgentDraft(value: AiAgentValue, args: AIAgentConfig): AiAgentValue {
	const { agent: _agent, ...rest } = value
	const brain = agentArgsToTransforms(args)
	for (const key of AGENT_HISTORY_KEYS) delete brain[key]
	return {
		...rest,
		tools: (args.tools ?? []) as AgentTool[],
		input_transforms: {
			...brain,
			...flowLocalInputs(value.input_transforms as Record<string, InputTransform>)
		}
	} as AiAgentValue
}

/**
 * Replace every linked agent step that has a draft with the draft's own configuration, so a preview
 * runs what the agent editor is showing rather than the deployed resource. Returns a new value: the
 * flow editor hands its live store object to previews.
 */
export function inlineAgentDrafts(
	value: FlowValue,
	drafts: Map<string, LinkedAgentDraft>
): FlowValue {
	if (drafts.size === 0) return value
	// JSON rather than `structuredClone`: the flow editor's value is a Svelte `$state` proxy, which
	// `structuredClone` refuses outright. A flow value is JSON by definition — it is about to be
	// posted as one — so the round trip loses nothing this preview would have carried.
	const next = JSON.parse(JSON.stringify(value)) as FlowValue
	for (const module of dfs(next.modules ?? [], (m) => m)) {
		const v = module?.value as AiAgentValue | undefined
		if (v?.type !== 'aiagent' || !v.agent) continue
		const draft = drafts.get(normalizeAgentRef(v.agent))
		if (!draft) continue
		module.value = inlineAgentDraft(v, draft.args)
	}
	return next
}

/** Load the drafts this flow's linked agents have and inline them. The whole substitution, for a
 *  caller holding nothing but the value it is about to preview. */
export async function withAgentDrafts(
	value: FlowValue,
	workspace: string | undefined
): Promise<FlowValue> {
	const paths = linkedAgentPaths(value)
	if (paths.length === 0) return value
	return inlineAgentDrafts(value, await loadLinkedAgentDrafts(paths, workspace))
}