import { useCallback, useMemo, useReducer } from "react"; import { useTranslation } from "react-i18next"; import type { ComposerAttachment } from "@/attachments/types"; import { resolveComposerAttachmentSubmitFormat, splitComposerAttachmentsForSubmit, } from "@/composer/attachments/submit"; import { useCreateFlowStore } from "@/stores/create-flow-store"; import { useSessionStore } from "@/stores/session-store"; import { buildOptimisticUserMessage, generateMessageId, type StreamItem, type UserMessageImageAttachment, } from "@/types/stream"; import type { AgentAttachment } from "@getpaseo/protocol/messages"; const EMPTY_STREAM_ITEMS: StreamItem[] = []; interface CreateAttempt { clientMessageId: string; text: string; timestamp: Date; images?: UserMessageImageAttachment[]; attachments?: AgentAttachment[]; } type DraftAgentMachineState = | { tag: "draft"; errorMessage: string } | { tag: "creating"; attempt: CreateAttempt }; type DraftAgentMachineEvent = | { type: "DRAFT_SET_ERROR"; message: string } | { type: "SUBMIT"; attempt: CreateAttempt } | { type: "CREATE_FAILED"; message: string }; function assertNever(value: never): never { throw new Error(`Unhandled state: ${JSON.stringify(value)}`); } function reducer( state: DraftAgentMachineState, event: DraftAgentMachineEvent, ): DraftAgentMachineState { switch (event.type) { case "DRAFT_SET_ERROR": { if (state.tag !== "draft") { return state; } return { ...state, errorMessage: event.message }; } case "SUBMIT": { return { tag: "creating", attempt: event.attempt }; } case "CREATE_FAILED": { if (state.tag !== "creating") { return state; } return { tag: "draft", errorMessage: event.message }; } default: return assertNever(event); } } interface CreateRequestResult { agentId: string | null; result: TCreateResult; } interface SubmitContext { text: string; attachments: ComposerAttachment[]; cwd: string; } interface CreateRequestContext { attempt: CreateAttempt; text: string; images?: UserMessageImageAttachment[]; attachments?: AgentAttachment[]; cwd: string; } interface UseDraftAgentCreateFlowOptions { draftId: string; getPendingServerId: () => string | null; initialAttempt?: CreateAttempt | null; allowEmptyText?: boolean; validateBeforeSubmit?: (ctx: SubmitContext) => string | null; onBeforeSubmit?: (ctx: CreateRequestContext) => Promise | void; onCreateStart?: () => void; createRequest: (ctx: CreateRequestContext) => Promise>; buildDraftAgent: (attempt: CreateAttempt) => TDraftAgent; onCreateSuccess: (ctx: { result: TCreateResult; attempt: CreateAttempt }) => Promise | void; onCreateError?: (error: Error) => void; } export function useDraftAgentCreateFlow({ draftId, getPendingServerId, initialAttempt = null, allowEmptyText = false, validateBeforeSubmit, onBeforeSubmit, onCreateStart, createRequest, buildDraftAgent, onCreateSuccess, onCreateError, }: UseDraftAgentCreateFlowOptions) { const { t } = useTranslation(); const [machine, dispatch] = useReducer( reducer, initialAttempt, (attempt): DraftAgentMachineState => attempt ? { tag: "creating", attempt } : { tag: "draft", errorMessage: "", }, ); const setPendingCreateAttempt = useCreateFlowStore((state) => state.setPending); const updatePendingAgentId = useCreateFlowStore((state) => state.updateAgentId); const markPendingCreateLifecycle = useCreateFlowStore((state) => state.markLifecycle); const clearPendingCreateAttempt = useCreateFlowStore((state) => state.clear); const handoffCreatedAgentUserMessage = useSessionStore( (state) => state.handoffCreatedAgentUserMessage, ); const formErrorMessage = machine.tag === "draft" ? machine.errorMessage : ""; const isSubmitting = machine.tag === "creating"; const optimisticStreamItems = useMemo(() => { if (machine.tag !== "creating") { return EMPTY_STREAM_ITEMS; } if ( !machine.attempt.text && (!machine.attempt.images || machine.attempt.images.length === 0) && (!machine.attempt.attachments || machine.attempt.attachments.length === 0) ) { return EMPTY_STREAM_ITEMS; } return [ buildOptimisticUserMessage({ id: machine.attempt.clientMessageId, text: machine.attempt.text, timestamp: machine.attempt.timestamp, images: machine.attempt.images, attachments: machine.attempt.attachments, }), ]; }, [machine]); const draftAgent = useMemo(() => { if (machine.tag !== "creating") { return null; } return buildDraftAgent(machine.attempt); }, [buildDraftAgent, machine]); const runCreateAttempt = useCallback( async ({ attempt, cwd }: { attempt: CreateAttempt; cwd: string }) => { const pendingServerId = getPendingServerId(); if (!pendingServerId) { const error = new Error(t("composer.errors.noHostSelected")); dispatch({ type: "DRAFT_SET_ERROR", message: error.message }); throw error; } await onBeforeSubmit?.({ attempt, text: attempt.text, images: attempt.images, attachments: attempt.attachments, cwd, }); try { const createResult = await createRequest({ attempt, text: attempt.text, images: attempt.images, attachments: attempt.attachments, cwd, }); if (createResult.agentId) { updatePendingAgentId({ draftId, agentId: createResult.agentId }); handoffCreatedAgentUserMessage( pendingServerId, createResult.agentId, buildOptimisticUserMessage({ id: attempt.clientMessageId, text: attempt.text, timestamp: attempt.timestamp, images: attempt.images, attachments: attempt.attachments, }), ); markPendingCreateLifecycle({ draftId, lifecycle: "sent" }); } await onCreateSuccess({ result: createResult.result, attempt }); } catch (error) { const resolved = error instanceof Error ? error : new Error(t("composer.errors.failedToCreateAgent")); dispatch({ type: "CREATE_FAILED", message: resolved.message }); markPendingCreateLifecycle({ draftId, lifecycle: "abandoned" }); clearPendingCreateAttempt({ draftId }); onCreateError?.(resolved); throw error; } }, [ handoffCreatedAgentUserMessage, clearPendingCreateAttempt, createRequest, draftId, getPendingServerId, markPendingCreateLifecycle, onBeforeSubmit, onCreateError, onCreateSuccess, t, updatePendingAgentId, ], ); const handleCreateFromInput = useCallback( async ({ text, attachments, cwd }: SubmitContext) => { if (isSubmitting) { throw new Error(t("composer.errors.alreadyLoading")); } dispatch({ type: "DRAFT_SET_ERROR", message: "" }); const trimmedPrompt = text.trim(); const pendingServerId = getPendingServerId(); if (!pendingServerId) { const error = new Error(t("composer.errors.noHostSelected")); dispatch({ type: "DRAFT_SET_ERROR", message: error.message }); throw error; } const supportsForgeSearch = useSessionStore.getState().sessions[pendingServerId]?.serverInfo?.features?.forgeSearch === true; const wirePayload = splitComposerAttachmentsForSubmit(attachments, { format: resolveComposerAttachmentSubmitFormat({ supportsForgeAttachments: supportsForgeSearch, }), }); const images = wirePayload.images; const hasAttachmentContent = images.length > 0 || wirePayload.attachments.length > 0; if (!trimmedPrompt && !hasAttachmentContent && !allowEmptyText) { const error = new Error(t("composer.errors.initialPromptRequired")); dispatch({ type: "DRAFT_SET_ERROR", message: error.message }); throw error; } const validationError = validateBeforeSubmit?.({ text: trimmedPrompt, attachments, cwd, }); if (validationError) { const error = new Error(validationError); dispatch({ type: "DRAFT_SET_ERROR", message: validationError }); throw error; } const attempt: CreateAttempt = { clientMessageId: generateMessageId(), text: trimmedPrompt, timestamp: new Date(), ...(images && images.length > 0 ? { images } : {}), ...(wirePayload.attachments.length > 0 ? { attachments: wirePayload.attachments } : {}), }; setPendingCreateAttempt({ draftId, serverId: pendingServerId, agentId: null, clientMessageId: attempt.clientMessageId, text: attempt.text, timestamp: attempt.timestamp.getTime(), ...(attempt.images && attempt.images.length > 0 ? { images: attempt.images } : {}), ...(attempt.attachments && attempt.attachments.length > 0 ? { attachments: attempt.attachments } : {}), }); dispatch({ type: "SUBMIT", attempt }); onCreateStart?.(); await runCreateAttempt({ attempt, cwd }); }, [ allowEmptyText, draftId, getPendingServerId, isSubmitting, onCreateStart, runCreateAttempt, setPendingCreateAttempt, t, validateBeforeSubmit, ], ); const continueCreateFromAttempt = useCallback( async ({ attempt, cwd }: { attempt: CreateAttempt; cwd: string }) => { if (!isSubmitting) { dispatch({ type: "SUBMIT", attempt }); } await runCreateAttempt({ attempt, cwd }); }, [isSubmitting, runCreateAttempt], ); return { machine, formErrorMessage, isSubmitting, optimisticStreamItems, draftAgent, handleCreateFromInput, continueCreateFromAttempt, }; } export type { CreateAttempt as DraftCreateAttempt };