Add context info in ui

sia/web/api.py

@@ -58,6 +58,8 @@ class Api:
         self._app.router.add_post("/api/memory/entry/{id}/update", self._update_entry)
         self._app.router.add_post("/api/memory/load_iteration", self._load_iteration)
 
+        self._app.router.add_get("/api/context_info", self._get_context_info)
+
     async def _get_llms(self, request: web.Request) -> web.Response:
         return web.Response(
             text=json.dumps(
@@ -293,4 +295,11 @@ class Api:
                 
             return web.Response(status=200)
         except Exception as e:
-            return web.Response(status=400, text=str(e))
+            return web.Response(status=400, text=str(e))
+
+    async def _get_context_info(self, request: web.Request) -> web.Response:
+        """Get current context info."""
+        return web.Response(
+            text=json.dumps(self._agent.context_info),
+            content_type="application/json"
+        )

sia/web/context_info_websocket.py

@@ -0,0 +1,48 @@
+from aiohttp import web, WSMsgType
+from typing import Dict, Set
+
+from .util import wrap_async
+from ..web_agent import WebAgent
+
+class ContextInfoWebSocket:
+    """
+    WebSocket handler for context info updates.
+    Broadcasts context info to all connected clients when it changes.
+    """
+
+    def __init__(self, agent: WebAgent):
+        self._agent = agent
+        self._clients: Set[web.WebSocketResponse] = set()
+        self._agent.add_context_change_handler(wrap_async(self._handle_context_change))
+
+    async def _broadcast_message(self, message: Dict):
+        """Broadcast message to all connected clients."""
+        disconnected = set()
+        for ws in self._clients:
+            try:
+                await ws.send_json(message)
+            except ConnectionResetError:
+                disconnected.add(ws)
+        self._clients -= disconnected
+
+    async def _handle_context_change(self, context_info: Dict):
+        """Handle context changes from the agent."""
+        await self._broadcast_message(context_info)
+
+    async def handle_connection(self, request: web.Request) -> web.WebSocketResponse:
+        """Handle new WebSocket connections."""
+        ws = web.WebSocketResponse(heartbeat=30)
+        await ws.prepare(request)
+
+        try:
+            # Send initial context info
+            await ws.send_json(self._agent.context_info)
+            self._clients.add(ws)
+
+            async for msg in ws:
+                if msg.type == WSMsgType.ERROR:
+                    print(f"WebSocket connection closed with error: {ws.exception()}")
+        finally:
+            self._clients.discard(ws)
+
+        return ws

sia/web/websockets.py

@@ -6,6 +6,7 @@ from ..web_agent import WebAgent
 from ..working_memory import WorkingMemory
 from .auto_approver_websocket import AutoApproverWebSocket
 from .chat_websocket import ChatWebSocket
+from .context_info_websocket import ContextInfoWebSocket
 from .memory_websocket import MemoryWebSocket
 from .response_websocket import ResponseWebSocket
 from .state_websocket import StateWebSocket
@@ -14,12 +15,14 @@ class Websockets:
     def __init__(self, app: web.Application, agent: WebAgent, io_buffer: ChatIOBuffer, auto_approver: AutoApprover, working_memory: WorkingMemory):
         self._auto_approver_ws = AutoApproverWebSocket(auto_approver)
         self._chat_ws = ChatWebSocket(io_buffer)
+        self._context_info_ws = ContextInfoWebSocket(agent)
         self._memory_ws = MemoryWebSocket(working_memory)
         self._response_ws = ResponseWebSocket(agent)
         self._state_ws = StateWebSocket(agent)
 
         app.router.add_get("/ws/auto_approver", self._auto_approver_ws.handle_connection)
         app.router.add_get("/ws/chat", self._chat_ws.handle_connection)
+        app.router.add_get("/ws/context_info", self._context_info_ws.handle_connection)
         app.router.add_get("/ws/memory", self._memory_ws.handle_connection)
         app.router.add_get("/ws/response", self._response_ws.handle_connection)
         app.router.add_get("/ws/state", self._state_ws.handle_connection)

sia/web_agent.py

@@ -45,6 +45,7 @@ class WebAgent(BaseAgent):
         
         self._state_change_handlers: List[Callable[[AgentState], None]] = []
         self._selected_llm_change_handlers: List[Callable[[str], None]] = []
+        self._context_change_handlers: List[Callable[[Dict], None]] = []
 
         self._update_compiled_context()
         self._working_memory.add_change_handler(self._update_compiled_context)
@@ -81,6 +82,13 @@ class WebAgent(BaseAgent):
     def add_state_change_handler(self, handler: Callable[[AgentState], None]) -> None:
         self._state_change_handlers.append(handler)
 
+    @property
+    def context_info(self) -> Dict:
+        return dict(self._compiled_context.attrib)
+
+    def add_context_change_handler(self, handler: Callable[[Dict], None]) -> None:
+        self._context_change_handlers.append(handler)
+
     def run_inference(self) -> None:
         """Start inference on specified LLM"""
         if self._state != AgentState.IDLE:
@@ -132,4 +140,15 @@ class WebAgent(BaseAgent):
             handler(state)
     
     def _update_compiled_context(self):
-        self._compiled_context = self._compile_context(self._llms[self.active_llm])
+        self._compiled_context = self._compile_context(self._llms[self.active_llm])
+        self._notify_context_change()
+
+    def _notify_context_change(self):
+        """Notify all context change handlers."""
+        for handler in self._context_change_handlers:
+            handler(self.context_info)
+
+    def _notify_selected_llm_change(self):
+        """Notify all selected LLM change handlers."""
+        for handler in self._selected_llm_change_handlers:
+            handler(self._selected_llm)

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 (