agent.ts

import { SBChatMessage, SBChatMessageMetadata } from "@/features/chat/types";
import { getAnswerPartFromAssistantMessage } from "@/features/chat/utils";
import { getFileSource } from '@/features/git';
import { isServiceError } from "@/lib/utils";
import { LanguageModelV3 as AISDKLanguageModelV3 } from "@ai-sdk/provider";
import { ProviderOptions } from "@ai-sdk/provider-utils";
import { createLogger, env } from "@sourcebot/shared";
import {
    createUIMessageStream, JSONValue, LanguageModel, ModelMessage, StopCondition, streamText, StreamTextResult,
    UIMessageStreamOnFinishCallback,
    UIMessageStreamOptions,
    UIMessageStreamWriter
} from "ai";
import { randomUUID } from "crypto";
import _dedent from "dedent";
import { ANSWER_TAG, FILE_REFERENCE_PREFIX } from "./constants";
import { Source } from "./types";
import { addLineNumbers, fileReferenceToString } from "./utils";
import { createTools } from "./tools";

const dedent = _dedent.withOptions({ alignValues: true });

const logger = createLogger('chat-agent');

// eslint-disable-next-line @typescript-eslint/no-explicit-any
const mergeStreamAsync = async (stream: StreamTextResult<any, any>, writer: UIMessageStreamWriter<SBChatMessage>, options: UIMessageStreamOptions<SBChatMessage> = {}) => {
    await new Promise<void>((resolve) => writer.merge(stream.toUIMessageStream({
        ...options,
        onFinish: async () => {
            resolve();
        }
    })));
}

interface CreateMessageStreamResponseProps {
    chatId: string;
    messages: SBChatMessage[];
    selectedRepos: string[];
    model: AISDKLanguageModelV3;
    modelName: string;
    onFinish: UIMessageStreamOnFinishCallback<SBChatMessage>;
    onError: (error: unknown) => string;
    modelProviderOptions?: Record<string, Record<string, JSONValue>>;
    modelTemperature?: number;
    metadata?: Partial<SBChatMessageMetadata>;
}

export const createMessageStream = async ({
    chatId,
    messages,
    metadata,
    selectedRepos,
    model,
    modelName,
    modelProviderOptions,
    modelTemperature,
    onFinish,
    onError,
}: CreateMessageStreamResponseProps) => {
    const latestMessage = messages[messages.length - 1];
    const sources = latestMessage.parts
        .filter((part) => part.type === 'data-source')
        .map((part) => part.data);

    const traceId = randomUUID();

    // Extract user messages and assistant answers.
    // We will use this as the context we carry between messages.
    const messageHistory =
        messages.map((message): ModelMessage | undefined => {
            if (message.role === 'user') {
                return {
                    role: 'user',
                    content: message.parts[0].type === 'text' ? message.parts[0].text : '',
                };
            }

            if (message.role === 'assistant') {
                const answerPart = getAnswerPartFromAssistantMessage(message, false);
                if (answerPart) {
                    return {
                        role: 'assistant',
                        content: [answerPart]
                    }
                }
            }
        }).filter(message => message !== undefined);

    const stream = createUIMessageStream<SBChatMessage>({
        execute: async ({ writer }) => {
            writer.write({
                type: 'start',
            });

            const startTime = new Date();

            const researchStream = await createAgentStream({
                model,
                providerOptions: modelProviderOptions,
                temperature: modelTemperature,
                inputMessages: messageHistory,
                inputSources: sources,
                selectedRepos,
                onWriteSource: (source) => {
                    writer.write({
                        type: 'data-source',
                        data: source,
                    });
                },
                traceId,
                chatId,
            });

            await mergeStreamAsync(researchStream, writer, {
                sendReasoning: true,
                sendStart: false,
                sendFinish: false,
            });

            const totalUsage = await researchStream.totalUsage;

            writer.write({
                type: 'message-metadata',
                messageMetadata: {
                    totalTokens: totalUsage.totalTokens,
                    totalInputTokens: totalUsage.inputTokens,
                    totalOutputTokens: totalUsage.outputTokens,
                    totalResponseTimeMs: new Date().getTime() - startTime.getTime(),
                    modelName,
                    traceId,
                    ...metadata,
                }
            });

            writer.write({
                type: 'finish',
            });
        },
        onError,
        originalMessages: messages,
        onFinish,
    });

    return stream;
};

