feat: Jan manages Context Allocation dynamically

web-app/src/containers/MessageItem.tsx

@@ -1,5 +1,5 @@
 /* eslint-disable @typescript-eslint/no-explicit-any */
 import { memo, useState, useCallback } from 'react'
 import type { UIMessage, ChatStatus } from 'ai'
 import { RenderMarkdown } from './RenderMarkdown'
 import { cn } from '@/lib/utils'
@@ -9,33 +9,33 @@
   ReasoningContent,
   ReasoningTrigger,
 } from '@/components/ai-elements/reasoning'
 import {
   Tool,
   ToolContent,
   ToolHeader,
   ToolInput,
   ToolOutput,
 } from '@/components/ai-elements/tool'
 import { CopyButton } from './CopyButton'
 import { useModelProvider } from '@/hooks/useModelProvider'
 import { IconRefresh, IconPaperclip } from '@tabler/icons-react'
 import { EditMessageDialog } from '@/containers/dialogs/EditMessageDialog'
 import { DeleteMessageDialog } from '@/containers/dialogs/DeleteMessageDialog'
 import TokenSpeedIndicator from '@/containers/TokenSpeedIndicator'
 import { extractFilesFromPrompt, FileMetadata } from '@/lib/fileMetadata'
 import { useMemo } from 'react'
 import { Button } from '@/components/ui/button'
 
 const CHAT_STATUS = {
   STREAMING: 'streaming',
   SUBMITTED: 'submitted',
 } as const
 
 const CONTENT_TYPE = {
   TEXT: 'text',
   FILE: 'file',
   REASONING: 'reasoning',
 } as const
 
 export type MessageItemProps = {
   message: UIMessage
@@ -48,35 +48,39 @@
   onDelete?: (messageId: string) => void
   assistant?: { avatar?: React.ReactNode; name?: string }
   showAssistant?: boolean
+  isAnimating?: boolean
+  hideActions?: boolean
 }
 
 export const MessageItem = memo(
   ({
     message,
     isLastMessage,
     status,
+    isAnimating,
+    hideActions,
     reasoningContainerRef,
     onRegenerate,
     onEdit,
     onDelete,
   }: MessageItemProps) => {
     const selectedModel = useModelProvider((state) => state.selectedModel)
     const [previewImage, setPreviewImage] = useState<{
       url: string
       filename?: string
     } | null>(null)
 
 
     const handleRegenerate = useCallback(() => {
       onRegenerate?.(message.id)
     }, [onRegenerate, message.id])
 
     const handleEdit = useCallback(
       (newText: string) => {
         onEdit?.(message.id, newText)
       },
       [onEdit, message.id]
     )
 
     const handleDelete = useCallback(() => {
       onDelete?.(message.id)
@@ -185,6 +189,7 @@
                 content={part.text}
                 isStreaming={isStreaming && isLastPart}
                 messageId={message.id}
+                isAnimating={isAnimating}
               />
             </>
           )}
@@ -345,7 +350,7 @@
         })}
 
         {/* Message actions for user messages */}
