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,