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 867e6097b0
commit a061a64fff
6 changed files with 195 additions and 18 deletions

View File

@@ -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;