Merge branch 'main' into whoisthey/language-model-input-modalities

Author: whoisthey

CHANGELOG.md

@@ -9,9 +9,11 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Changed
 - [EE] Improved Ask Sourcebot prompt caching by splitting static and dynamic prompt sections and advancing cache breakpoints after every agent step instead of only after each message. [#1366](https://github.com/sourcebot-dev/sourcebot/pull/1366)
+- Refactored Ask Sourcebot user message text extraction into a shared helper that robustly handles non-text message parts. [#1371](https://github.com/sourcebot-dev/sourcebot/pull/1371)
 
 ### Added
 - Added per-step token cost tracking and estimated tool call token usage to Ask Sourcebot chat history. [#1353](https://github.com/sourcebot-dev/sourcebot/pull/1353)
+- [EE] Added a context-window usage gauge to the Ask Sourcebot chat details, showing how much of the selected model's context window each turn occupies. Window sizes are resolved from the models.dev catalog. [#1370](https://github.com/sourcebot-dev/sourcebot/pull/1370)
 - Added optional `inputModalities` and `supportedDocumentTypes` configuration for language models, exposing model input-modality and document capabilities (defaults to text-only, no documents). [#1372](https://github.com/sourcebot-dev/sourcebot/pull/1372)
 
 ### Fixed

packages/web/src/app/(app)/chat/[id]/page.tsx

@@ -14,6 +14,7 @@ import { __unsafePrisma } from '@/prisma';
 import { ChatVisibility } from '@sourcebot/db';
 import { Metadata } from 'next';
 import { SBChatMessage } from '@/features/chat/types';
+import { getUserMessageText } from '@/features/chat/utils';
 import { env } from '@sourcebot/shared';
 import { hasEntitlement } from '@/lib/entitlements';
 import { ChatEntitlementMessage } from '@/features/chat/components/chatEntitlementMessage';
@@ -54,11 +55,11 @@ export const generateMetadata = async ({ params }: PageProps): Promise<Metadata>
 
     let description = 'A chat on Sourcebot';
     if (firstUserMessage) {
-        const textPart = firstUserMessage.parts.find(p => p.type === 'text');
-        if (textPart && textPart.type === 'text') {
-            description = textPart.text.length > 160
-                ? textPart.text.substring(0, 160).trim() + '...'
-                : textPart.text;
+        const text = getUserMessageText(firstUserMessage);
+        if (text) {
+            description = text.length > 160
+                ? text.substring(0, 160).trim() + '...'
+                : text;
         }
     }
 

packages/web/src/app/api/(server)/ee/chat/route.ts

@@ -6,6 +6,7 @@ import { additionalChatRequestParamsSchema } from "@/features/chat/types";
 import { getLanguageModelKey } from "@/features/chat/utils";
 import { checkAskEntitlement, getConfiguredLanguageModels, isOwnerOfChat, updateChatMessages } from "@/features/chat/utils.server";
 import { getAISDKLanguageModelAndOptions } from "@/features/chat/llm.server";
+import { resolveContextWindow } from "@/features/chat/modelContextWindow.server";
 import { apiHandler } from "@/lib/apiHandler";
 import { ErrorCode } from "@/lib/errorCodes";
 import { captureEvent } from "@/lib/posthog";
@@ -89,6 +90,11 @@ export const POST = apiHandler(async (req: NextRequest) => {
 
             const { model, providerOptions, temperature } = await getAISDKLanguageModelAndOptions(languageModelConfig);
 
+            // Total context window for the selected model, used as the
+            // denominator for the UI's context-usage gauge. Undefined when
+            // unknown (e.g. self-hosted models).
+            const contextWindow = await resolveContextWindow(languageModelConfig);
+
             // No-op for non-Anthropic providers / when caching is disabled, so
             // it never perturbs other providers' requests.
             const promptCacheStrategy = getPromptCacheStrategy(
@@ -139,6 +145,7 @@ export const POST = apiHandler(async (req: NextRequest) => {
                 disabledMcpServerIds,
                 model,
                 modelName: languageModelConfig.displayName ?? languageModelConfig.model,
+                contextWindow,
                 promptCacheStrategy,
                 modelProviderOptions: providerOptions,
                 modelTemperature: temperature,

packages/web/src/ee/features/chat/agent.ts

@@ -22,7 +22,7 @@ import { randomUUID } from "crypto";
 import _dedent from "dedent";
 import { ANSWER_TAG, FILE_REFERENCE_PREFIX } from "@/features/chat/constants";
 import { Source } from "@/features/chat/types";
-import { addLineNumbers, fileReferenceToString, getAnswerPartFromAssistantMessage, getTurnProgressState } from "@/features/chat/utils";
+import { addLineNumbers, fileReferenceToString, getAnswerPartFromAssistantMessage, getTurnProgressState, getUserMessageText } from "@/features/chat/utils";
 import { createTools } from "./tools";
 import { getConnectedMcpClients } from "@/ee/features/chat/mcp/mcpClientFactory";
 import { getMcpTools, McpToolsResult } from "@/ee/features/chat/mcp/mcpToolSets";
@@ -54,6 +54,7 @@ interface CreateMessageStreamResponseProps {
     disabledMcpServerIds?: string[];
     model: AISDKLanguageModelV3;
     modelName: string;
+    contextWindow?: number;
     promptCacheStrategy: PromptCacheStrategy;
     onFinish: UIMessageStreamOnFinishCallback<SBChatMessage>;
     onError: (error: unknown) => string;
@@ -73,6 +74,7 @@ export const createMessageStream = async ({
     disabledMcpServerIds,
     model,
     modelName,
+    contextWindow,
     promptCacheStrategy,
     modelProviderOptions,
     modelTemperature,
@@ -105,7 +107,7 @@ export const createMessageStream = async ({
             if (message.role === 'user') {
                 return {
                     role: 'user',
-                    content: message.parts[0].type === 'text' ? message.parts[0].text : '',
+                    content: getUserMessageText(message),
                 };
             }
 
@@ -279,6 +281,7 @@ export const createMessageStream = async ({
                     // phases so earlier phases' steps are preserved in order.
                     stepTokenUsage: [...(priorMetadata?.stepTokenUsage ?? []), ...stepTokenUsage],
                     modelName,
+                    contextWindow,
                     traceId,
                 }
             });

packages/web/src/ee/features/chat/components/chatThread/chatThread.tsx

@@ -5,7 +5,7 @@ import { Button } from '@/components/ui/button';
 import { Separator } from '@/components/ui/separator';
 import { CustomSlateEditor } from '@/features/chat/customSlateEditor';
 import { AdditionalChatRequestParams, CustomEditor, LanguageModelInfo, SBChatMessage, SearchScope, Source } from '@/features/chat/types';
-import { createUIMessage, getAllMentionElements, getTurnProgressState, resetEditor, slateContentToString } from '@/features/chat/utils';
+import { createUIMessage, getAllMentionElements, getTurnProgressState, getUserMessageText, resetEditor, slateContentToString } from '@/features/chat/utils';
 import { useChat } from '@ai-sdk/react';
 import { CreateUIMessage, DefaultChatTransport, lastAssistantMessageIsCompleteWithApprovalResponses } from 'ai';
 import { ArrowDownIcon, CopyIcon } from 'lucide-react';
@@ -204,16 +204,16 @@ export const ChatThread = ({
             } satisfies AdditionalChatRequestParams,
         });
 
+        const userMessageText = getUserMessageText(message);
         if (
             messages.length === 0 &&
-            message.parts.length > 0 &&
-            message.parts[0].type === 'text'
+            userMessageText.length > 0
         ) {
             generateAndUpdateChatNameFromMessage(
                 {
                     chatId,
                     languageModelId: selectedLanguageModel.model,
-                    message: message.parts[0].text,
+                    message: userMessageText,
                 },
             ).then((response) => {
                 if (isServiceError(response)) {

packages/web/src/ee/features/chat/components/chatThread/chatThreadListItem.tsx

@@ -8,7 +8,7 @@ import { CSSProperties, forwardRef, memo, useCallback, useEffect, useMemo, useRe
 import scrollIntoView from 'scroll-into-view-if-needed';
 import { Reference, referenceSchema, SBChatMessage, Source } from "@/features/chat/types";
 import { useExtractReferences } from '../../useExtractReferences';
-import { getAnswerPartFromAssistantMessage, getLastStepParts, groupMessageIntoSteps, isSBChatToolPart, repairReferences, tryResolveFileReference } from '@/features/chat/utils';
+import { getAnswerPartFromAssistantMessage, getLastStepParts, getUserMessageText, groupMessageIntoSteps, isSBChatToolPart, repairReferences, tryResolveFileReference } from '@/features/chat/utils';
 import { AnswerCard } from './answerCard';
 import { DetailsCard } from './detailsCard';
 import { ApprovalRequestedToolPart, ToolApprovalBanner } from './toolApprovalBanner';
@@ -49,7 +49,7 @@ const ChatThreadListItemComponent = forwardRef<HTMLDivElement, ChatThreadListIte
     const userHasManuallyExpanded = useRef(false);
 
     const userQuestion = useMemo(() => {
-        return userMessage.parts.length > 0 && userMessage.parts[0].type === 'text' ? userMessage.parts[0].text : '';
+        return getUserMessageText(userMessage);
     }, [userMessage]);
 
     // Take the assistant message and repair any references that are not properly formatted.

packages/web/src/ee/features/chat/components/chatThread/detailsCard.tsx

@@ -86,6 +86,18 @@ const DetailsCardComponent = ({
         ? Math.round((cacheReadTokens / inputTokens) * 100)
         : 0;
 
+    // Context-window usage gauge. "In use" is the input the model saw on its
+    // most recent step — i.e. the full accumulated prompt occupying the window
+    // right now — not the cumulative totalInputTokens.
+    const stepTokenUsage = metadata?.stepTokenUsage;
+    const currentContextTokens = stepTokenUsage && stepTokenUsage.length > 0
+        ? stepTokenUsage[stepTokenUsage.length - 1].inputTokens
+        : undefined;
+    const contextWindow = metadata?.contextWindow;
+    const contextUsagePercent = currentContextTokens !== undefined && contextWindow !== undefined && contextWindow > 0
+        ? Math.min(100, Math.round((currentContextTokens / contextWindow) * 100))
+        : undefined;
+
     const handleExpandedChanged = useCallback((next: boolean) => {
         captureEvent('wa_chat_details_card_toggled', { chatId, isExpanded: next });
         onExpandedChanged(next);
@@ -193,6 +205,23 @@ const DetailsCardComponent = ({
                                                 )}
                                             </div>
                                         )}
+                                        {contextUsagePercent !== undefined && currentContextTokens !== undefined && contextWindow !== undefined && (
+                                            <Tooltip>
+                                                <TooltipTrigger asChild>
+                                                    <div className="cursor-help">
+                                                        <ContextWindowGauge
+                                                            total={contextWindow}
+                                                            percent={contextUsagePercent}
+                                                        />
+                                                    </div>
+                                                </TooltipTrigger>
+                                                <TooltipContent side="bottom">
+                                                    <div className="max-w-xs text-xs">
+                                                        The most recent step&apos;s prompt used {currentContextTokens.toLocaleString()} of the model&apos;s {contextWindow.toLocaleString()}-token context window ({contextUsagePercent}%).
+                                                    </div>
+                                                </TooltipContent>
+                                            </Tooltip>
+                                        )}
                                         {metadata?.totalResponseTimeMs && (
                                             <div className="flex items-center text-xs">
                                                 <Clock className="w-3 h-3 mr-1 flex-shrink-0" />
@@ -367,6 +396,61 @@ const StepTokenUsage = ({ usage, label = 'step' }: { usage: StepTokenUsageEntry,
     );
 }
 
+
+const CONTEXT_USAGE_YELLOW_PERCENT = 50;
+const CONTEXT_USAGE_RED_PERCENT = 80;
+
+const getContextUsageColorClass = (percent: number): string => {
+    if (percent >= CONTEXT_USAGE_RED_PERCENT) {
+        return "text-red-500";
+    }
+    if (percent >= CONTEXT_USAGE_YELLOW_PERCENT) {
+        return "text-yellow-500";
+    }
+    return "text-[#6cb38f]";
+};
+
+const ContextWindowGauge = ({ total, percent }: { total: number, percent: number }) => {
+    const size = 14;
+    const strokeWidth = 2;
+    const radius = (size - strokeWidth) / 2;
+    const circumference = 2 * Math.PI * radius;
+    const dashOffset = circumference * (1 - Math.min(100, percent) / 100);
+    const colorClass = getContextUsageColorClass(percent);
+
+    return (
+        <div className="flex items-center gap-1.5 text-xs whitespace-nowrap">
+            <svg width={size} height={size} className="-rotate-90 flex-shrink-0">
+                {/* Neutral gray track. */}
+                <circle
+                    cx={size / 2}
+                    cy={size / 2}
+                    r={radius}
+                    fill="none"
+                    stroke="currentColor"
+                    strokeWidth={strokeWidth}
+                    className="text-zinc-500"
+                />
+                {/* Progress arc. */}
+                <circle
+                    cx={size / 2}
+                    cy={size / 2}
+                    r={radius}
+                    fill="none"
+                    stroke="currentColor"
+                    strokeWidth={strokeWidth}
+                    strokeLinecap="round"
+                    strokeDasharray={circumference}
+                    strokeDashoffset={dashOffset}
+                    className={cn("transition-all duration-300", colorClass)}
+                />
+            </svg>
+            <span className={cn("font-semibold", colorClass)}>{percent}%</span>
+            <span className="text-muted-foreground">of {getShortenedNumberDisplayString(total, 0).toUpperCase()}</span>
+        </div>
+    );
+}
+
 type GuardedToolType =
     | 'tool-read_file'
     | 'tool-grep'

packages/web/src/ee/features/mcp/askCodebase.ts

@@ -2,6 +2,7 @@ import { sew } from "@/middleware/sew";
 import { getConfiguredLanguageModels, updateChatMessages, checkAskEntitlement } from "@/features/chat/utils.server";
 import { generateChatNameFromMessage } from "@/ee/features/chat/llm.server";
 import { getAISDKLanguageModelAndOptions } from "@/features/chat/llm.server";
+import { resolveContextWindow } from "@/features/chat/modelContextWindow.server";
 import { LanguageModelInfo, SBChatMessage, SearchScope } from "@/features/chat/types";
 import { convertLLMOutputToPortableMarkdown, getAnswerPartFromAssistantMessage, getLanguageModelKey } from "@/features/chat/utils";
 import { resolveModelInputModalities, resolveModelSupportedDocumentTypes } from "@/features/chat/modelCapabilities";
@@ -85,6 +86,7 @@ export const askCodebase = (params: AskCodebaseParams): Promise<AskCodebaseResul
 
             const { model, providerOptions, temperature } = await getAISDKLanguageModelAndOptions(languageModelConfig);
             const modelName = languageModelConfig.displayName ?? languageModelConfig.model;
+            const contextWindow = await resolveContextWindow(languageModelConfig);
 
             // No-op for non-Anthropic providers / when caching is disabled.
             const promptCacheStrategy = getPromptCacheStrategy(
@@ -183,6 +185,7 @@ export const askCodebase = (params: AskCodebaseParams): Promise<AskCodebaseResul
                 prisma,
                 model,
                 modelName,
+                contextWindow,
                 promptCacheStrategy,
                 modelProviderOptions: providerOptions,
                 modelTemperature: temperature,