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:
@@ -12,6 +12,7 @@ import {
|
|||||||
import { Server, Socket } from 'socket.io';
|
import { Server, Socket } from 'socket.io';
|
||||||
import type { AgentSessionEvent } from '@mariozechner/pi-coding-agent';
|
import type { AgentSessionEvent } from '@mariozechner/pi-coding-agent';
|
||||||
import type { Auth } from '@mosaic/auth';
|
import type { Auth } from '@mosaic/auth';
|
||||||
|
import type { SetThinkingPayload } from '@mosaic/types';
|
||||||
import { AgentService } from '../agent/agent.service.js';
|
import { AgentService } from '../agent/agent.service.js';
|
||||||
import { AUTH } from '../auth/auth.tokens.js';
|
import { AUTH } from '../auth/auth.tokens.js';
|
||||||
import { v4 as uuid } from 'uuid';
|
import { v4 as uuid } from 'uuid';
|
||||||
@@ -112,6 +113,21 @@ export class ChatGateway implements OnGatewayInit, OnGatewayConnection, OnGatewa
|
|||||||
// Track channel connection
|
// Track channel connection
|
||||||
this.agentService.addChannel(conversationId, `websocket:${client.id}`);
|
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
|
// Send acknowledgment
|
||||||
client.emit('message:ack', { conversationId, messageId: uuid() });
|
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 {
|
private relayEvent(client: Socket, conversationId: string, event: AgentSessionEvent): void {
|
||||||
if (!client.connected) {
|
if (!client.connected) {
|
||||||
this.logger.warn(
|
this.logger.warn(
|
||||||
@@ -143,9 +196,31 @@ export class ChatGateway implements OnGatewayInit, OnGatewayConnection, OnGatewa
|
|||||||
client.emit('agent:start', { conversationId });
|
client.emit('agent:start', { conversationId });
|
||||||
break;
|
break;
|
||||||
|
|
||||||
case 'agent_end':
|
case 'agent_end': {
|
||||||
client.emit('agent:end', { conversationId });
|
// 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;
|
break;
|
||||||
|
}
|
||||||
|
|
||||||
case 'message_update': {
|
case 'message_update': {
|
||||||
const assistantEvent = event.assistantMessageEvent;
|
const assistantEvent = event.assistantMessageEvent;
|
||||||
|
|||||||
@@ -23,10 +23,22 @@ export function TuiApp({ gatewayUrl, conversationId, sessionCookie }: TuiAppProp
|
|||||||
initialConversationId: conversationId,
|
initialConversationId: conversationId,
|
||||||
});
|
});
|
||||||
|
|
||||||
useInput((_ch, key) => {
|
useInput((ch, key) => {
|
||||||
if (key.ctrl && _ch === 'c') {
|
if (key.ctrl && ch === 'c') {
|
||||||
exit();
|
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 (
|
return (
|
||||||
@@ -54,6 +66,7 @@ export function TuiApp({ gatewayUrl, conversationId, sessionCookie }: TuiAppProp
|
|||||||
connecting={socket.connecting}
|
connecting={socket.connecting}
|
||||||
modelName={socket.modelName}
|
modelName={socket.modelName}
|
||||||
providerName={socket.providerName}
|
providerName={socket.providerName}
|
||||||
|
thinkingLevel={socket.thinkingLevel}
|
||||||
/>
|
/>
|
||||||
</Box>
|
</Box>
|
||||||
);
|
);
|
||||||
|
|||||||
@@ -10,6 +10,7 @@ export interface BottomBarProps {
|
|||||||
connecting: boolean;
|
connecting: boolean;
|
||||||
modelName: string | null;
|
modelName: string | null;
|
||||||
providerName: string | null;
|
providerName: string | null;
|
||||||
|
thinkingLevel: string;
|
||||||
}
|
}
|
||||||
|
|
||||||
function formatTokens(n: number): string {
|
function formatTokens(n: number): string {
|
||||||
@@ -34,6 +35,7 @@ export function BottomBar({
|
|||||||
connecting,
|
connecting,
|
||||||
modelName,
|
modelName,
|
||||||
providerName,
|
providerName,
|
||||||
|
thinkingLevel,
|
||||||
}: BottomBarProps) {
|
}: BottomBarProps) {
|
||||||
const gatewayStatus = connected ? 'Connected' : connecting ? 'Connecting…' : 'Disconnected';
|
const gatewayStatus = connected ? 'Connected' : connecting ? 'Connecting…' : 'Disconnected';
|
||||||
const gatewayColor = connected ? 'green' : connecting ? 'yellow' : 'red';
|
const gatewayColor = connected ? 'green' : connecting ? 'yellow' : 'red';
|
||||||
@@ -97,6 +99,7 @@ export function BottomBar({
|
|||||||
<Text dimColor>
|
<Text dimColor>
|
||||||
{providerName ? `(${providerName}) ` : ''}
|
{providerName ? `(${providerName}) ` : ''}
|
||||||
{modelName ?? 'awaiting model'}
|
{modelName ?? 'awaiting model'}
|
||||||
|
{thinkingLevel !== 'off' ? ` • ${thinkingLevel}` : ''}
|
||||||
</Text>
|
</Text>
|
||||||
</Box>
|
</Box>
|
||||||
</Box>
|
</Box>
|
||||||
|
|||||||
@@ -4,10 +4,12 @@ import type {
|
|||||||
ServerToClientEvents,
|
ServerToClientEvents,
|
||||||
ClientToServerEvents,
|
ClientToServerEvents,
|
||||||
MessageAckPayload,
|
MessageAckPayload,
|
||||||
|
AgentEndPayload,
|
||||||
AgentTextPayload,
|
AgentTextPayload,
|
||||||
AgentThinkingPayload,
|
AgentThinkingPayload,
|
||||||
ToolStartPayload,
|
ToolStartPayload,
|
||||||
ToolEndPayload,
|
ToolEndPayload,
|
||||||
|
SessionInfoPayload,
|
||||||
ErrorPayload,
|
ErrorPayload,
|
||||||
} from '@mosaic/types';
|
} from '@mosaic/types';
|
||||||
|
|
||||||
@@ -53,12 +55,26 @@ export interface UseSocketReturn {
|
|||||||
tokenUsage: TokenUsage;
|
tokenUsage: TokenUsage;
|
||||||
modelName: string | null;
|
modelName: string | null;
|
||||||
providerName: string | null;
|
providerName: string | null;
|
||||||
|
thinkingLevel: string;
|
||||||
|
availableThinkingLevels: string[];
|
||||||
sendMessage: (content: string) => void;
|
sendMessage: (content: string) => void;
|
||||||
|
setThinkingLevel: (level: string) => void;
|
||||||
connectionError: string | null;
|
connectionError: string | null;
|
||||||
}
|
}
|
||||||
|
|
||||||
type TypedSocket = Socket<ServerToClientEvents, ClientToServerEvents>;
|
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 {
|
export function useSocket(opts: UseSocketOptions): UseSocketReturn {
|
||||||
const { gatewayUrl, sessionCookie, initialConversationId } = opts;
|
const { gatewayUrl, sessionCookie, initialConversationId } = opts;
|
||||||
|
|
||||||
@@ -70,22 +86,16 @@ export function useSocket(opts: UseSocketOptions): UseSocketReturn {
|
|||||||
const [currentStreamText, setCurrentStreamText] = useState('');
|
const [currentStreamText, setCurrentStreamText] = useState('');
|
||||||
const [currentThinkingText, setCurrentThinkingText] = useState('');
|
const [currentThinkingText, setCurrentThinkingText] = useState('');
|
||||||
const [activeToolCalls, setActiveToolCalls] = useState<ToolCall[]>([]);
|
const [activeToolCalls, setActiveToolCalls] = useState<ToolCall[]>([]);
|
||||||
// TODO: wire up once gateway emits token-usage and model-info events
|
const [tokenUsage, setTokenUsage] = useState<TokenUsage>(EMPTY_USAGE);
|
||||||
const tokenUsage: TokenUsage = {
|
const [modelName, setModelName] = useState<string | null>(null);
|
||||||
input: 0,
|
const [providerName, setProviderName] = useState<string | null>(null);
|
||||||
output: 0,
|
const [thinkingLevel, setThinkingLevelState] = useState<string>('off');
|
||||||
total: 0,
|
const [availableThinkingLevels, setAvailableThinkingLevels] = useState<string[]>([]);
|
||||||
cacheRead: 0,
|
|
||||||
cacheWrite: 0,
|
|
||||||
cost: 0,
|
|
||||||
contextPercent: 0,
|
|
||||||
contextWindow: 0,
|
|
||||||
};
|
|
||||||
const modelName: string | null = null;
|
|
||||||
const providerName: string | null = null;
|
|
||||||
const [connectionError, setConnectionError] = useState<string | null>(null);
|
const [connectionError, setConnectionError] = useState<string | null>(null);
|
||||||
|
|
||||||
const socketRef = useRef<TypedSocket | null>(null);
|
const socketRef = useRef<TypedSocket | null>(null);
|
||||||
|
const conversationIdRef = useRef(conversationId);
|
||||||
|
conversationIdRef.current = conversationId;
|
||||||
|
|
||||||
useEffect(() => {
|
useEffect(() => {
|
||||||
const socket = io(`${gatewayUrl}/chat`, {
|
const socket = io(`${gatewayUrl}/chat`, {
|
||||||
@@ -121,6 +131,13 @@ export function useSocket(opts: UseSocketOptions): UseSocketReturn {
|
|||||||
setConversationId(data.conversationId);
|
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', () => {
|
socket.on('agent:start', () => {
|
||||||
setIsStreaming(true);
|
setIsStreaming(true);
|
||||||
setCurrentStreamText('');
|
setCurrentStreamText('');
|
||||||
@@ -153,7 +170,7 @@ export function useSocket(opts: UseSocketOptions): UseSocketReturn {
|
|||||||
);
|
);
|
||||||
});
|
});
|
||||||
|
|
||||||
socket.on('agent:end', () => {
|
socket.on('agent:end', (data: AgentEndPayload) => {
|
||||||
setCurrentStreamText((prev) => {
|
setCurrentStreamText((prev) => {
|
||||||
if (prev) {
|
if (prev) {
|
||||||
setMessages((msgs) => [
|
setMessages((msgs) => [
|
||||||
@@ -166,6 +183,23 @@ export function useSocket(opts: UseSocketOptions): UseSocketReturn {
|
|||||||
setCurrentThinkingText('');
|
setCurrentThinkingText('');
|
||||||
setActiveToolCalls([]);
|
setActiveToolCalls([]);
|
||||||
setIsStreaming(false);
|
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) => {
|
socket.on('error', (data: ErrorPayload) => {
|
||||||
@@ -196,6 +230,15 @@ export function useSocket(opts: UseSocketOptions): UseSocketReturn {
|
|||||||
[conversationId, isStreaming],
|
[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 {
|
return {
|
||||||
connected,
|
connected,
|
||||||
connecting,
|
connecting,
|
||||||
@@ -208,7 +251,10 @@ export function useSocket(opts: UseSocketOptions): UseSocketReturn {
|
|||||||
tokenUsage,
|
tokenUsage,
|
||||||
modelName,
|
modelName,
|
||||||
providerName,
|
providerName,
|
||||||
|
thinkingLevel,
|
||||||
|
availableThinkingLevels,
|
||||||
sendMessage,
|
sendMessage,
|
||||||
|
setThinkingLevel,
|
||||||
connectionError,
|
connectionError,
|
||||||
};
|
};
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -9,6 +9,26 @@ export interface AgentStartPayload {
|
|||||||
|
|
||||||
export interface AgentEndPayload {
|
export interface AgentEndPayload {
|
||||||
conversationId: string;
|
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 {
|
export interface AgentTextPayload {
|
||||||
@@ -44,6 +64,21 @@ export interface ChatMessagePayload {
|
|||||||
content: string;
|
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 */
|
/** Socket.IO typed event map: server → client */
|
||||||
export interface ServerToClientEvents {
|
export interface ServerToClientEvents {
|
||||||
'message:ack': (payload: MessageAckPayload) => void;
|
'message:ack': (payload: MessageAckPayload) => void;
|
||||||
@@ -53,10 +88,12 @@ export interface ServerToClientEvents {
|
|||||||
'agent:thinking': (payload: AgentThinkingPayload) => void;
|
'agent:thinking': (payload: AgentThinkingPayload) => void;
|
||||||
'agent:tool:start': (payload: ToolStartPayload) => void;
|
'agent:tool:start': (payload: ToolStartPayload) => void;
|
||||||
'agent:tool:end': (payload: ToolEndPayload) => void;
|
'agent:tool:end': (payload: ToolEndPayload) => void;
|
||||||
|
'session:info': (payload: SessionInfoPayload) => void;
|
||||||
error: (payload: ErrorPayload) => void;
|
error: (payload: ErrorPayload) => void;
|
||||||
}
|
}
|
||||||
|
|
||||||
/** Socket.IO typed event map: client → server */
|
/** Socket.IO typed event map: client → server */
|
||||||
export interface ClientToServerEvents {
|
export interface ClientToServerEvents {
|
||||||
message: (data: ChatMessagePayload) => void;
|
message: (data: ChatMessagePayload) => void;
|
||||||
|
'set:thinking': (data: SetThinkingPayload) => void;
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -7,6 +7,9 @@ export type {
|
|||||||
AgentThinkingPayload,
|
AgentThinkingPayload,
|
||||||
ToolStartPayload,
|
ToolStartPayload,
|
||||||
ToolEndPayload,
|
ToolEndPayload,
|
||||||
|
SessionUsagePayload,
|
||||||
|
SessionInfoPayload,
|
||||||
|
SetThinkingPayload,
|
||||||
ErrorPayload,
|
ErrorPayload,
|
||||||
ChatMessagePayload,
|
ChatMessagePayload,
|
||||||
ServerToClientEvents,
|
ServerToClientEvents,
|
||||||
|
|||||||
Reference in New Issue
Block a user