diff --git a/apps/gateway/src/chat/chat.gateway.ts b/apps/gateway/src/chat/chat.gateway.ts index 0f802c3..01969eb 100644 --- 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; diff --git a/packages/cli/src/tui/app.tsx b/packages/cli/src/tui/app.tsx index 2061edf..3b169ac 100644 --- 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} /> ); diff --git a/packages/cli/src/tui/components/bottom-bar.tsx b/packages/cli/src/tui/components/bottom-bar.tsx index 1fa3c63..cccb50b 100644 --- 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({ {providerName ? `(${providerName}) ` : ''} {modelName ?? 'awaiting model'} + {thinkingLevel !== 'off' ? ` • ${thinkingLevel}` : ''} diff --git a/packages/cli/src/tui/hooks/use-socket.ts b/packages/cli/src/tui/hooks/use-socket.ts index 31cd210..933a255 100644 --- 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; +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([]); - // 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(EMPTY_USAGE); + const [modelName, setModelName] = useState(null); + const [providerName, setProviderName] = useState(null); + const [thinkingLevel, setThinkingLevelState] = useState('off'); + const [availableThinkingLevels, setAvailableThinkingLevels] = useState([]); const [connectionError, setConnectionError] = useState(null); const socketRef = useRef(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, }; } diff --git a/packages/types/src/chat/events.ts b/packages/types/src/chat/events.ts index 1f61358..074c96e 100644 --- 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; } diff --git a/packages/types/src/chat/index.ts b/packages/types/src/chat/index.ts index 551a146..7d039a7 100644 --- 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,