-        {message.role === 'user' && (
+        {message.role === 'user' && !hideActions && (
           <div className="flex items-center justify-end gap-1 text-muted-foreground text-xs mt-4">
             <CopyButton text={getFullTextContent()} />
 
@@ -369,7 +374,7 @@
               <div
                 className={cn(
                   'flex items-center gap-1',
-                  isStreaming && 'hidden'
+                  (isStreaming || hideActions) && 'hidden'
                 )}
               >
                 <CopyButton text={getFullTextContent()} />
@@ -434,7 +439,8 @@
       prevProps.isFirstMessage === nextProps.isFirstMessage &&
       prevProps.isLastMessage === nextProps.isLastMessage &&
       prevProps.status === nextProps.status &&
-      prevProps.showAssistant === nextProps.showAssistant
+      prevProps.showAssistant === nextProps.showAssistant &&
+      prevProps.hideActions === nextProps.hideActions
     )
   }
 )

web-app/src/containers/RenderMarkdown.tsx

@@ -21,6 +21,7 @@ interface MarkdownProps {
   isUser?: boolean
   isStreaming?: boolean
   messageId?: string
+  isAnimating?: boolean
 }
 
 // Cache for normalized LaTeX content
@@ -85,6 +86,7 @@ function RenderMarkdownComponent({
   isUser,
   components,
   messageId,
+  isAnimating
 }: MarkdownProps) {
 
   // Memoize the normalized content to avoid reprocessing on every render
@@ -101,7 +103,7 @@ function RenderMarkdownComponent({
       )}
     >
       <Streamdown
-        animate={true}
+        animate={isAnimating ?? true}
         animationDuration={500}
         linkSafety={{
           enabled: false,

web-app/src/hooks/use-chat.ts

@@ -126,6 +126,10 @@ export function useChat(
     }
   }, [mcpToolNames, ragToolNames])
 
+  const setContinueFromContent = useCallback((content: string) => {
+    transportRef.current?.setContinueFromContent(content)
+  }, [])
+
   // Expose method to update RAG tools availability
   const updateRagToolsAvailability = useCallback(
     async (
@@ -147,5 +151,6 @@ export function useChat(
   return {
     ...chatResult,
     updateRagToolsAvailability,
+    setContinueFromContent,
   }
 }

web-app/src/lib/custom-chat-transport.ts

@@ -48,6 +48,38 @@ export type ServiceHub = {
   }
 }
 
+/**
+ * Wraps a UIMessageChunk stream so that when the first `text-start` chunk
+ * arrives, a `text-delta` carrying `prefixText` is immediately injected into
+ * the same text block. This makes the new message show the partial content
+ * right away while continuation tokens stream in after it.
+ */
+function prependTextDeltaToUIStream(
+  stream: ReadableStream<UIMessageChunk>,
+  prefixText: string
+): ReadableStream<UIMessageChunk> {
+  const reader = stream.getReader()
+  let prefixEmitted = false
+  return new ReadableStream<UIMessageChunk>({
+    async pull(controller) {
+      const { done, value } = await reader.read()
+      if (done) {
+        controller.close()
+        return
+      }
+      controller.enqueue(value)
+      if (!prefixEmitted && (value as { type: string }).type === 'text-start') {
+        prefixEmitted = true
+        const id = (value as { type: 'text-start'; id: string }).id
+        controller.enqueue({ type: 'text-delta', id, delta: prefixText } as UIMessageChunk)
+      }
+    },
+    cancel() {
+      reader.cancel()
+    },
+  })
+}
+
 export class CustomChatTransport implements ChatTransport<UIMessage> {
   public model: LanguageModel | null = null
   private tools: Record<string, Tool> = {}
@@ -58,6 +90,7 @@ export class CustomChatTransport implements ChatTransport<UIMessage> {
   private systemMessage?: string
   private serviceHub: ServiceHub | null
   private threadId?: string
+  private continueFromContent: string | null = null
 
   constructor(systemMessage?: string, threadId?: string) {
     this.systemMessage = systemMessage
@@ -213,6 +246,14 @@ export class CustomChatTransport implements ChatTransport<UIMessage> {
     return this.tools
   }
 
+  /**
+   * Set partial assistant content to send as a prefill on the next request,
+   * so the model continues generation from where it left off.
+   */
+  setContinueFromContent(content: string) {
+    this.continueFromContent = content
+  }
+
   async sendMessages(
     options: {
       chatId: string
@@ -258,10 +299,18 @@ export class CustomChatTransport implements ChatTransport<UIMessage> {
     }
 
     // Convert UI messages to model messages
-    const modelMessages = convertToModelMessages(
+    const baseMessages = convertToModelMessages(
       this.mapUserInlineAttachments(options.messages)
     )
 
+    // If continuing a truncated response, append the partial assistant content as a
+    // prefill so the model resumes from where it left off rather than regenerating.
+    const continueContent = this.continueFromContent
+    this.continueFromContent = null
+    const modelMessages = continueContent
+      ? [...baseMessages, { role: 'assistant' as const, content: continueContent }]
+      : baseMessages
+
     // Include tools only if we have tools loaded AND model supports them
     const hasTools = Object.keys(this.tools).length > 0
     const selectedModel = useModelProvider.getState().selectedModel
@@ -282,7 +331,7 @@ export class CustomChatTransport implements ChatTransport<UIMessage> {
 
     let tokensPerSecond = 0
 
-    return result.toUIMessageStream({
+    const uiStream = result.toUIMessageStream({
       messageMetadata: ({ part }) => {
         // Track stream start time on start
         if (part.type === 'start' && !streamStartTime) {
@@ -320,6 +369,7 @@ export class CustomChatTransport implements ChatTransport<UIMessage> {
           }
 
           return {
+            finishReason: finishPart.finishReason,
             usage: {
               inputTokens: inputTokens,
               outputTokens: outputTokens,
@@ -364,6 +414,13 @@ export class CustomChatTransport implements ChatTransport<UIMessage> {
         }
       },
     })
+
+    // When continuing a truncated response, inject the partial content as the
+    // very first text-delta so the new message immediately shows it and the
+    // user sees a seamless continuation rather than an empty box.
+    return continueContent
+      ? prependTextDeltaToUIStream(uiStream, continueContent)
+      : uiStream
   }
 
   async reconnectToStream(