Add context info in ui

web/src/components/Layout/MobileControls.tsx

@@ -5,17 +5,17 @@ const MobileControls: React.FC = () => {
   return (
     <div className="p-4 space-y-4 h-full overflow-y-auto bg-white">
       <h3 className="text-lg font-semibold mb-4">Agent Controls</h3>
-      
+
       <AgentSection className="pb-4 border-b" />
-      
+
       <div className="pt-4">
         <LLMSection className="pb-4 border-b" />
       </div>
-      
+
       <div className="pt-4">
         <AutoApproverSection className="pb-4 border-b" />
       </div>
-      
+
       <div className="pt-4">
         <MemorySection />
       </div>

web/src/components/Layout/Ribbon.tsx

@@ -4,9 +4,9 @@ import { AgentSection, LLMSection, AutoApproverSection, MemorySection } from '..
 
 const Ribbon: React.FC = () => {
   const screenSize = useScreenSize();
-  
+
   if (screenSize === 'mobile') return null;
-  
+
   return (
     <div className="bg-white border-b">
       <div className="flex items-stretch">

web/src/components/Memory/MemoryEditor.tsx

@@ -1,7 +1,8 @@
 import React, { useState } from 'react';
 import { useWebSocketContext } from '../../contexts/WebSocketContext';
 import { api, EntryTypes, getInitialEntryData } from '../../services/api';
-import { 
+import {
+  ContextInfoEntry,
   ReasoningEntry,
   SingleEntry,
   RepeatEntry,
@@ -253,7 +254,9 @@ const MemoryEditor: React.FC<MemoryEditorProps> = ({ className = '' }) => {
             <div className="animate-pulse bg-blue-300 h-full w-24"></div>
           </div>
         )}
-        
+
+        <ContextInfoEntry />
+
         {memory.length === 0 ? (
           <div className="flex items-center justify-center h-full">
             <p className="text-gray-500">No memory entries</p>

web/src/components/Memory/entries/ContextInfoEntry.tsx

@@ -0,0 +1,30 @@
+import React from 'react';
+import { useWebSocketContext } from '../../../contexts/WebSocketContext';
+
+const ContextInfoEntry: React.FC = () => {
+  const { contextInfo } = useWebSocketContext();
+
+  if (!contextInfo) {
+    return null;
+  }
+
+  return (
+    <div className="mb-4 border rounded-lg bg-gray-50">
+      <div className="px-3 py-2 bg-gray-100 border-b rounded-t-lg">
+        <span className="text-sm font-medium">Context Info</span>
+      </div>
+      <div className="p-3">
+        <div className="space-y-1 text-sm font-mono">
+          {Object.entries(contextInfo).map(([key, value]) => (
+            <div key={key} className="flex">
+              <span className="text-gray-600 mr-2">{key}:</span>
+              <span className="text-gray-900">{String(value)}</span>
+            </div>
+          ))}
+        </div>
+      </div>
+    </div>
+  );
+};
+
+export default ContextInfoEntry;

web/src/components/Memory/entries/index.ts

@@ -1,5 +1,6 @@
 export { default as BaseEntry } from './BaseEntry';
 export { default as BackgroundEntry } from './BackgroundEntry';
+export { default as ContextInfoEntry } from './ContextInfoEntry';
 export { default as ParseErrorEntry } from './ParseErrorEntry';
 export { default as ReadEntry } from './ReadEntry';
 export { default as ReasoningEntry } from './ReasoningEntry';

web/src/contexts/WebSocketContext.tsx

@@ -1,6 +1,12 @@
 import React, { createContext, useContext, useEffect, useRef, useState } from 'react';
 import { AutoApproverConfig, ChatMessage } from '../types';
 
+interface ContextInfo {
+  [key: string]: string | number;
+  token_count: number;
+  token_limit: number;
+}
+
 interface WebSocketContextType {
   state: {
     agentState: 'IDLE' | 'INFERENCE' | 'PROCESSING_RESPONSE';
@@ -17,6 +23,7 @@ interface WebSocketContextType {
   llms: {name: string}[];
   setActiveLLM: (llm: string) => Promise<void>;
   connectionStatus: Record<string, string>;
+  contextInfo: ContextInfo | null;
 }
 
 const defaultState: WebSocketContextType = {
@@ -40,6 +47,7 @@ const defaultState: WebSocketContextType = {
   llms: [],
   setActiveLLM: async () => {},
   connectionStatus: {},
+  contextInfo: null,
 };
 
 const WebSocketContext = createContext<WebSocketContextType>(defaultState);
@@ -56,6 +64,7 @@ export const WebSocketProvider: React.FC<{children: React.ReactNode}> = ({ child
   const [autoApprover, setAutoApprover] = useState<AutoApproverConfig>(defaultState.autoApprover);
   const [llms, setLLMs] = useState<{name: string}[]>([]);
   const [connectionStatus, setConnectionStatus] = useState<Record<string, string>>({});
+  const [contextInfo, setContextInfo] = useState<ContextInfo | null>(null);
   
   const sockets = useRef<{[key: string]: WebSocket | null}>({});
   
@@ -101,7 +110,7 @@ export const WebSocketProvider: React.FC<{children: React.ReactNode}> = ({ child
   }, []);
 
   useEffect(() => {
-    const endpoints = ['ws/state', 'ws/response', 'ws/memory', 'ws/auto_approver', 'ws/chat'];
+    const endpoints = ['ws/state', 'ws/response', 'ws/memory', 'ws/auto_approver', 'ws/chat', 'ws/context_info'];
     const newConnectionStatus: Record<string, string> = {};
     
     const createWebSocket = (endpoint: string) => {
@@ -157,20 +166,22 @@ export const WebSocketProvider: React.FC<{children: React.ReactNode}> = ({ child
                 setMessages(prevMessages => {
                   const updatedIds = data.messages.map((m: ChatMessage) => m.id);
                   const filteredMessages = prevMessages.filter(m => !updatedIds.includes(m.id));
-                  
+
                   return [...filteredMessages, ...data.messages];
                 });
               }
-              
+
               if (data.deleted_message_ids) {
-                setMessages(prevMessages => 
+                setMessages(prevMessages =>
                   prevMessages.filter(m => !data.deleted_message_ids.includes(m.id))
                 );
               }
-              
+
               if (data.last_read_timestamp) {
                 setLastReadTimestamp(data.last_read_timestamp);
               }
+            } else if (endpoint === 'ws/context_info') {
+              setContextInfo(data);
             }
           } catch (error) {
             console.error(`Failed to parse message from ${endpoint}:`, error);
@@ -232,7 +243,8 @@ export const WebSocketProvider: React.FC<{children: React.ReactNode}> = ({ child
     isConnected,
     llms,
     setActiveLLM,
-    connectionStatus
+    connectionStatus,
+    contextInfo
   };
   
   return (