ctx-mgmt: ctx session management (dev mode only) (#2415)

This commit is contained in:
Lily Delalande
2025-05-02 16:12:56 -04:00
committed by GitHub
parent 2366a3ad01
commit 8ba40bdccc
19 changed files with 710 additions and 118 deletions
@@ -5,21 +5,30 @@ import { useChatContextManager } from './ContextManager';
interface ContextLengthExceededHandlerProps {
messages: Message[];
messageId: string;
chatId: string;
workingDir: string;
}
export const ContextLengthExceededHandler: React.FC<ContextLengthExceededHandlerProps> = ({
messages,
messageId,
chatId,
workingDir,
}) => {
const { fetchSummary, summaryContent, isLoadingSummary, errorLoadingSummary, openSummaryModal } =
useChatContextManager();
const {
summaryContent,
isLoadingSummary,
errorLoadingSummary,
openSummaryModal,
handleContextLengthExceeded,
} = useChatContextManager();
const [hasFetchStarted, setHasFetchStarted] = useState(false);
// Find the relevant message to check if it's the latest
const isCurrentMessageLatest =
messageId === messages[messages.length - 1].id ||
messageId === messages[messages.length - 1].created.toString();
messageId === messages[messages.length - 1]?.id ||
messageId === String(messages[messages.length - 1]?.created);
// Only allow interaction for the most recent context length exceeded event
const shouldAllowSummaryInteraction = isCurrentMessageLatest;
@@ -27,31 +36,60 @@ export const ContextLengthExceededHandler: React.FC<ContextLengthExceededHandler
// Use a ref to track if we've started the fetch
const fetchStartedRef = useRef(false);
// Function to trigger the async operation properly
const triggerContextLengthExceeded = () => {
setHasFetchStarted(true);
fetchStartedRef.current = true;
// Call the async function without awaiting it in useEffect
handleContextLengthExceeded(messages, chatId, workingDir).catch((err) => {
console.error('Error handling context length exceeded:', err);
});
};
useEffect(() => {
// Automatically fetch summary if conditions are met
if (
!summaryContent &&
!hasFetchStarted &&
shouldAllowSummaryInteraction &&
!fetchStartedRef.current
) {
setHasFetchStarted(true);
fetchStartedRef.current = true;
fetchSummary(messages);
// Use the wrapper function instead of calling the async function directly
triggerContextLengthExceeded();
}
}, [fetchSummary, hasFetchStarted, messages, shouldAllowSummaryInteraction, summaryContent]);
// eslint-disable-next-line react-hooks/exhaustive-deps
}, [
hasFetchStarted,
messages,
shouldAllowSummaryInteraction,
summaryContent,
chatId,
workingDir,
]);
// Handle retry
// Handle retry - Call the async function properly
const handleRetry = () => {
if (!shouldAllowSummaryInteraction) return;
fetchSummary(messages);
// Reset states for retry
setHasFetchStarted(false);
fetchStartedRef.current = false;
// Trigger the process again
triggerContextLengthExceeded();
};
// Render the notification UI
return (
<div className="flex flex-col items-start mt-1 pl-4">
{/* Horizontal line with text in the middle - shown regardless of loading state */}
<div className="relative flex items-center py-2 w-full">
<div className="flex-grow border-t border-gray-300"></div>
<div className="flex-grow border-t border-gray-300"></div>
</div>
{isLoadingSummary && shouldAllowSummaryInteraction ? (
// Only show loading indicator during loading state
// Show loading indicator during loading state
<div className="flex items-center text-xs text-gray-400">
<span className="mr-2">Preparing summary...</span>
<span className="animate-spin h-3 w-3 border-2 border-gray-400 rounded-full border-t-transparent"></span>
@@ -59,15 +97,17 @@ export const ContextLengthExceededHandler: React.FC<ContextLengthExceededHandler
) : (
// Show different UI based on whether it's already handled
<>
<span className="text-xs text-gray-400 italic">{'Session summarized'}</span>
<span className="text-xs text-gray-400">{`Your conversation has exceeded the model's context capacity`}</span>
<span className="text-xs text-gray-400">{`Messages above this line remain viewable but are not included in the active context`}</span>
{/* Only show the button if its last message */}
{shouldAllowSummaryInteraction && (
<button
onClick={() => (errorLoadingSummary ? handleRetry() : openSummaryModal())}
className="text-xs text-textStandard hover:text-textSubtle transition-colors mt-1 flex items-center"
>
{errorLoadingSummary ? 'Retry loading summary' : 'View or edit summary'}
{errorLoadingSummary
? 'Retry loading summary'
: 'View or edit summary (you may continue your conversation based on the summary)'}
</button>
)}
</>
@@ -3,7 +3,7 @@ import { Message } from '../../types/message';
import { manageContextFromBackend, convertApiMessageToFrontendMessage } from './index';
// Define the context management interface
interface ContextManagerState {
interface ChatContextManagerState {
summaryContent: string;
summarizedThread: Message[];
isSummaryModalOpen: boolean;
@@ -11,31 +11,73 @@ interface ContextManagerState {
errorLoadingSummary: boolean;
}
interface ContextManagerActions {
interface ChatContextManagerActions {
fetchSummary: (messages: Message[]) => Promise<void>;
updateSummary: (newSummaryContent: string) => void;
resetMessagesWithSummary: (
messages: Message[],
setMessages: (messages: Message[]) => void
setMessages: (messages: Message[]) => void,
ancestorMessages: Message[],
setAncestorMessages: (messages: Message[]) => void,
summaryContent: string
) => void;
openSummaryModal: () => void;
closeSummaryModal: () => void;
hasContextLengthExceededContent: (message: Message) => boolean;
handleContextLengthExceeded: (
messages: Message[],
chatId: string,
workingDir: string
) => Promise<void>;
}
// Create the context
const ContextManagerContext = createContext<
(ContextManagerState & ContextManagerActions) | undefined
const ChatContextManagerContext = createContext<
(ChatContextManagerState & ChatContextManagerActions) | undefined
>(undefined);
// Create the provider component
export const ContextManagerProvider: React.FC<{ children: React.ReactNode }> = ({ children }) => {
export const ChatContextManagerProvider: React.FC<{ children: React.ReactNode }> = ({
children,
}) => {
const [summaryContent, setSummaryContent] = useState<string>('');
const [summarizedThread, setSummarizedThread] = useState<Message[]>([]);
const [isSummaryModalOpen, setIsSummaryModalOpen] = useState<boolean>(false);
const [isLoadingSummary, setIsLoadingSummary] = useState<boolean>(false);
const [errorLoadingSummary, setErrorLoadingSummary] = useState<boolean>(false);
const handleContextLengthExceeded = async (messages: Message[]): Promise<void> => {
setIsLoadingSummary(true);
setErrorLoadingSummary(false);
try {
// 2. Now get the summary from the backend
const summaryResponse = await manageContextFromBackend({
messages: messages,
manageAction: 'summarize',
});
// Convert API messages to frontend messages
const convertedMessages = summaryResponse.messages.map(
(apiMessage) => convertApiMessageToFrontendMessage(apiMessage, false, true) // do not show to user but send to llm
);
// Extract summary from the first message
const summaryMessage = convertedMessages[0].content[0];
if (summaryMessage.type === 'text') {
const summary = summaryMessage.text;
setSummaryContent(summary);
setSummarizedThread(convertedMessages);
}
setIsLoadingSummary(false);
} catch (err) {
console.error('Error handling context length exceeded:', err);
setErrorLoadingSummary(true);
setIsLoadingSummary(false);
}
};
const fetchSummary = async (messages: Message[]) => {
setIsLoadingSummary(true);
setErrorLoadingSummary(false);
@@ -47,8 +89,8 @@ export const ContextManagerProvider: React.FC<{ children: React.ReactNode }> = (
});
// Convert API messages to frontend messages
const convertedMessages = response.messages.map((apiMessage) =>
convertApiMessageToFrontendMessage(apiMessage)
const convertedMessages = response.messages.map(
(apiMessage) => convertApiMessageToFrontendMessage(apiMessage, false, true) // do not show to user but send to llm
);
// Extract the summary text from the first message
@@ -101,27 +143,69 @@ export const ContextManagerProvider: React.FC<{ children: React.ReactNode }> = (
const resetMessagesWithSummary = (
messages: Message[],
setMessages: (messages: Message[]) => void
setMessages: (messages: Message[]) => void,
ancestorMessages: Message[],
setAncestorMessages: (messages: Message[]) => void,
summaryContent: string
) => {
// Update summarizedThread with metadata
const updatedSummarizedThread = summarizedThread.map((msg) => ({
...msg,
display: false,
sendToLLM: true,
}));
// Create a copy of the summarized thread
const updatedSummarizedThread = [...summarizedThread];
// Update list of messages with other metadata
const updatedMessages = messages.map((msg) => ({
...msg,
display: true,
sendToLLM: false,
}));
// Make sure there's at least one message in the summarized thread
if (updatedSummarizedThread.length > 0) {
// Get the first message
const firstMessage = { ...updatedSummarizedThread[0] };
// Make a copy that combines both
const newMessages = [...updatedMessages, ...updatedSummarizedThread];
// Make a copy of the content array
const contentCopy = [...firstMessage.content];
// Assuming the first content item is of type TextContent
if (contentCopy.length > 0 && 'text' in contentCopy[0]) {
// Update the text with the new summary content
contentCopy[0] = {
...contentCopy[0],
text: summaryContent,
};
// Update the first message with the new content
firstMessage.content = contentCopy;
// Update the first message in the thread
updatedSummarizedThread[0] = firstMessage;
}
}
// Update metadata for the summarized thread
const finalUpdatedThread = updatedSummarizedThread.map((msg, index) => ({
...msg,
display: index === 0, // First message has display: true, others false
sendToLLM: true, // All messages have sendToLLM: true
}));
// Update the messages state
setMessages(newMessages);
setMessages(finalUpdatedThread);
// If ancestorMessages already has items, extend it instead of replacing it
if (ancestorMessages.length > 0) {
// Convert current messages to ancestor format
const newAncestorMessages = messages.map((msg) => ({
...msg,
display: true,
sendToLLM: false,
}));
// Append new ancestor messages to existing ones
setAncestorMessages([...ancestorMessages, ...newAncestorMessages]);
} else {
// Initial set of ancestor messages
const newAncestorMessages = messages.map((msg) => ({
...msg,
display: true,
sendToLLM: false,
}));
setAncestorMessages(newAncestorMessages);
}
// Clear the summarized thread and content
setSummarizedThread([]);
@@ -155,14 +239,19 @@ export const ContextManagerProvider: React.FC<{ children: React.ReactNode }> = (
openSummaryModal,
closeSummaryModal,
hasContextLengthExceededContent,
handleContextLengthExceeded,
};
return <ContextManagerContext.Provider value={value}>{children}</ContextManagerContext.Provider>;
return (
<ChatContextManagerContext.Provider value={value}>
{children}
</ChatContextManagerContext.Provider>
);
};
// Create a hook to use the context
export const useChatContextManager = () => {
const context = useContext(ContextManagerContext);
const context = useContext(ChatContextManagerContext);
if (context === undefined) {
throw new Error('useContextManager must be used within a ContextManagerProvider');
}
@@ -93,6 +93,12 @@ export function SessionSummaryModal({
ref={textareaRef}
defaultValue={summaryContent}
className="bg-gray-50 dark:bg-gray-800 p-4 rounded-lg text-gray-700 dark:text-gray-300 border border-gray-200 dark:border-gray-700 text-sm w-full min-h-[200px] focus:outline-none focus:ring-2 focus:ring-blue-500 focus:border-transparent"
style={{
textRendering: 'optimizeLegibility',
WebkitFontSmoothing: 'antialiased',
MozOsxFontSmoothing: 'grayscale',
transform: 'translateZ(0)', // Force hardware acceleration
}}
/>
</div>
);
@@ -49,10 +49,14 @@ export async function manageContextFromBackend({
}
// Function to convert API Message to frontend Message
export function convertApiMessageToFrontendMessage(apiMessage: ApiMessage): FrontendMessage {
export function convertApiMessageToFrontendMessage(
apiMessage: ApiMessage,
display?: boolean,
sendToLLM?: boolean
): FrontendMessage {
return {
display: false,
sendToLLM: false,
display: display ?? true,
sendToLLM: sendToLLM ?? true,
id: generateId(),
role: apiMessage.role,
created: apiMessage.created,