feat(gateway,cli,types): wire token usage, model info, and thinking levels

Gateway:
- Emit session:info on session creation with provider, model, thinking level
- Include SessionUsagePayload in agent:end with token stats, cost, context usage
- Handle set:thinking client event to cycle thinking levels
- Respond with updated session:info after thinking level change

Types (@mosaic/types):
- Add SessionUsagePayload (tokens, cost, context) to AgentEndPayload
- Add SessionInfoPayload (provider, model, thinking level, available levels)
- Add SetThinkingPayload and set:thinking to ClientToServerEvents
- Add session:info to ServerToClientEvents

CLI TUI:
- useSocket now tracks tokenUsage, modelName, providerName, thinkingLevel
- Updates from both session:info and agent:end usage payload
- Ctrl+T cycles thinking level via set:thinking socket event
- Footer shows thinking level next to model (e.g. 'claude-opus-4-6 • medium')
- Token stats populate with real ↑in ↓out Rcache Wcache $cost ctx%
This commit is contained in:
2026-03-15 14:05:33 -05:00
parent 06887e59c7
commit f0b31d9983
5 changed files with 180 additions and 16 deletions

View File

@@ -10,6 +10,7 @@ export interface BottomBarProps {
connecting: boolean;
modelName: string | null;
providerName: string | null;
thinkingLevel: string;
}
function formatTokens(n: number): string {
@@ -34,6 +35,7 @@ export function BottomBar({
connecting,
modelName,
providerName,
thinkingLevel,
}: BottomBarProps) {
const gatewayStatus = connected ? 'Connected' : connecting ? 'Connecting…' : 'Disconnected';
const gatewayColor = connected ? 'green' : connecting ? 'yellow' : 'red';
@@ -97,6 +99,7 @@ export function BottomBar({
<Text dimColor>
{providerName ? `(${providerName}) ` : ''}
{modelName ?? 'awaiting model'}
{thinkingLevel !== 'off' ? `${thinkingLevel}` : ''}
</Text>
</Box>
</Box>

View File

@@ -4,10 +4,12 @@ import type {
ServerToClientEvents,
ClientToServerEvents,
MessageAckPayload,
AgentEndPayload,
AgentTextPayload,
AgentThinkingPayload,
ToolStartPayload,
ToolEndPayload,
SessionInfoPayload,
ErrorPayload,
} from '@mosaic/types';
@@ -53,12 +55,26 @@ export interface UseSocketReturn {
tokenUsage: TokenUsage;
modelName: string | null;
providerName: string | null;
thinkingLevel: string;
availableThinkingLevels: string[];
sendMessage: (content: string) => void;
setThinkingLevel: (level: string) => void;
connectionError: string | null;
}
type TypedSocket = Socket<ServerToClientEvents, ClientToServerEvents>;
const EMPTY_USAGE: TokenUsage = {
input: 0,
output: 0,
total: 0,
cacheRead: 0,
cacheWrite: 0,
cost: 0,
contextPercent: 0,
contextWindow: 0,
};
export function useSocket(opts: UseSocketOptions): UseSocketReturn {
const { gatewayUrl, sessionCookie, initialConversationId } = opts;
@@ -70,22 +86,16 @@ export function useSocket(opts: UseSocketOptions): UseSocketReturn {
const [currentStreamText, setCurrentStreamText] = useState('');
const [currentThinkingText, setCurrentThinkingText] = useState('');
const [activeToolCalls, setActiveToolCalls] = useState<ToolCall[]>([]);
// TODO: wire up once gateway emits token-usage and model-info events
const tokenUsage: TokenUsage = {
input: 0,
output: 0,
total: 0,
cacheRead: 0,
cacheWrite: 0,
cost: 0,
contextPercent: 0,
contextWindow: 0,
};
const modelName: string | null = null;
const providerName: string | null = null;
const [tokenUsage, setTokenUsage] = useState<TokenUsage>(EMPTY_USAGE);
const [modelName, setModelName] = useState<string | null>(null);
const [providerName, setProviderName] = useState<string | null>(null);
const [thinkingLevel, setThinkingLevelState] = useState<string>('off');
const [availableThinkingLevels, setAvailableThinkingLevels] = useState<string[]>([]);
const [connectionError, setConnectionError] = useState<string | null>(null);
const socketRef = useRef<TypedSocket | null>(null);
const conversationIdRef = useRef(conversationId);
conversationIdRef.current = conversationId;
useEffect(() => {
const socket = io(`${gatewayUrl}/chat`, {
@@ -121,6 +131,13 @@ export function useSocket(opts: UseSocketOptions): UseSocketReturn {
setConversationId(data.conversationId);
});
socket.on('session:info', (data: SessionInfoPayload) => {
setProviderName(data.provider);
setModelName(data.modelId);
setThinkingLevelState(data.thinkingLevel);
setAvailableThinkingLevels(data.availableThinkingLevels);
});
socket.on('agent:start', () => {
setIsStreaming(true);
setCurrentStreamText('');
@@ -153,7 +170,7 @@ export function useSocket(opts: UseSocketOptions): UseSocketReturn {
);
});
socket.on('agent:end', () => {
socket.on('agent:end', (data: AgentEndPayload) => {
setCurrentStreamText((prev) => {
if (prev) {
setMessages((msgs) => [
@@ -166,6 +183,23 @@ export function useSocket(opts: UseSocketOptions): UseSocketReturn {
setCurrentThinkingText('');
setActiveToolCalls([]);
setIsStreaming(false);
// Update usage from the payload
if (data.usage) {
setProviderName(data.usage.provider);
setModelName(data.usage.modelId);
setThinkingLevelState(data.usage.thinkingLevel);
setTokenUsage({
input: data.usage.tokens.input,
output: data.usage.tokens.output,
total: data.usage.tokens.total,
cacheRead: data.usage.tokens.cacheRead,
cacheWrite: data.usage.tokens.cacheWrite,
cost: data.usage.cost,
contextPercent: data.usage.context.percent ?? 0,
contextWindow: data.usage.context.window,
});
}
});
socket.on('error', (data: ErrorPayload) => {
@@ -196,6 +230,15 @@ export function useSocket(opts: UseSocketOptions): UseSocketReturn {
[conversationId, isStreaming],
);
const setThinkingLevel = useCallback((level: string) => {
const cid = conversationIdRef.current;
if (!socketRef.current?.connected || !cid) return;
socketRef.current.emit('set:thinking', {
conversationId: cid,
level,
});
}, []);
return {
connected,
connecting,
@@ -208,7 +251,10 @@ export function useSocket(opts: UseSocketOptions): UseSocketReturn {
tokenUsage,
modelName,
providerName,
thinkingLevel,
availableThinkingLevels,
sendMessage,
setThinkingLevel,
connectionError,
};
}