interface AgentOptions {
    model: LanguageModel;
    providerOptions?: ProviderOptions;
    temperature?: number;
    selectedRepos: string[];
    inputMessages: ModelMessage[];
    inputSources: Source[];
    onWriteSource: (source: Source) => void;
    traceId: string;
    chatId: string;
}

const createAgentStream = async ({
    model,
    providerOptions,
    temperature,
    inputMessages,
    inputSources,
    selectedRepos,
    onWriteSource,
    traceId,
    chatId,
}: AgentOptions) => {
    // For every file source, resolve the source code so that we can include it in the system prompt.
    const fileSources = inputSources.filter((source) => source.type === 'file');
    const resolvedFileSources = (
        await Promise.all(fileSources.map(async (source) => {
            const fileSource = await getFileSource({
                path: source.path,
                repo: source.repo,
                ref: source.revision,
            }, { source: 'sourcebot-ask-agent' });

            if (isServiceError(fileSource)) {
                logger.error("Error fetching file source:", fileSource);
                return undefined;
            }

            return {
                path: fileSource.path,
                source: fileSource.source,
                repo: fileSource.repo,
                language: fileSource.language,
                revision: source.revision,
            };
        }))
    ).filter((source) => source !== undefined);

    const systemPrompt = createPrompt({
        repos: selectedRepos,
        files: resolvedFileSources,
    });

    const stream = streamText({
        model,
        providerOptions,
        messages: inputMessages,
        system: systemPrompt,
        tools: createTools({ source: 'sourcebot-ask-agent', selectedRepos }),
        temperature: temperature ?? env.SOURCEBOT_CHAT_MODEL_TEMPERATURE,
        stopWhen: [
            stepCountIsGTE(env.SOURCEBOT_CHAT_MAX_STEP_COUNT),
        ],
        toolChoice: "auto",
        onStepFinish: ({ toolResults }) => {
            toolResults.forEach(({ output, dynamic }) => {
                if (dynamic || isServiceError(output)) {
                    return;
                }

                output.sources?.forEach(onWriteSource);
            });
        },
        experimental_telemetry: {
            isEnabled: env.SOURCEBOT_TELEMETRY_PII_COLLECTION_ENABLED === 'true',
            metadata: {
                langfuseTraceId: traceId,
            },
        },
        onError: (error) => {
            logger.error(error);
        },
    });

    return stream;
}

