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%
2026-03-15 14:05:33 -05:00
parent 867e6097b0
commit a061a64fff
6 changed files with 195 additions and 18 deletions
--- a/apps/gateway/src/chat/chat.gateway.ts
+++ b/apps/gateway/src/chat/chat.gateway.ts
@@ -12,6 +12,7 @@ import {
 import { Server, Socket } from 'socket.io';
 import type { AgentSessionEvent } from '@mariozechner/pi-coding-agent';
 import type { Auth } from '@mosaic/auth';
+import type { SetThinkingPayload } from '@mosaic/types';
 import { AgentService } from '../agent/agent.service.js';
 import { AUTH } from '../auth/auth.tokens.js';
 import { v4 as uuid } from 'uuid';
@@ -112,6 +113,21 @@ export class ChatGateway implements OnGatewayInit, OnGatewayConnection, OnGatewa
    // Track channel connection
    this.agentService.addChannel(conversationId, `websocket:${client.id}`);

+    // Send session info so the client knows the model/provider
+    {
+      const agentSession = this.agentService.getSession(conversationId);
+      if (agentSession) {
+        const piSession = agentSession.piSession;
+        client.emit('session:info', {
+          conversationId,
+          provider: agentSession.provider,
+          modelId: agentSession.modelId,
+          thinkingLevel: piSession.thinkingLevel,
+          availableThinkingLevels: piSession.getAvailableThinkingLevels(),
+        });
+      }
+    }
+
    // Send acknowledgment
    client.emit('message:ack', { conversationId, messageId: uuid() });

@@ -130,6 +146,43 @@ export class ChatGateway implements OnGatewayInit, OnGatewayConnection, OnGatewa
    }
  }

