import type { Dispatch, MutableRefObject, SetStateAction } from 'react' import { appendLocalResponse, appendLocalMarker, appendStreamingMessage, buildFormattedMessage, createMissingAgentResponse, type AgentStatus, type AgentExecutionContext, type AiAgentMessage, type PendingUserPrompt, } from './aiAgentConversation' import type { AgentFileCallbacks } from './aiAgentFileOperations' import { createStreamCallbacks } from './aiAgentStreamCallbacks' import type { ToolInvocation } from './aiAgentMessageState' import { trackAiAgentMessageBlocked, trackAiAgentMessageSent, trackAiAgentResponseStopped } from './productAnalytics' import { streamAiAgent } from '../utils/streamAiAgent' import { streamAiModel } from '../utils/streamAiModel' import { hydrateNoteReferences } from '../utils/ai-reference-content' import { createTranslator } from './i18n' export interface AiAgentAbortState { aborted: boolean controller?: AbortController } export interface AiAgentSessionRuntime { setMessages: Dispatch> setStatus: Dispatch> abortRef: MutableRefObject responseAccRef: MutableRefObject fileCallbacksRef: MutableRefObject toolInputMapRef: MutableRefObject> messagesRef: MutableRefObject statusRef: MutableRefObject } interface SendAgentMessageOptions { runtime: AiAgentSessionRuntime context: AgentExecutionContext prompt: PendingUserPrompt } interface RegenerateAgentMessageOptions { runtime: AiAgentSessionRuntime context: AgentExecutionContext messageId: string } interface SelectedTargetStreamRequest { context: AgentExecutionContext formattedMessage: string systemPrompt: string callbacks: ReturnType signal?: AbortSignal } function normalizePrompt(prompt: PendingUserPrompt): PendingUserPrompt { return { text: prompt.text.trim(), references: prompt.references && prompt.references.length > 0 ? prompt.references : undefined, } } function completedMessageCount(messages: AiAgentMessage[]): number { return messages.filter((message) => !message.isStreaming && !message.localMarker).length } function shouldIgnorePrompt(status: AgentStatus, prompt: PendingUserPrompt): boolean { return !prompt.text || status === 'thinking' || status === 'tool-executing' } function blockMissingVault(runtime: AiAgentSessionRuntime, context: AgentExecutionContext, prompt: PendingUserPrompt): void { trackAiAgentMessageBlocked(context.agent, 'missing_vault') appendLocalResponse(runtime.setMessages, prompt, 'No vault loaded. Open a vault first.') } function blockUnavailableAgent(runtime: AiAgentSessionRuntime, context: AgentExecutionContext, prompt: PendingUserPrompt): void { trackAiAgentMessageBlocked(context.agent, 'agent_unavailable') appendLocalResponse( runtime.setMessages, prompt, createMissingAgentResponse(context.agent), ) } async function streamWithSelectedTarget({ context, formattedMessage, systemPrompt, callbacks, signal, }: SelectedTargetStreamRequest): Promise { if (context.target?.kind === 'api_model') { await streamAiModel({ provider: context.target.provider, model: context.target.model, message: formattedMessage, systemPrompt, vaultPath: context.vaultPath, vaultPaths: context.vaultPaths, callbacks, }) return } await streamAiAgent({ agent: context.agent, message: formattedMessage, systemPrompt, vaultPath: context.vaultPath, vaultPaths: context.vaultPaths, permissionMode: context.permissionMode, callbacks, signal, }) } function stoppedResponseText(response: string, locale: AgentExecutionContext['locale']): string { const stopped = createTranslator(locale ?? 'en')('ai.panel.stoppedResponse') const partial = response.trim() return partial ? `${partial}\n\n${stopped}` : stopped } export async function sendAgentMessage({ runtime, context, prompt, }: SendAgentMessageOptions): Promise { const currentStatus = runtime.statusRef.current const normalizedPrompt = normalizePrompt(prompt) if (shouldIgnorePrompt(currentStatus, normalizedPrompt)) return if (!context.vaultPath) { blockMissingVault(runtime, context, normalizedPrompt) return } if (!context.ready) { blockUnavailableAgent(runtime, context, normalizedPrompt) return } trackAiAgentMessageSent({ agent: context.agent, permissionMode: context.permissionMode, hasContext: !!context.systemPromptOverride, referenceCount: normalizedPrompt.references?.length ?? 0, historyMessageCount: completedMessageCount(runtime.messagesRef.current), }) const controller = new AbortController() const abortState: AiAgentAbortState = { aborted: false, controller, } runtime.abortRef.current = abortState runtime.responseAccRef.current = '' runtime.toolInputMapRef.current = new Map() const messageId = appendStreamingMessage(runtime.setMessages, normalizedPrompt) runtime.setStatus('thinking') const promptForAgent = { ...normalizedPrompt, references: await hydrateNoteReferences(normalizedPrompt.references), } const { formattedMessage, systemPrompt } = buildFormattedMessage( context, runtime.messagesRef.current, promptForAgent, ) const callbacks = createStreamCallbacks({ agent: context.agent, responseSourceLabel: context.target?.label, locale: context.locale, messageId, vaultPath: context.vaultPath, setMessages: runtime.setMessages, setStatus: runtime.setStatus, abortRef: { current: abortState }, responseAccRef: runtime.responseAccRef, toolInputMapRef: runtime.toolInputMapRef, fileCallbacksRef: runtime.fileCallbacksRef, }) await streamWithSelectedTarget({ context, formattedMessage, systemPrompt, callbacks, signal: controller.signal, }) } export async function regenerateAgentMessage({ runtime, context, messageId, }: RegenerateAgentMessageOptions): Promise { const currentMessages = runtime.messagesRef.current const messageIndex = currentMessages.findIndex((message) => message.id === messageId) const message = currentMessages[messageIndex] if (!message || message.localMarker || runtime.statusRef.current === 'thinking' || runtime.statusRef.current === 'tool-executing') return const preservedMessages = currentMessages.slice(0, messageIndex) runtime.abortRef.current = { aborted: false } runtime.responseAccRef.current = '' runtime.toolInputMapRef.current = new Map() runtime.messagesRef.current = preservedMessages runtime.statusRef.current = 'idle' runtime.setMessages(preservedMessages) runtime.setStatus('idle') await sendAgentMessage({ runtime, context, prompt: { text: message.userMessage, references: message.references, }, }) } export function addAgentLocalMarker( runtime: Pick, text: string, ): void { appendLocalMarker(runtime.setMessages, text) } export function clearAgentConversation(runtime: Pick): void { runtime.abortRef.current.aborted = true runtime.abortRef.current.controller?.abort() runtime.responseAccRef.current = '' runtime.toolInputMapRef.current = new Map() runtime.setMessages([]) runtime.setStatus('idle') } export function stopAgentMessage( runtime: AiAgentSessionRuntime, context: Pick, ): void { if (!runtime.abortRef.current.controller || runtime.abortRef.current.aborted) return runtime.abortRef.current.aborted = true runtime.abortRef.current.controller.abort() const response = runtime.responseAccRef.current const toolCount = runtime.toolInputMapRef.current.size trackAiAgentResponseStopped(context.agent, response, toolCount) runtime.setMessages((current) => current.map((message) => ( message.isStreaming ? { ...message, isStreaming: false, reasoningDone: true, response: stoppedResponseText(response, context.locale), actions: message.actions.map((action) => ( action.status === 'pending' ? { ...action, status: 'error' as const } : action )), } : message ))) runtime.responseAccRef.current = '' runtime.toolInputMapRef.current = new Map() runtime.setStatus('idle') }