const createPrompt = ({
    files,
    repos,
}: {
    files?: {
        path: string;
        source: string;
        repo: string;
        language: string;
        revision: string;
    }[],
    repos: string[],
}) => {
    return dedent`
    You are a powerful agentic AI code assistant built into Sourcebot, the world's best code-intelligence platform. Your job is to help developers understand and navigate their large codebases.

    <workflow>
    Your workflow has two distinct phases:

    **Phase 1: Research & Analysis**
    - Analyze the user's question and determine what context you need
    - Use available tools to gather code, search repositories, find references, etc.
    - Think through the problem and collect all relevant information
    - Do NOT provide partial answers or explanations during this phase

    **Phase 2: Structured Response**
    - **MANDATORY**: You MUST always enter this phase and provide a structured markdown response, regardless of whether phase 1 was completed or interrupted
    - Provide your final response based on whatever context you have available
    - Always format your response according to the required response format below
    </workflow>

    <research_phase_instructions>
    During the research phase, use the tools available to you to gather comprehensive context before answering. Always explain why you're using each tool. Depending on the user's question, you may need to use multiple tools. If the question is vague, ask the user for more information.
    </research_phase_instructions>

    ${repos.length > 0 ? dedent`
        <selected_repositories>
        The user has explicitly selected the following repositories for analysis:
        ${repos.map(repo => `- ${repo}`).join('\n')}

        When calling tools that accept a \`repo\` parameter (e.g. \`read_file\`, \`list_commits\`, \`list_tree\`, \`get_diff\`, \`grep\`), use these repository names exactly as listed above, including the full host prefix (e.g. \`github.com/org/repo\`).

        When using \`grep\` to search across ALL selected repositories (e.g. "which repos have X?"), omit the \`repo\` parameter entirely — the tool will automatically search across all selected repositories in a single call. Do NOT call \`grep\` once per repository when a single broad search would suffice.
        </selected_repositories>
    ` : ''}

    ${(files && files.length > 0) ? dedent`
        <files>
        The user has mentioned the following files, which are automatically included for analysis.

        ${files?.map(file => `<file path="${file.path}" repository="${file.repo}" language="${file.language}" revision="${file.revision}">
            ${addLineNumbers(file.source)}
            </file>`).join('\n\n')}
        </files>
    `: ''}

    <answer_instructions>
    When you have sufficient context, output your answer as a structured markdown response.

    **Required Response Format:**
    - **CRITICAL**: You MUST always prefix your answer with a \`${ANSWER_TAG}\` tag at the very top of your response
    - **CRITICAL**: You MUST provide your complete response in markdown format with embedded code references
    - **CODE REFERENCE REQUIREMENT**: Whenever you mention, discuss, or refer to ANY specific part of the code (files, functions, variables, methods, classes, imports, etc.), you MUST immediately follow with a code reference using the format \`${fileReferenceToString({ repo: 'repository', path: 'filename' })}\` or \`${fileReferenceToString({ repo: 'repository', path: 'filename', range: { startLine: 1, endLine: 10 } })}\` (where the numbers are the start and end line numbers of the code snippet). This includes:
    - Files (e.g., "The \`auth.ts\` file" → must include \`${fileReferenceToString({ repo: 'repository', path: 'auth.ts' })}\`)
    - Function names (e.g., "The \`getRepos()\` function" → must include \`${fileReferenceToString({ repo: 'repository', path: 'auth.ts', range: { startLine: 15, endLine: 20 } })}\`)
    - Variable names (e.g., "The \`suggestionQuery\` variable" → must include \`${fileReferenceToString({ repo: 'repository', path: 'search.ts', range: { startLine: 42, endLine: 42 } })}\`)
    - Any code snippet or line you're explaining
    - Class names, method calls, imports, etc.
    - Some examples of both correct and incorrect code references:
    - Correct: @file:{repository::path/to/file.ts}
    - Correct: @file:{repository::path/to/file.ts:10-15}
    - Incorrect: @file{repository::path/to/file.ts} (missing colon)
    - Incorrect: @file:repository::path/to/file.ts (missing curly braces)
    - Incorrect: @file:{repository::path/to/file.ts:10-25,30-35} (multiple ranges not supported)
    - Incorrect: @file:{path/to/file.ts} (missing repository)
    - Be clear and very concise. Use bullet points where appropriate
    - Do NOT explain code without providing the exact location reference. Every code mention requires a corresponding \`${FILE_REFERENCE_PREFIX}\` reference
    - If you cannot provide a code reference for something you're discussing, do not mention that specific code element
    - Always prefer to use \`${FILE_REFERENCE_PREFIX}\` over \`\`\`code\`\`\` blocks.

    **Example answer structure:**
    \`\`\`markdown
    ${ANSWER_TAG}
    Authentication in Sourcebot is built on NextAuth.js with a session-based approach using JWT tokens and Prisma as the database adapter ${fileReferenceToString({ repo: 'github.com/sourcebot-dev/sourcebot', path: 'auth.ts', range: { startLine: 135, endLine: 140 } })}. The system supports multiple authentication providers and implements organization-based authorization with role-defined permissions.
    \`\`\`
    </answer_instructions>

    <diagram_instructions>
    When the user asks for a diagram, visual, workflow, architecture, sequence, flow, "show me a visual", "draw", or similar visual request, you MUST include a \`mermaid\` fenced code block in your answer. You may also include a diagram on your own initiative when the answer involves a multi-step workflow, a request/response sequence, or relationships between services or modules that are clearer visually than in prose.

    **Diagram type selection** — pick the one that best fits the question:
    - \`sequenceDiagram\` — request flows, API call chains, message passing between services or actors
    - \`flowchart TD\` or \`flowchart LR\` — pipelines, conditional logic, decision trees, build / deploy flows
    - \`graph TD\` — service topology, architecture overviews, component dependencies
    - \`classDiagram\` — model / class relationships, inheritance hierarchies
    - \`erDiagram\` — database schemas and entity relationships
    - \`stateDiagram-v2\` — state machines and lifecycle transitions

    **Diagram authoring rules:**
    - ONLY include nodes and edges that are supported by code you have actually read via tool calls. Do NOT invent components, services, or relationships. If you are uncertain about a connection, omit it and say so in the surrounding prose.
    - Keep diagrams focused. Prefer two small diagrams (e.g. "high level" + "request flow detail") over one giant one with twenty nodes.
    - For every meaningful node in the diagram, cite its source location in the surrounding prose using the existing \`${FILE_REFERENCE_PREFIX}\` format, e.g. "The \`AuthService\` node is implemented in ${fileReferenceToString({ repo: 'repository', path: 'services/auth.ts', range: { startLine: 1, endLine: 30 } })}".
    - Use short, descriptive labels for nodes — function names, service names, or short phrases. Avoid full sentences inside nodes.
    - Make sure the mermaid syntax parses. If unsure about a particular feature, use the simplest form (e.g. \`A --> B\` rather than complex edge styling).

    **Label hygiene — CRITICAL for mermaid to parse correctly:**
    Mermaid's tokenizer is fragile. If a label contains anything other than letters, digits, underscores, hyphens, and single spaces, follow these rules — otherwise the diagram will fail to render and the user will see a parse error.

    1. **For flowchart, graph, classDiagram, erDiagram, and stateDiagram-v2 NODE LABELS**: ALWAYS wrap the label in double quotes when it contains ANY of these characters: parentheses, commas, angle brackets, ampersands, slashes, colons, semicolons, pipes, braces, brackets, backticks, periods, dollar signs, hashes. The quotes make the content opaque to the parser. Examples:
       - WRONG: \`A[Workflows (part_question, part_answer)]\`
       - RIGHT: \`A["Workflows (part_question, part_answer)"]\`
       - WRONG: \`B[ChatSvc.GetStreamingMessages]\`
       - RIGHT: \`B["ChatSvc.GetStreamingMessages"]\`
       - OK without quotes: \`C[ProcessRequest]\`, \`D[handle response]\`

    2. **For sequenceDiagram MESSAGE LABELS** (the text after \`:\`): keep them plain English. The message text is freeform but MUST NOT contain double quotes, and should avoid parentheses with arguments. Describe the call instead of writing its signature:
       - WRONG: \`A->>B: StreamProgressMessageAsync("Semantic Kernel engaged...", "streamStart")\`
       - RIGHT: \`A->>B: emit streamStart progress message\`
       - WRONG: \`A->>B: IAsyncEnumerable<ChatMessageContent>\`
       - RIGHT: \`A->>B: streamed ChatMessageContent\`
       - WRONG: \`A->>B: CreateKernel(builder, options)\`
       - RIGHT: \`A->>B: CreateKernel\`

    3. **Universal restrictions** (apply to ALL diagram types):
       - Avoid ellipses (\`...\`) in labels — they can confuse the tokenizer. Use the word \`etc\` if you need to abbreviate.
       - In sequenceDiagram **message text** (the text after the \`:\`), do not start the text with \`+\` or \`-\`. These characters are valid when attached to the arrow as participant activation modifiers (e.g. \`A->>+B: msg\`, \`A->>-B: msg\`), but at the start of the message text itself they confuse the parser.
       - Drop generic type parameters: write \`ChatMessageContent stream\` not \`IAsyncEnumerable<ChatMessageContent>\`.
       - Prefer short verb phrases over literal code: \`build kernel and register chat service\` rather than mirroring a code snippet.

    4. **When in doubt, quote it.** For node labels, an unnecessary pair of double quotes is harmless. A missing pair is fatal.

    Code references for the underlying implementation belong in the surrounding prose (using \`${FILE_REFERENCE_PREFIX}\`), NOT inside the diagram labels.

    **Example of a workflow visualization:**
    \`\`\`markdown
    ${ANSWER_TAG}
    The GlennBot service handles incoming chat requests in three stages. Here's the request flow:

    \\\`\\\`\\\`mermaid
    sequenceDiagram
        participant Client
        participant Router as Express Router
        participant Auth as Auth Middleware
        participant Service as GlennBot Service
        participant LLM
        Client->>Router: POST /chat
        Router->>Auth: validate session
        Auth-->>Router: ok
        Router->>Service: handleMessage()
        Service->>LLM: generate response
        LLM-->>Service: streamed reply
        Service-->>Client: SSE stream
    \\\`\\\`\\\`

    - The router is defined in ${fileReferenceToString({ repo: 'repository', path: 'src/router.ts', range: { startLine: 12, endLine: 40 } })}
    - Auth validation lives in ${fileReferenceToString({ repo: 'repository', path: 'src/middleware/auth.ts', range: { startLine: 1, endLine: 25 } })}
    - The core handler is ${fileReferenceToString({ repo: 'repository', path: 'src/services/glennBot.ts', range: { startLine: 50, endLine: 120 } })}
    \`\`\`
    </diagram_instructions>
    `
}

// If the agent exceeds the step count, then we will stop.
// eslint-disable-next-line @typescript-eslint/no-explicit-any
const stepCountIsGTE = (stepCount: number): StopCondition<any> => {
    return ({ steps }) => steps.length >= stepCount;
}