+  @SubscribeMessage('set:thinking')
+  handleSetThinking(
+    @ConnectedSocket() client: Socket,
+    @MessageBody() data: SetThinkingPayload,
+  ): void {
+    const session = this.agentService.getSession(data.conversationId);
+    if (!session) {
+      client.emit('error', {
+        conversationId: data.conversationId,
+        error: 'No active session for this conversation.',
+      });
+      return;
+    }
+
+    const validLevels = session.piSession.getAvailableThinkingLevels();
+    if (!validLevels.includes(data.level as never)) {
+      client.emit('error', {
+        conversationId: data.conversationId,
+        error: `Invalid thinking level "${data.level}". Available: ${validLevels.join(', ')}`,
+      });
+      return;
+    }
+
+    session.piSession.setThinkingLevel(data.level as never);
+    this.logger.log(
+      `Thinking level set to "${data.level}" for conversation ${data.conversationId}`,
+    );
+
+    client.emit('session:info', {
+      conversationId: data.conversationId,
+      provider: session.provider,
+      modelId: session.modelId,
+      thinkingLevel: session.piSession.thinkingLevel,
+      availableThinkingLevels: session.piSession.getAvailableThinkingLevels(),
+    });
+  }
+
  private relayEvent(client: Socket, conversationId: string, event: AgentSessionEvent): void {
    if (!client.connected) {
      this.logger.warn(
@@ -143,9 +196,31 @@ export class ChatGateway implements OnGatewayInit, OnGatewayConnection, OnGatewa
        client.emit('agent:start', { conversationId });
        break;

-      case 'agent_end':
-        client.emit('agent:end', { conversationId });
+      case 'agent_end': {
+        // Gather usage stats from the Pi session
+        const agentSession = this.agentService.getSession(conversationId);
+        const piSession = agentSession?.piSession;
+        const stats = piSession?.getSessionStats();
+        const contextUsage = piSession?.getContextUsage();
+
+        client.emit('agent:end', {
+          conversationId,
+          usage: stats
+            ? {
+                provider: agentSession?.provider ?? 'unknown',
+                modelId: agentSession?.modelId ?? 'unknown',
+                thinkingLevel: piSession?.thinkingLevel ?? 'off',
+                tokens: stats.tokens,
+                cost: stats.cost,
+                context: {
+                  percent: contextUsage?.percent ?? null,
+                  window: contextUsage?.contextWindow ?? 0,
+                },
+              }
+            : undefined,
+        });
        break;
+      }

      case 'message_update': {
        const assistantEvent = event.assistantMessageEvent;
--- a/packages/cli/src/tui/app.tsx
+++ b/packages/cli/src/tui/app.tsx
@@ -23,10 +23,22 @@ export function TuiApp({ gatewayUrl, conversationId, sessionCookie }: TuiAppProp
    initialConversationId: conversationId,
  });

-  useInput((_ch, key) => {
-    if (key.ctrl && _ch === 'c') {
+  useInput((ch, key) => {
+    if (key.ctrl && ch === 'c') {
      exit();
    }
+    // Ctrl+T: cycle thinking level
+    if (key.ctrl && ch === 't') {
+      const levels = socket.availableThinkingLevels;
+      if (levels.length > 0) {
+        const currentIdx = levels.indexOf(socket.thinkingLevel);
+        const nextIdx = (currentIdx + 1) % levels.length;
+        const next = levels[nextIdx];
+        if (next) {
+          socket.setThinkingLevel(next);
+        }
+      }
+    }
  });

  return (
@@ -54,6 +66,7 @@ export function TuiApp({ gatewayUrl, conversationId, sessionCookie }: TuiAppProp
        connecting={socket.connecting}
        modelName={socket.modelName}
        providerName={socket.providerName}
+        thinkingLevel={socket.thinkingLevel}
      />
    </Box>
  );
--- a/packages/cli/src/tui/components/bottom-bar.tsx
+++ b/packages/cli/src/tui/components/bottom-bar.tsx
@@ -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>
--- a/packages/cli/src/tui/hooks/use-socket.ts
+++ b/packages/cli/src/tui/hooks/use-socket.ts
@@ -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,
  };
 }
--- a/packages/types/src/chat/events.ts
+++ b/packages/types/src/chat/events.ts
@@ -9,6 +9,26 @@ export interface AgentStartPayload {

 export interface AgentEndPayload {
  conversationId: string;
+  usage?: SessionUsagePayload;
+}
+
+/** Session metadata emitted with agent:end and on session:info */
+export interface SessionUsagePayload {
+  provider: string;
+  modelId: string;
+  thinkingLevel: string;
+  tokens: {
+    input: number;
+    output: number;
+    cacheRead: number;
+    cacheWrite: number;
+    total: number;
+  };
+  cost: number;
+  context: {
+    percent: number | null;
+    window: number;
+  };
 }

 export interface AgentTextPayload {
@@ -44,6 +64,21 @@ export interface ChatMessagePayload {
  content: string;
 }

+/** Session info pushed when session is created or model changes */
+export interface SessionInfoPayload {
+  conversationId: string;
+  provider: string;
+  modelId: string;
+  thinkingLevel: string;
+  availableThinkingLevels: string[];
+}
+
+/** Client request to change thinking level */
+export interface SetThinkingPayload {
+  conversationId: string;
+  level: string;
+}
+
 /** Socket.IO typed event map: server → client */
 export interface ServerToClientEvents {
  'message:ack': (payload: MessageAckPayload) => void;
@@ -53,10 +88,12 @@ export interface ServerToClientEvents {
  'agent:thinking': (payload: AgentThinkingPayload) => void;
  'agent:tool:start': (payload: ToolStartPayload) => void;
  'agent:tool:end': (payload: ToolEndPayload) => void;
+  'session:info': (payload: SessionInfoPayload) => void;
  error: (payload: ErrorPayload) => void;
 }

 /** Socket.IO typed event map: client → server */
 export interface ClientToServerEvents {
  message: (data: ChatMessagePayload) => void;
+  'set:thinking': (data: SetThinkingPayload) => void;
 }
--- a/packages/types/src/chat/index.ts
+++ b/packages/types/src/chat/index.ts
@@ -7,6 +7,9 @@ export type {
  AgentThinkingPayload,
  ToolStartPayload,
  ToolEndPayload,
+  SessionUsagePayload,
+  SessionInfoPayload,
+  SetThinkingPayload,
  ErrorPayload,
  ChatMessagePayload,
  ServerToClientEvents,