import { useCallback, useRef } from "react"; import { getWebSocketClient } from "@/lib/ws/connection"; import { MessageSendError } from "@/lib/chat/message-send-error"; import { generateUUID } from "@/lib/utils"; import { useAppStoreApi } from "@/components/state-provider"; import { useQueue } from "./domains/session/use-queue"; import type { ChatSubmitPayload, MessageAttachment, } from "@/components/task/chat/chat-input-container"; import type { ActiveDocument } from "@/lib/state/slices/ui/types"; import type { PlanComment } from "@/lib/state/slices/comments"; import type { ContextFile } from "@/lib/state/context-files-store"; import type { CustomPrompt, Message, TaskPlanCommentRef } from "@/lib/types/http"; import type { TaskMentionData } from "@/hooks/use-inline-mention"; import type { AppState } from "@/lib/state/store"; import type { EntityReference } from "@/lib/types/entity-reference"; import { planCommentAdmissionConflict, toTaskPlanCommentRefs } from "@/lib/plan-comment-refs"; import { collectPromptReferenceExpansions, formatPromptReferenceExpansions, sanitizePromptReferenceSystemText, } from "@/lib/prompts/expand-prompt-references"; import { deriveSessionInputMode, type SessionInputMode, } from "./domains/session/session-input-mode"; import { t } from "@/lib/i18n"; import { getTaskPlanComments } from "@/lib/api/domains/plan-comment-api"; import { listTaskSessions } from "@/lib/api/domains/session-api"; export function buildDocumentContext( activeDocument: ActiveDocument | null, planModeEnabled: boolean, ): string { if (!activeDocument) return ""; if (activeDocument.type === "plan") { if (!planModeEnabled) return ""; // i18n-exempt: agent-facing prompt sent verbatim to the model, never rendered. return `\n\n\nACTIVE DOCUMENT: The user is editing the task plan side-by-side with this chat.\nRead the current plan using the get_task_plan_kandev MCP tool to understand the context before responding.\nAny plan modifications should use the update_task_plan_kandev MCP tool.\n`; } // i18n-exempt: agent-facing prompt sent verbatim to the model, never rendered. return `\n\n\nACTIVE DOCUMENT: The user is editing "${activeDocument.name}" (${activeDocument.path}) side-by-side with this chat.\nRead this file to understand the context before responding.\n`; } function resolveStepTitle(stepId: string, state: AppState): string { const step = state.kanban.steps.find((s) => s.id === stepId); if (step) return step.title; for (const snap of Object.values(state.kanbanMulti.snapshots)) { const found = (snap.steps ?? []).find((s) => s.id === stepId); if (found) return found.title; } return t("common:step"); } // Strips characters that could break out of the block when // task strings are interpolated verbatim — newlines (close-tag injection) // and angle brackets. Task titles can come from Jira/Linear sync or other // users in a shared workspace, so the data is not trusted. function sanitizeForPrompt(value: string): string { return value.replace(/[\r\n<>]/g, " "); } export function buildTaskMentionsContext(tasks: TaskMentionData[], state: AppState): string { if (tasks.length === 0) return ""; const lines = tasks.map((t) => { const stepTitle = resolveStepTitle(t.workflowStepId, state); const title = sanitizeForPrompt(t.title); const taskId = sanitizeForPrompt(t.taskId); const workflowId = sanitizeForPrompt(t.workflowId); const step = sanitizeForPrompt(stepTitle); const stateSuffix = t.state ? `, state: ${sanitizeForPrompt(t.state)}` : ""; return `- ${title} (id: ${taskId}, workflow_id: ${workflowId}, step: ${step}${stateSuffix})`; }); return ( `\n\n\n` + `REFERENCED TASKS: The user mentioned the following tasks. Use these IDs with the kandev MCP tools ` + `(e.g. \`get_task_conversation_kandev\`, \`update_task_kandev\`, \`get_task_plan_kandev\`) when the user asks you to act on them.\n` + lines.join("\n") + `\n` ); } export function buildContextFilesContext( contextFiles: ContextFile[], prompts: CustomPrompt[], ): string { const files = contextFiles.filter( (f) => !f.path.startsWith("prompt:") && f.path !== "plan:context", ); const promptFiles = contextFiles.filter((f) => f.path.startsWith("prompt:")); let context = ""; if (files.length > 0) { const pathList = files .map((f) => `- ${f.isDirectory ? "directory" : "file"}: ${sanitizeForPrompt(f.path)}`) .join("\n"); context += `\n\n\nCONTEXT PATHS: The user has attached the following file and directory paths as context. Inspect these paths to understand what the user is referring to:\n${pathList}\n`; } if (promptFiles.length > 0) { const promptsById = new Map(prompts.map((p) => [p.id, p])); const selectedPrompts = promptFiles .map((f) => promptsById.get(f.path.replace("prompt:", ""))) .filter((prompt): prompt is CustomPrompt => Boolean(prompt)); const selectedPromptNames = new Set(selectedPrompts.map((prompt) => prompt.name)); const promptExpansions = new Map(); const resolved = selectedPrompts .map((prompt) => { for (const expansion of collectPromptReferenceExpansions( prompt.content, prompts, prompt.name, selectedPromptNames, )) { if (!promptExpansions.has(expansion.name)) { promptExpansions.set(expansion.name, expansion.content); } } return `### ${sanitizePromptReferenceSystemText(prompt.name)}\n${sanitizePromptReferenceSystemText(prompt.content)}`; }) .filter(Boolean); if (resolved.length > 0) { const expansions = Array.from(promptExpansions, ([name, content]) => ({ name, content })); const expansionContext = formatPromptReferenceExpansions(expansions); context += `\n\n\nCONTEXT PROMPTS: The user has included the following prompt instructions as context:\n${resolved.join("\n\n")}${expansionContext ? "\n\n" + expansionContext : ""}\n`; } } return context; } export interface UseMessageHandlerParams { resolvedSessionId: string | null; taskId: string | null; sessionModel: string | null; activeModel: string | null; planModeEnabled?: boolean; hasPendingClarification?: boolean; activeDocument?: ActiveDocument | null; planComments?: PlanComment[]; contextFiles?: ContextFile[]; prompts?: CustomPrompt[]; } type SendMessagePayload = { taskId: string; resolvedSessionId: string; clientMessageId?: string; finalMessage: string; modelToSend: string | undefined; planMode: boolean; hasReviewComments?: boolean; attachments?: MessageAttachment[]; contextFilesMeta?: Array<{ path: string; name: string; is_directory?: boolean }>; entityReferences?: EntityReference[]; planCommentRefs?: TaskPlanCommentRef[]; requirePrimarySession?: boolean; }; type MessageListResponse = { messages?: Message[] }; function isUncertainMessageTransportError(error: unknown): boolean { if (!(error instanceof Error)) return false; const message = error.message.toLowerCase(); return ( message.includes("websocket request timed out") || message === "websocket connection closed" ); } async function findMessageByID( client: ReturnType, taskId: string, sessionId: string, messageId: string, ): Promise { if (!client) return undefined; const sessionIds = [sessionId]; try { const response = await listTaskSessions(taskId); for (const session of response.sessions ?? []) { if (session.id && !sessionIds.includes(session.id)) sessionIds.push(session.id); } } catch { // The submitted session remains a useful reconciliation fallback. } try { for (const candidateSessionId of sessionIds) { const response = await client.request( "message.list", { session_id: candidateSessionId, limit: 100, sort: "desc" }, 5000, ); const found = response.messages?.find((message) => message.id === messageId); if (found) return found; } return undefined; } catch { return undefined; } } async function waitForConnected(client: NonNullable>) { const getStatus = client.getStatus?.bind(client); if (!getStatus || getStatus() === "connected") return true; const deadline = Date.now() + 3000; while (Date.now() < deadline) { await new Promise((resolve) => setTimeout(resolve, 100)); if (getStatus() === "connected") return true; } return false; } type MessageReconciliation = { client: NonNullable>; taskId: string; sessionId: string; messageId: string; request: () => Promise; originalError: unknown; }; async function reconcileUncertainMessage({ client, taskId, sessionId, messageId, request, originalError, }: MessageReconciliation) { const committed = await findMessageByID(client, taskId, sessionId, messageId); if (committed) return committed; if (!(await waitForConnected(client))) throw originalError; try { return await request(); } catch (retryError) { const retriedMessage = await findMessageByID(client, taskId, sessionId, messageId); if (retriedMessage) return retriedMessage; throw retryError; } } export async function sendMessageRequest( payload: SendMessagePayload, ): Promise { const client = getWebSocketClient(); if (!client) { throw new MessageSendError( "connection-unavailable", "Connection unavailable. Reconnect and try again.", ); } const { taskId, resolvedSessionId, clientMessageId, finalMessage, modelToSend, planMode, hasReviewComments, attachments, contextFilesMeta, entityReferences, planCommentRefs, requirePrimarySession, } = payload; const hasAttachments = attachments && attachments.length > 0; const stableMessageId = clientMessageId ?? generateUUID(); const requestPayload = { task_id: taskId, session_id: resolvedSessionId, client_message_id: stableMessageId, content: finalMessage, ...(modelToSend && { model: modelToSend }), ...(planMode && { plan_mode: true }), ...(hasReviewComments && { has_review_comments: true }), ...(hasAttachments && { attachments }), ...(contextFilesMeta && { context_files: contextFilesMeta }), ...(entityReferences && { entity_references: entityReferences }), ...(planCommentRefs?.length && { plan_comment_refs: planCommentRefs }), ...(requirePrimarySession && { require_primary_session: true }), }; const request = () => client.request( "message.add", requestPayload, hasAttachments ? 30000 : 10000, ); try { return await request(); } catch (error) { if (!isUncertainMessageTransportError(error)) throw error; return reconcileUncertainMessage({ client, taskId, sessionId: resolvedSessionId, messageId: stableMessageId, request, originalError: error, }); } } const TERMINAL_SESSION_STATES = new Set(["FAILED", "CANCELLED", "COMPLETED"]); function requireSessionInputMode(state: AppState, selectedSessionId: string): SessionInputMode { const selectedSession = state.taskSessions.items[selectedSessionId] ?? null; const queuedCount = state.queue.metaBySessionId[selectedSessionId]?.count ?? 0; const inputMode = deriveSessionInputMode(selectedSession, queuedCount); if (inputMode === "unavailable") { // A terminal session row (agent process has exited) gets the backend's // actionable copy; a missing row keeps the generic message since there is // nothing session-specific to say. const message = selectedSession && TERMINAL_SESSION_STATES.has(selectedSession.state) ? t("task:sessionEndedCreateNew") : t("task:sessionNotAvailableForInput"); throw new MessageSendError("session-unavailable", message); } return inputMode; } function buildQueueAttachments(attachments?: MessageAttachment[]) { return attachments?.map((att) => ({ type: att.type, ...(att.attachment_id ? { attachment_id: att.attachment_id } : { data: att.data ?? "" }), mime_type: att.mime_type, name: att.name, size_bytes: att.size_bytes, delivery_mode: att.delivery_mode, })); } function normalizePlanCommentSendError( error: unknown, taskId: string, storeApi: ReturnType, ): unknown { const conflict = planCommentAdmissionConflict(error); if (!conflict) return error; if (conflict.snapshot) storeApi.getState().setTaskPlanComments(taskId, conflict.snapshot); return conflict.code === "plan_comments_changed" ? new MessageSendError("plan-comments-changed", t("task:planCommentsChangedRetry")) : new MessageSendError("primary-session-changed", t("task:primarySessionChangedRetry")); } function buildContextFilesMetadata(contextFiles: ContextFile[]) { const realFiles = contextFiles.filter( (file) => !file.path.startsWith("prompt:") && file.path !== "plan:context", ); if (realFiles.length === 0) return undefined; return realFiles.map((file) => ({ path: file.path, name: file.name, ...(file.isDirectory !== undefined ? { is_directory: file.isDirectory } : {}), })); } async function deliverComposedMessage({ payload, taskId, resolvedSessionId, finalMessage, modelToSend, planModeEnabled, hasPendingClarification, planCommentRefs, contextFilesMeta, inputMode, queue, storeApi, clientAdmissionId, }: { payload: ChatSubmitPayload; taskId: string; resolvedSessionId: string; finalMessage: string; modelToSend: string | undefined; planModeEnabled: boolean; hasPendingClarification: boolean; planCommentRefs: TaskPlanCommentRef[]; contextFilesMeta: ReturnType; inputMode: SessionInputMode; queue: ReturnType["queue"]; storeApi: ReturnType; clientAdmissionId: string; }) { try { if (hasPendingClarification || inputMode === "queue") { const accepted = await queue({ taskId, content: finalMessage, model: modelToSend, planMode: planModeEnabled, attachments: buildQueueAttachments(payload.attachments), entityReferences: payload.entityReferences, clientQueueId: clientAdmissionId, ...(planCommentRefs.length > 0 ? { planCommentRefs } : {}), ...(contextFilesMeta ? { contextFilesMeta } : {}), }); if (!accepted) { return false; } await refreshAcceptedPlanComments(taskId, planCommentRefs, storeApi); return; } const created = await sendMessageRequest({ taskId, resolvedSessionId, clientMessageId: clientAdmissionId, finalMessage, modelToSend, planMode: planModeEnabled, hasReviewComments: !!payload.reviewComments?.length, attachments: payload.attachments, contextFilesMeta, entityReferences: payload.entityReferences, planCommentRefs, }); if (created?.id && created.session_id) storeApi.getState().addMessage(created); await refreshAcceptedPlanComments(taskId, planCommentRefs, storeApi); } catch (error) { throw normalizePlanCommentSendError(error, taskId, storeApi); } } async function refreshAcceptedPlanComments( taskId: string, refs: TaskPlanCommentRef[], storeApi: ReturnType, ) { if (refs.length === 0) return; try { const snapshot = await getTaskPlanComments(taskId); storeApi.getState().setTaskPlanComments(taskId, snapshot); } catch (error) { // i18n-exempt: accepted delivery remains successful; foreground recovery retries this refresh. console.error("Failed to refresh task plan comments after delivery:", error); } } type PendingMessageAdmission = { key: string; id: string }; function messageAdmissionKey(parts: { taskId: string; resolvedSessionId: string; finalMessage: string; modelToSend?: string; planModeEnabled: boolean; hasReviewComments: boolean; planCommentRefs: TaskPlanCommentRef[]; contextFilesMeta: ReturnType; attachments: ChatSubmitPayload["attachments"]; entityReferences: ChatSubmitPayload["entityReferences"]; }) { return JSON.stringify(parts); } async function recoverPendingMessageAdmission( admission: PendingMessageAdmission | null, taskId: string, sessionId: string, refs: TaskPlanCommentRef[], storeApi: ReturnType, ) { if (!admission) return false; const client = getWebSocketClient(); const committed = client ? await findMessageByID(client, taskId, sessionId, admission.id) : undefined; if (!committed) return false; storeApi.getState().addMessage(committed); await refreshAcceptedPlanComments(taskId, refs, storeApi); return true; } function buildFinalMessageForSubmit({ payload, contextFiles, activeDocument, planModeEnabled, prompts, state, }: { payload: ChatSubmitPayload; contextFiles: ContextFile[]; activeDocument: ActiveDocument | null; planModeEnabled: boolean; prompts: CustomPrompt[]; state: AppState; }) { const allContextFiles = [...contextFiles, ...(payload.inlineMentions || [])]; const documentContext = buildDocumentContext(activeDocument, planModeEnabled); const contextFilesContext = buildContextFilesContext(allContextFiles, prompts); const taskMentionsContext = payload.inlineTaskMentions?.length ? buildTaskMentionsContext(payload.inlineTaskMentions, state) : ""; return { finalMessage: payload.message.trim() + documentContext + contextFilesContext + taskMentionsContext, allContextFiles, }; } export function useMessageHandler({ resolvedSessionId, taskId, sessionModel, activeModel, planModeEnabled = false, hasPendingClarification = false, activeDocument = null, planComments = [], contextFiles = [], prompts = [], }: UseMessageHandlerParams) { const { queue } = useQueue(resolvedSessionId); const storeApi = useAppStoreApi(); const pendingAdmissionRef = useRef(null); const handleSendMessage = useCallback( async (payload: ChatSubmitPayload) => { if (!taskId || !resolvedSessionId) { const error = new MessageSendError( "no-active-session", "No active task session. Start an agent before sending a message.", ); console.error(error.message); throw error; } const { finalMessage, allContextFiles } = buildFinalMessageForSubmit({ payload, contextFiles, activeDocument, planModeEnabled, prompts, state: storeApi.getState(), }); const modelToSend = activeModel && activeModel !== sessionModel ? activeModel : undefined; const planCommentRefs = payload.planCommentRefs ?? toTaskPlanCommentRefs(planComments); const contextFilesMeta = buildContextFilesMetadata(allContextFiles); const inputMode = requireSessionInputMode(storeApi.getState(), resolvedSessionId); const admissionKey = messageAdmissionKey({ taskId, resolvedSessionId, finalMessage, modelToSend, planModeEnabled, hasReviewComments: !!payload.reviewComments?.length, planCommentRefs, contextFilesMeta, attachments: payload.attachments, entityReferences: payload.entityReferences, }); const previousAdmission = pendingAdmissionRef.current?.key === admissionKey ? pendingAdmissionRef.current : null; const admission = previousAdmission ?? { key: admissionKey, id: generateUUID() }; pendingAdmissionRef.current = admission; if ( await recoverPendingMessageAdmission( previousAdmission, taskId, resolvedSessionId, planCommentRefs, storeApi, ) ) { if (pendingAdmissionRef.current === admission) pendingAdmissionRef.current = null; return; } const delivered = await deliverComposedMessage({ payload, taskId, resolvedSessionId, finalMessage, modelToSend, planModeEnabled, hasPendingClarification, planCommentRefs, contextFilesMeta, inputMode, queue, storeApi, clientAdmissionId: admission.id, }); if (delivered === false) return false; if (pendingAdmissionRef.current === admission) pendingAdmissionRef.current = null; }, [ resolvedSessionId, taskId, activeModel, sessionModel, planModeEnabled, hasPendingClarification, queue, storeApi, planComments, contextFiles, activeDocument, prompts, ], ); return { handleSendMessage }; }