Blame for backend/src/lib/documentMode.ts

import type { ChatMessage, DocIndex } from "./chatTools";
import type { RunTree } from "langsmith";
import type { UserApiKeys } from "./llm/types";
import {
    buildSaudiLawRoutingResult,
    classifySaudiLawNeed,
    type SaudiLawRoutingResult,
} from "./saudiLawRouting";
import { classifyWebSearchNeed, detectWebSearchIntentHeuristic } from "./webSearchRouting";

/**
 * Explicit assistant-mode contract. Every chat request resolves to exactly
 * one of these. `auto` means "the server may infer a mode from heuristics" —
 * every other value is an explicit user choice that is already authorized by
 * the user's selection and must never be silently upgraded to a mutating mode.
 */
export type RequestedAssistantMode =
    | "auto"
    | "answer"
    | "comments"
    | "document_edit"
    | "workflow_review"
    | "document_generation"
    | "web_search";

const REQUESTED_ASSISTANT_MODES: readonly RequestedAssistantMode[] = [
    "auto",
    "answer",
    "comments",
    "document_edit",
    "workflow_review",
    "document_generation",
    "web_search",
];

/**
 * Validate an untrusted `requested_mode` request field. Throws a tagged error
 * whose message contains `invalid_requested_mode` so routes can translate it to
 * an HTTP 400 before any model/provider/tool/storage/database side effect.
 */
export function parseRequestedAssistantMode(
    value: unknown,
): RequestedAssistantMode {
    if (value === undefined || value === null || value === "") return "auto";
    if (
        typeof value === "string" &&
        (REQUESTED_ASSISTANT_MODES as readonly string[]).includes(value)
    ) {
        return value as RequestedAssistantMode;
    }
    throw new Error(`invalid_requested_mode: ${String(value)}`);
}

export type AssistantTurnMode =
    | "answer"
    | "comments"
    | "document_direct_edit"
    | "document_edit_with_web_search"
    | "web_search"
    | "workflow_review"
    | "document_generation"
    | "project_chat";

export type DocumentModeDecision = {
    trackChanges: boolean;
    reason:
        | "explicit_request_true"
        | "explicit_comments_only"
        | "auto_docx_edit_intent"
        | "no_editable_docx"
        | "default_comments_only"
        | "explicit_non_mutating_mode";
    editIntent: boolean;
    commentsOnlyIntent: boolean;
    editableDocxCount: number;
    pdfCount: number;
};

export type TurnModeSource = "explicit" | "inferred";

export type TurnModeReason =
    | "explicit_mode"
    | "explicit_workflow_execution"
    | "user_requested_workflow"
    | "docx_mutation_with_web_search"
    | "direct_docx_mutation"
    | "web_search_intent"
    | "web_search_routing"
    | "document_generation_intent"
    | "normal_answer"
    | "saudi_law_requires_web_search"
    | "web_search_blocked_docs_attached";

export type TurnModeDecision = {
    turnMode: AssistantTurnMode;
    requestedMode: RequestedAssistantMode;
    source: TurnModeSource;
    documentMode: DocumentModeDecision;
    enableGoogleSearch: boolean;
    requiresConfirmation: boolean;
    reason: TurnModeReason;
    saudiLawRouting?: SaudiLawRoutingResult;
    routingContext?: {
        followUpDetected: boolean;
        inheritedWebSearch: boolean;
        routingText: string;
        inheritanceSource: "deterministic" | "none";
    };
};

const COMMENTS_ONLY_RE =
    /\b(comments? only|feedback only|review only|marginalia|do not edit|don't edit|no edits|no redlines|without changing|do not change|leave the text unchanged)\b/i;

const EDIT_INTENT_RE =
    /\b(rewrite|redline|red-line|edit|revise|amend|modify|change|replace|insert|add|delete|remove|fix|correct|tracked changes?|make .* favourable|make .* favorable|more favourable|more favorable)\b/i;

const COMMENT_EDIT_INTENT_RE =
    /\b(add|insert|place|leave|make|create)\b.{0,40}\b(comment|note|annotation|margin note)\b/i;

const UPDATE_DOCUMENT_INTENT_RE =
    /\bupdate\b.{0,80}\b(document|docx|file|agreement|contract|clause|section|paragraph|text|wording|draft)\b|\b(document|docx|file|agreement|contract|clause|section|paragraph|text|wording|draft)\b.{0,80}\bupdate\b/i;

const WORKFLOW_REVIEW_RE =
    /\b(run|execute|start|apply|use)\b.{0,80}\b(workflow|playbook|review workflow|buyer playbook|seller playbook|zamil playbook)\b/i;

const DOCUMENT_GENERATION_RE =
    /\b(draft|generate|create|produce|prepare|write)\b.{0,80}\b(document|docx|agreement|contract|letter|memo|notice|template|deed|resolution)\b/i;

const FOLLOW_UP_FRAGMENT_RE =
    /^\s*(for|in|at|near|around|and|also|same(?:\s+for)?|what about|how about|now|today|tomorrow|yesterday|next week|this week|source|sources|cite|citation|citations|verify|fact[- ]?check|in arabic|arabic|translate)\b/i;

const REFERENTIAL_FOLLOW_UP_RE =
    /\b(that|this|it|those|these|same|above|previous|earlier|there|here)\b/i;

const TOPIC_RESET_RE =
    /\b(new topic|different question|unrelated|forget that|instead)\b/i;

const PRIOR_TURN_CONTEXT_CHAR_LIMIT = 600;

export type PriorTurnRoutingState = {
    turnMode?: AssistantTurnMode | string | null;
    enableGoogleSearch?: boolean | null;
    reason?: string | null;
};

export function extractPriorTurnRoutingState(
    content: unknown,
): PriorTurnRoutingState | null {
    if (!Array.isArray(content)) return null;
    const modeEvent = content.find((event): event is Record<string, unknown> => {
        return (
            typeof event === "object" &&
            event !== null &&
            (event as Record<string, unknown>).type === "mode_decision"
        );
    });
    const decision = modeEvent?.decision;
    if (typeof decision !== "object" || decision === null) return null;
    const record = decision as Record<string, unknown>;
    return {
        turnMode:
            typeof record.turnMode === "string" ? record.turnMode : undefined,
        enableGoogleSearch:
            typeof record.enableGoogleSearch === "boolean"
                ? record.enableGoogleSearch
                : undefined,
        reason: typeof record.reason === "string" ? record.reason : undefined,
    };
}

function detectWebSearchIntent(
    lastUserText: string,
    hasAttachedDocs = false,
): boolean {
    return detectWebSearchIntentHeuristic(lastUserText, { hasAttachedDocs });
}

type DocIndexEntry = {
    filename?: string;
    file_type?: string | null;
};

function isDocx(info: DocIndexEntry): boolean {
    const type = info.file_type?.toLowerCase();
    const name = info.filename?.toLowerCase() ?? "";
    return type === "docx" || name.endsWith(".docx");
}

function isPdf(info: DocIndexEntry): boolean {
    const type = info.file_type?.toLowerCase();
    const name = info.filename?.toLowerCase() ?? "";
    return type === "pdf" || name.endsWith(".pdf");
}

function latestUserText(messages: ChatMessage[]): string {
    const lastUser = [...messages].reverse().find((m) => m.role === "user");
    if (!lastUser?.content) return "";
    return String(lastUser.content);
}

function previousUserText(messages: ChatMessage[]): string {
    let seenLatest = false;
    for (let i = messages.length - 1; i >= 0; i--) {
        const message = messages[i];
        if (message.role !== "user") continue;
        if (!seenLatest) {
            seenLatest = true;
            continue;
        }
        return String(message.content ?? "");
    }
    return "";
}

function priorStateNeedsWebSearch(
    prior: PriorTurnRoutingState | null | undefined,
): boolean {
    if (!prior) return false;
    if (prior.enableGoogleSearch === true) return true;
    return (
        prior.turnMode === "web_search" ||
        prior.turnMode === "document_edit_with_web_search"
    );
}

function isLikelyFollowUpFragment(text: string): boolean {
    const trimmed = text.trim();
    if (!trimmed) return false;
    if (TOPIC_RESET_RE.test(trimmed)) return false;
    if (WORKFLOW_REVIEW_RE.test(trimmed) || DOCUMENT_GENERATION_RE.test(trimmed)) {
        return false;
    }
    if (EDIT_INTENT_RE.test(trimmed) || COMMENT_EDIT_INTENT_RE.test(trimmed)) {
        return false;
    }
    if (FOLLOW_UP_FRAGMENT_RE.test(trimmed)) return true;
    if (REFERENTIAL_FOLLOW_UP_RE.test(trimmed) && trimmed.length <= 140) {
        return true;
    }
    if (
        trimmed.length <= 32 &&
        !/\b(what|who|why|when|where|how|draft|generate|rewrite|review|explain|tell|define)\b/i.test(
            trimmed,
        )
    ) {
        return true;
    }
    return false;
}

function buildFollowUpRoutingContext(params: {
    messages: ChatMessage[];
    priorTurn?: PriorTurnRoutingState | null;
}): {
    followUpDetected: boolean;
    inheritedWebSearch: boolean;
    routingText: string;
    inheritanceSource: "deterministic" | "none";
} {
    const latest = latestUserText(params.messages).trim();
    const previous = previousUserText(params.messages).trim();
    if (!latest || !previous || !isLikelyFollowUpFragment(latest)) {
        return {
            followUpDetected: false,
            inheritedWebSearch: false,
            routingText: latest,
            inheritanceSource: "none",
        };
    }

    const previousNeededWeb =
        priorStateNeedsWebSearch(params.priorTurn) ||
        detectWebSearchIntent(previous);
    const routingText = `${previous.slice(0, PRIOR_TURN_CONTEXT_CHAR_LIMIT)}\nFollow-up: ${latest}`;

    return {
        followUpDetected: true,
        inheritedWebSearch: previousNeededWeb,
        routingText,
        inheritanceSource: previousNeededWeb ? "deterministic" : "none",
    };
}

function latestUserHasWorkflow(messages: ChatMessage[]): boolean {
    const lastUser = [...messages].reverse().find((m) => m.role === "user");
    return !!lastUser?.workflow?.id;
}

export function resolveDocumentMode(params: {
    requestedTrackChanges?: boolean;
    messages: ChatMessage[];
    docIndex: DocIndex;
}): DocumentModeDecision {
    const lastUserText = latestUserText(params.messages);
    const docs = Object.values(params.docIndex);
    const editableDocxCount = docs.filter(isDocx).length;
    const pdfCount = docs.filter(isPdf).length;
    const hasEditableDocx = editableDocxCount > 0;

    const commentsOnlyIntent = COMMENTS_ONLY_RE.test(lastUserText);
    const editIntent =
        EDIT_INTENT_RE.test(lastUserText) ||
        COMMENT_EDIT_INTENT_RE.test(lastUserText) ||
        UPDATE_DOCUMENT_INTENT_RE.test(lastUserText);

    if (
        params.requestedTrackChanges === true &&
        hasEditableDocx &&
        !commentsOnlyIntent
    ) {
        return {
            trackChanges: true,
            reason: "explicit_request_true",
            editIntent,
            commentsOnlyIntent,
            editableDocxCount,
            pdfCount,
        };
    }

    if (commentsOnlyIntent) {
        return {
            trackChanges: false,
            reason: "explicit_comments_only",
            editIntent,
            commentsOnlyIntent,
            editableDocxCount,
            pdfCount,
        };
    }

    if (!hasEditableDocx) {
        return {
            trackChanges: false,
            reason: "no_editable_docx",
            editIntent,
            commentsOnlyIntent,
            editableDocxCount,
            pdfCount,
        };
    }

    if (editIntent) {
        return {
            trackChanges: true,
            reason: "auto_docx_edit_intent",
            editIntent,
            commentsOnlyIntent,
            editableDocxCount,
            pdfCount,
        };
    }

    return {
        trackChanges: false,
        reason: "default_comments_only",
        editIntent,
        commentsOnlyIntent,
        editableDocxCount,
        pdfCount,
    };
}

function explicitTurnModeFor(
    requestedMode: Exclude<RequestedAssistantMode, "auto">,
): AssistantTurnMode {
    switch (requestedMode) {
        case "answer":
            return "answer";
        case "comments":
            return "comments";
        case "document_edit":
            return "document_direct_edit";
        case "workflow_review":
            return "workflow_review";
        case "document_generation":
            return "document_generation";
        case "web_search":
            return "web_search";
    }
}

function resolveExplicitTurnMode(
    requestedMode: Exclude<RequestedAssistantMode, "auto">,
    params: {
        messages: ChatMessage[];
        docIndex: DocIndex;
        projectId?: string | null;
    },
): TurnModeDecision {
    const allowTrackChanges = requestedMode === "document_edit";
    let documentMode = resolveDocumentMode({
        requestedTrackChanges: allowTrackChanges,
        messages: params.messages,
        docIndex: params.docIndex,
    });

    // Explicit non-mutating modes must never mutate document bytes. Force the
    // document mode flat regardless of what the legacy heuristics inferred so a
    // user who explicitly selected (e.g.) "answer" can never be silently
    // upgraded to tracked-change editing.
    if (!allowTrackChanges) {
        documentMode = {
            ...documentMode,
            trackChanges: false,
            reason: "explicit_non_mutating_mode",
        };
    }

    const lastUserText = latestUserText(params.messages);
    const hasAttachedDocs = Object.keys(params.docIndex).length > 0;
    const wantsWebSearch = detectWebSearchIntent(lastUserText, hasAttachedDocs);

    if (requestedMode === "web_search" && hasAttachedDocs) {
        return {
            turnMode: params.projectId ? "project_chat" : "answer",
            requestedMode,
            source: "explicit",
            documentMode,
            enableGoogleSearch: false,
            requiresConfirmation: false,
            reason: "web_search_blocked_docs_attached",
        };
    }

    return {
        turnMode: explicitTurnModeFor(requestedMode),
        requestedMode,
        source: "explicit",
        documentMode,
        enableGoogleSearch:
            !hasAttachedDocs &&
            (requestedMode === "web_search" ||
                (requestedMode === "answer" && wantsWebSearch)),
        requiresConfirmation: false,
        reason: "explicit_mode",
    };
}

function resolveAutoTurnMode(params: {
    requestedTrackChanges?: boolean;
    messages: ChatMessage[];
    docIndex: DocIndex;
    isExecutingWorkflow?: boolean;
    projectId?: string | null;
}): TurnModeDecision {
    const documentMode = resolveDocumentMode({
        requestedTrackChanges: params.requestedTrackChanges,
        messages: params.messages,
        docIndex: params.docIndex,
    });

    const lastUserText = latestUserText(params.messages);
    const hasWorkflowAttachment = latestUserHasWorkflow(params.messages);
    const hasAttachedDocs = Object.keys(params.docIndex).length > 0;
    const wantsWebSearch = detectWebSearchIntent(lastUserText, hasAttachedDocs);

    const finalize = (
        turnMode: AssistantTurnMode,
        enableGoogleSearch: boolean,
        reason: TurnModeReason,
    ): TurnModeDecision => ({
        turnMode,
        requestedMode: "auto",
        source: "inferred",
        documentMode,
        enableGoogleSearch,
        requiresConfirmation:
            turnMode === "document_direct_edit" ||
            turnMode === "document_edit_with_web_search",
        reason,
    });

    if (params.isExecutingWorkflow || hasWorkflowAttachment) {
        return finalize("workflow_review", false, "explicit_workflow_execution");
    }

    if (WORKFLOW_REVIEW_RE.test(lastUserText)) {
        return finalize("workflow_review", false, "user_requested_workflow");
    }

    if (
        !hasAttachedDocs &&
        documentMode.trackChanges &&
        documentMode.editableDocxCount > 0 &&
        wantsWebSearch
    ) {
        return finalize(
            "document_edit_with_web_search",
            true,
            "docx_mutation_with_web_search",
        );
    }

    if (documentMode.trackChanges && documentMode.editableDocxCount > 0) {
        return finalize("document_direct_edit", false, "direct_docx_mutation");
    }

    if (DOCUMENT_GENERATION_RE.test(lastUserText)) {
        return finalize(
            "document_generation",
            !hasAttachedDocs && wantsWebSearch,
            "document_generation_intent",
        );
    }

    if (!hasAttachedDocs && wantsWebSearch) {
        return finalize("web_search", true, "web_search_intent");
    }

    return finalize(
        params.projectId ? "project_chat" : "answer",
        false,
        "normal_answer",
    );
}

export function resolveTurnMode(params: {
    requestedMode?: RequestedAssistantMode;
    requestedTrackChanges?: boolean;
    messages: ChatMessage[];
    docIndex: DocIndex;
    isExecutingWorkflow?: boolean;
    projectId?: string | null;
}): TurnModeDecision {
    const requestedMode = params.requestedMode ?? "auto";
    if (requestedMode !== "auto") {
        return resolveExplicitTurnMode(requestedMode, {
            messages: params.messages,
            docIndex: params.docIndex,
            projectId: params.projectId,
        });
    }
    return resolveAutoTurnMode(params);
}

export function applySaudiLawRouting(
    decision: TurnModeDecision,
    saudi: SaudiLawRoutingResult,
    requestedMode: RequestedAssistantMode,
): TurnModeDecision {
    if (!saudi.isSaudiLawQuestion) return decision;

    if (
        decision.turnMode === "workflow_review" ||
        decision.requiresConfirmation
    ) {
        return decision;
    }

    const base = {
        ...decision,
        enableGoogleSearch: true,
        saudiLawRouting: saudi,
    };

    if (decision.turnMode === "document_direct_edit") {
        return {
            ...base,
            turnMode: "document_edit_with_web_search",
            reason: "saudi_law_requires_web_search",
        };
    }

    if (requestedMode === "answer") {
        return {
            ...base,
            reason: "saudi_law_requires_web_search",
        };
    }

    if (decision.enableGoogleSearch) {
        return {
            ...base,
            reason: "saudi_law_requires_web_search",
        };
    }

    if (requestedMode === "auto") {
        return {
            ...base,
            turnMode: "web_search",
            reason: "saudi_law_requires_web_search",
        };
    }

    return {
        ...base,
        reason: "saudi_law_requires_web_search",
    };
}

function applyWebSearchRouting(
    decision: TurnModeDecision,
    requestedMode: RequestedAssistantMode,
    routing: Awaited<ReturnType<typeof classifyWebSearchNeed>>,
): TurnModeDecision {
    if (!routing.needsWebSearch || decision.enableGoogleSearch) return decision;

    if (requestedMode === "answer") {
        return {
            ...decision,
            enableGoogleSearch: true,
            reason: routing.source === "llm" ? "web_search_routing" : decision.reason,
        };
    }

    return {
        ...decision,
        turnMode: "web_search",
        enableGoogleSearch: true,
        reason: routing.source === "llm" ? "web_search_routing" : "web_search_intent",
    };
}

export async function resolveTurnModeAsync(params: {
    requestedMode?: RequestedAssistantMode;
    requestedTrackChanges?: boolean;
    messages: ChatMessage[];
    docIndex: DocIndex;
    isExecutingWorkflow?: boolean;
    projectId?: string | null;
    apiKeys?: UserApiKeys;
    routingTrace?: RunTree | null;
    saudiRoutingTrace?: RunTree | null;
    priorTurn?: PriorTurnRoutingState | null;
    /** Eval/test hook — forwarded to classifySaudiLawNeed. */
    saudiClassifierOverride?: (args: {
        userText: string;
        apiKeys?: UserApiKeys;
        trace?: RunTree | null;
    }) => Promise<string>;
}): Promise<TurnModeDecision> {
    const requestedMode = params.requestedMode ?? "auto";
    let decision = resolveTurnMode(params);
    const lastUserText = latestUserText(params.messages);
    const followUpRouting = buildFollowUpRoutingContext({
        messages: params.messages,
        priorTurn: params.priorTurn,
    });
    const routingText = followUpRouting.routingText || lastUserText;
    const hasAttachedDocs = Object.keys(params.docIndex).length > 0;

    if (
        followUpRouting.inheritedWebSearch &&
        !hasAttachedDocs &&
        !decision.enableGoogleSearch &&
        !decision.requiresConfirmation &&
        decision.turnMode !== "workflow_review" &&
        decision.turnMode !== "document_generation"
    ) {
        decision = {
            ...decision,
            turnMode: requestedMode === "auto" ? "web_search" : decision.turnMode,
            enableGoogleSearch: true,
            reason: "web_search_routing",
        };
    }

    decision = {
        ...decision,
        routingContext: followUpRouting,
    };

    const skipGenericRouting =
        decision.enableGoogleSearch ||
        decision.requiresConfirmation ||
        decision.turnMode === "workflow_review" ||
        requestedMode === "web_search";

    const saudiPromise =
        process.env.SAUDI_LAW_ROUTING !== "0"
            ? classifySaudiLawNeed({
                  userText: routingText,
                  apiKeys: params.apiKeys,
                  trace: params.saudiRoutingTrace ?? params.routingTrace,
                  hasAttachedDocs,
                  classifierOverride: params.saudiClassifierOverride,
              })
            : Promise.resolve(
                  buildSaudiLawRoutingResult({
                      isSaudiLawQuestion: false,
                      source: "skip",
                      domain: "unknown",
                      confidence: "high",
                  }),
              );

    const genericPromise = skipGenericRouting
        ? Promise.resolve(null)
        : classifyWebSearchNeed({
              userText: routingText,
              apiKeys: params.apiKeys,
              trace: params.routingTrace,
              hasAttachedDocs,
          });

    const [saudi, routing] = await Promise.all([
        saudiPromise,
        genericPromise,
    ]);

    if (process.env.SAUDI_LAW_ROUTING_LOG === "1") {
        console.log("[saudiLawRouting]", saudi);
    }

    if (saudi.isSaudiLawQuestion) {
        return applySaudiLawRouting(decision, saudi, requestedMode);
    }

    if (skipGenericRouting || routing === null) {
        return decision;
    }

    return applyWebSearchRouting(decision, requestedMode, routing);
}