1226 lines
42 KiB
TypeScript
1226 lines
42 KiB
TypeScript
import { createHash } from 'node:crypto';
|
|
import { Inject, Logger, Optional } from '@nestjs/common';
|
|
import {
|
|
WebSocketGateway,
|
|
WebSocketServer,
|
|
SubscribeMessage,
|
|
OnGatewayConnection,
|
|
OnGatewayDisconnect,
|
|
type OnGatewayInit,
|
|
ConnectedSocket,
|
|
MessageBody,
|
|
} from '@nestjs/websockets';
|
|
import { Server, Socket } from 'socket.io';
|
|
import type { AgentSessionEvent } from '@mariozechner/pi-coding-agent';
|
|
import {
|
|
verifyDiscordIngressEnvelope,
|
|
parseDiscordInteractionBindings,
|
|
resolveDiscordInteractionActorId,
|
|
resolveDiscordInteractionBinding,
|
|
type DiscordIngressEnvelope,
|
|
type DiscordIngressPayload,
|
|
} from '@mosaicstack/discord-plugin';
|
|
import type { Auth } from '@mosaicstack/auth';
|
|
import type { Brain } from '@mosaicstack/brain';
|
|
import { redactSensitiveContent } from '@mosaicstack/log';
|
|
import type {
|
|
SetThinkingPayload,
|
|
SlashCommandApprovalResultPayload,
|
|
SlashCommandPayload,
|
|
SystemReloadPayload,
|
|
RoutingDecisionInfo,
|
|
AbortPayload,
|
|
} from '@mosaicstack/types';
|
|
import { AgentService, type ConversationHistoryMessage } from '../agent/agent.service.js';
|
|
import { RuntimeProviderService } from '../agent/runtime-provider-registry.service.js';
|
|
import { AUTH } from '../auth/auth.tokens.js';
|
|
import {
|
|
scopeFromUser,
|
|
type ActorTenantScope,
|
|
type AuthenticatedUserLike,
|
|
} from '../auth/session-scope.js';
|
|
import { BRAIN } from '../brain/brain.tokens.js';
|
|
import { CommandRegistryService } from '../commands/command-registry.service.js';
|
|
import { CommandExecutorService } from '../commands/command-executor.service.js';
|
|
import { CommandAuthorizationService } from '../commands/command-authorization.service.js';
|
|
import { RoutingEngineService } from '../agent/routing/routing-engine.service.js';
|
|
import { v4 as uuid } from 'uuid';
|
|
import { ChatSocketMessageDto } from './chat.dto.js';
|
|
import { validateDiscordServiceToken, validateSocketSession } from './chat.gateway-auth.js';
|
|
import { DiscordReplayProtector } from '../plugin/discord-replay-protector.js';
|
|
|
|
/** Per-client state tracking streaming accumulation for persistence. */
|
|
interface ClientSession {
|
|
conversationId: string;
|
|
cleanup: () => void;
|
|
/** Accumulated assistant response text for the current turn. */
|
|
assistantText: string;
|
|
/** Tool calls observed during the current turn. */
|
|
toolCalls: Array<{ toolCallId: string; toolName: string; args: unknown; isError: boolean }>;
|
|
/** Tool calls in-flight (started but not ended yet). */
|
|
pendingToolCalls: Map<string, { toolName: string; args: unknown }>;
|
|
/** Server-derived owner/tenant scope for this socket's conversation attachment. */
|
|
scope: ActorTenantScope;
|
|
/** Last routing decision made for this session (M4-008) */
|
|
lastRoutingDecision?: RoutingDecisionInfo;
|
|
}
|
|
|
|
/**
|
|
* Per-conversation model overrides set via /model command (M4-007).
|
|
* Keyed by conversationId, value is the model name to use.
|
|
*/
|
|
const modelOverrides = new Map<string, string>();
|
|
const MAX_REDACTION_BUFFER_LENGTH = 8_192;
|
|
|
|
function isDiscordIngressEnvelope(value: unknown): value is DiscordIngressEnvelope {
|
|
if (typeof value !== 'object' || value === null) return false;
|
|
const envelope = value as { payload?: unknown; signature?: unknown };
|
|
if (
|
|
typeof envelope.signature !== 'string' ||
|
|
typeof envelope.payload !== 'object' ||
|
|
envelope.payload === null
|
|
) {
|
|
return false;
|
|
}
|
|
const payload = envelope.payload as Record<string, unknown>;
|
|
return [
|
|
payload['correlationId'],
|
|
payload['messageId'],
|
|
payload['guildId'],
|
|
payload['channelId'],
|
|
payload['userId'],
|
|
payload['conversationId'],
|
|
payload['content'],
|
|
].every((field: unknown): boolean => typeof field === 'string');
|
|
}
|
|
|
|
function isChatSocketMessage(value: unknown): value is ChatSocketMessageDto {
|
|
if (typeof value !== 'object' || value === null) return false;
|
|
const payload = value as { content?: unknown; conversationId?: unknown };
|
|
return (
|
|
typeof payload.content === 'string' &&
|
|
(payload.conversationId === undefined || typeof payload.conversationId === 'string')
|
|
);
|
|
}
|
|
|
|
@WebSocketGateway({
|
|
cors: {
|
|
origin: process.env['GATEWAY_CORS_ORIGIN'] ?? 'http://localhost:3000',
|
|
},
|
|
namespace: '/chat',
|
|
})
|
|
export class ChatGateway implements OnGatewayInit, OnGatewayConnection, OnGatewayDisconnect {
|
|
@WebSocketServer()
|
|
server!: Server;
|
|
|
|
private readonly logger = new Logger(ChatGateway.name);
|
|
private readonly clientSessions = new Map<string, ClientSession>();
|
|
/** Raw stream fragments are kept in memory only until they are safe to redact and emit. */
|
|
private readonly textEgressBuffers = new Map<string, string>();
|
|
private readonly thinkingEgressBuffers = new Map<string, string>();
|
|
private readonly overflowedEgress = new Set<string>();
|
|
private readonly discordReplayProtector = new DiscordReplayProtector();
|
|
|
|
constructor(
|
|
@Inject(AgentService) private readonly agentService: AgentService,
|
|
@Inject(AUTH) private readonly auth: Auth,
|
|
@Inject(BRAIN) private readonly brain: Brain,
|
|
@Inject(CommandRegistryService) private readonly commandRegistry: CommandRegistryService,
|
|
@Inject(CommandExecutorService) private readonly commandExecutor: CommandExecutorService,
|
|
@Inject(RoutingEngineService) private readonly routingEngine: RoutingEngineService,
|
|
@Optional()
|
|
@Inject(CommandAuthorizationService)
|
|
private readonly commandAuthorization: CommandAuthorizationService | null = null,
|
|
@Optional()
|
|
@Inject(RuntimeProviderService)
|
|
private readonly runtimeRegistry: RuntimeProviderService | null = null,
|
|
) {}
|
|
|
|
afterInit(): void {
|
|
this.logger.log('Chat WebSocket gateway initialized');
|
|
}
|
|
|
|
async handleConnection(client: Socket): Promise<void> {
|
|
const serviceToken = client.handshake.auth['discordServiceToken'];
|
|
if (validateDiscordServiceToken(serviceToken, process.env['DISCORD_SERVICE_TOKEN'])) {
|
|
client.data.discordService = true;
|
|
this.logger.log(`Authenticated Discord service connected: ${client.id}`);
|
|
return;
|
|
}
|
|
|
|
const session = await validateSocketSession(client.handshake.headers, this.auth);
|
|
if (!session) {
|
|
this.logger.warn(`Rejected unauthenticated WebSocket client: ${client.id}`);
|
|
client.disconnect();
|
|
return;
|
|
}
|
|
|
|
client.data.user = session.user;
|
|
client.data.session = session.session;
|
|
this.logger.log(`Client connected: ${client.id}`);
|
|
client.emit('commands:manifest', { manifest: this.commandRegistry.getManifest() });
|
|
}
|
|
|
|
handleDisconnect(client: Socket): void {
|
|
this.logger.log(`Client disconnected: ${client.id}`);
|
|
const session = this.clientSessions.get(client.id);
|
|
if (session) {
|
|
session.cleanup();
|
|
this.agentService.removeChannel(
|
|
session.conversationId,
|
|
`websocket:${client.id}`,
|
|
session.scope,
|
|
);
|
|
this.clientSessions.delete(client.id);
|
|
}
|
|
this.textEgressBuffers.delete(client.id);
|
|
this.thinkingEgressBuffers.delete(client.id);
|
|
this.overflowedEgress.delete(this.egressKey(client, 'agent:text'));
|
|
this.overflowedEgress.delete(this.egressKey(client, 'agent:thinking'));
|
|
}
|
|
|
|
private getClientScope(client: Socket): ActorTenantScope | null {
|
|
const user = client.data.user as AuthenticatedUserLike | undefined;
|
|
if (!user?.id) return null;
|
|
return scopeFromUser(user);
|
|
}
|
|
|
|
private modelOverrideKey(conversationId: string, scope: ActorTenantScope): string {
|
|
return `${scope.tenantId}:${scope.userId}:${conversationId}`;
|
|
}
|
|
|
|
private scopesEqual(a: ActorTenantScope, b: ActorTenantScope): boolean {
|
|
return a.userId === b.userId && a.tenantId === b.tenantId;
|
|
}
|
|
|
|
@SubscribeMessage('message')
|
|
async handleMessage(
|
|
@ConnectedSocket() client: Socket,
|
|
@MessageBody() rawData: unknown,
|
|
): Promise<void> {
|
|
let discordIngress: DiscordIngressPayload | null = null;
|
|
let data: ChatSocketMessageDto;
|
|
if (client.data.discordService) {
|
|
if (!isDiscordIngressEnvelope(rawData)) {
|
|
this.logger.warn(`Rejected malformed Discord ingress from ${client.id}`);
|
|
return;
|
|
}
|
|
discordIngress = this.resolveDiscordIngress(client, rawData);
|
|
if (!discordIngress) return;
|
|
data = { conversationId: discordIngress.conversationId, content: discordIngress.content };
|
|
} else {
|
|
if (!isChatSocketMessage(rawData)) {
|
|
this.logger.warn(`Rejected malformed chat message from ${client.id}`);
|
|
return;
|
|
}
|
|
data = rawData;
|
|
}
|
|
const conversationId = data.conversationId ?? uuid();
|
|
const discordServiceUserId = process.env['DISCORD_SERVICE_USER_ID'];
|
|
if (discordIngress && !discordServiceUserId) {
|
|
this.logger.warn(
|
|
`Rejected Discord ingress without configured service owner from ${client.id}`,
|
|
);
|
|
return;
|
|
}
|
|
const scope = discordIngress
|
|
? {
|
|
userId: discordServiceUserId!,
|
|
tenantId: process.env['DISCORD_SERVICE_TENANT_ID'] ?? discordServiceUserId!,
|
|
}
|
|
: this.getClientScope(client);
|
|
if (!scope) {
|
|
client.emit('error', { conversationId, error: 'Authenticated user scope is required.' });
|
|
return;
|
|
}
|
|
const userId = scope.userId;
|
|
const correlationId = discordIngress?.correlationId;
|
|
|
|
this.logger.log(
|
|
`Message from ${client.id} in conversation ${conversationId}${correlationId ? ` correlation=${correlationId}` : ''}`,
|
|
);
|
|
|
|
// Ensure agent session exists for this conversation
|
|
let sessionRoutingDecision: RoutingDecisionInfo | undefined;
|
|
try {
|
|
let agentSession = this.agentService.getSession(conversationId, scope);
|
|
if (!agentSession) {
|
|
// When resuming an existing conversation, load prior messages to inject as context (M1-004)
|
|
const conversationHistory = await this.loadConversationHistory(conversationId, userId);
|
|
|
|
// M5-004: Check if there's an existing sessionId bound to this conversation
|
|
let existingSessionId: string | undefined;
|
|
if (userId) {
|
|
existingSessionId = await this.getConversationSessionId(conversationId, userId);
|
|
if (existingSessionId) {
|
|
this.logger.log(
|
|
`Resuming existing sessionId=${existingSessionId} for conversation=${conversationId}`,
|
|
);
|
|
}
|
|
}
|
|
|
|
// Determine provider/model via routing engine or per-session /model override (M4-012 / M4-007)
|
|
let resolvedProvider = data.provider;
|
|
let resolvedModelId = data.modelId;
|
|
|
|
const modelOverride = modelOverrides.get(this.modelOverrideKey(conversationId, scope));
|
|
if (modelOverride) {
|
|
// /model override bypasses routing engine (M4-007)
|
|
resolvedModelId = modelOverride;
|
|
this.logger.log(
|
|
`Using /model override "${modelOverride}" for conversation=${conversationId}`,
|
|
);
|
|
} else if (!resolvedProvider && !resolvedModelId) {
|
|
// No explicit provider/model from client — use routing engine (M4-012)
|
|
try {
|
|
const routingDecision = await this.routingEngine.resolve(data.content, userId);
|
|
resolvedProvider = routingDecision.provider;
|
|
resolvedModelId = routingDecision.model;
|
|
sessionRoutingDecision = {
|
|
model: routingDecision.model,
|
|
provider: routingDecision.provider,
|
|
ruleName: routingDecision.ruleName,
|
|
reason: routingDecision.reason,
|
|
};
|
|
this.logger.log(
|
|
`Routing decision for conversation=${conversationId}: ${routingDecision.provider}/${routingDecision.model} (rule="${routingDecision.ruleName}")`,
|
|
);
|
|
} catch (routingErr) {
|
|
this.logger.warn(
|
|
`Routing engine failed for conversation=${conversationId}, using defaults`,
|
|
routingErr instanceof Error ? routingErr.message : String(routingErr),
|
|
);
|
|
}
|
|
}
|
|
|
|
// M5-004: Use existingSessionId as sessionId when available (session reuse)
|
|
const sessionIdToCreate = existingSessionId ?? conversationId;
|
|
agentSession = await this.agentService.createSession(sessionIdToCreate, {
|
|
provider: resolvedProvider,
|
|
modelId: resolvedModelId,
|
|
agentConfigId: data.agentId,
|
|
userId,
|
|
tenantId: scope.tenantId,
|
|
conversationHistory: conversationHistory.length > 0 ? conversationHistory : undefined,
|
|
});
|
|
|
|
if (conversationHistory.length > 0) {
|
|
this.logger.log(
|
|
`Loaded ${conversationHistory.length} prior messages for conversation=${conversationId}`,
|
|
);
|
|
}
|
|
}
|
|
} catch (err) {
|
|
this.logger.error(
|
|
`Session creation failed for client=${client.id}, conversation=${conversationId}`,
|
|
err instanceof Error ? err.stack : String(err),
|
|
);
|
|
client.emit('error', {
|
|
conversationId,
|
|
error: 'Failed to start agent session. Please try again.',
|
|
});
|
|
return;
|
|
}
|
|
|
|
// Ensure conversation record exists in the DB before persisting messages
|
|
// M5-004: Also bind the sessionId to the conversation record
|
|
if (userId) {
|
|
await this.ensureConversation(conversationId, userId);
|
|
await this.bindSessionToConversation(conversationId, userId, conversationId);
|
|
}
|
|
|
|
// M5-007: Count the user message
|
|
this.agentService.recordMessage(conversationId);
|
|
|
|
// Persist the user message
|
|
if (userId) {
|
|
try {
|
|
await this.brain.conversations.addMessage(
|
|
{
|
|
conversationId,
|
|
role: 'user',
|
|
content: redactSensitiveContent(data.content).content,
|
|
metadata: {
|
|
timestamp: new Date().toISOString(),
|
|
...(correlationId
|
|
? {
|
|
correlationId,
|
|
discordMessageId: discordIngress?.messageId,
|
|
discordUserId: discordIngress?.userId,
|
|
}
|
|
: {}),
|
|
classifications: redactSensitiveContent(data.content).classifications,
|
|
},
|
|
},
|
|
userId,
|
|
);
|
|
} catch (err) {
|
|
this.logger.error(
|
|
`Failed to persist user message for conversation=${conversationId}`,
|
|
err instanceof Error ? err.stack : String(err),
|
|
);
|
|
}
|
|
}
|
|
|
|
// Always clean up previous listener to prevent leak
|
|
const existing = this.clientSessions.get(client.id);
|
|
if (existing) {
|
|
existing.cleanup();
|
|
}
|
|
|
|
// Subscribe to agent events and relay to client
|
|
const cleanup = this.agentService.onEvent(
|
|
conversationId,
|
|
(event: AgentSessionEvent) => {
|
|
this.relayEvent(client, conversationId, event);
|
|
},
|
|
scope,
|
|
);
|
|
|
|
// Preserve routing decision from the existing client session if we didn't get a new one
|
|
const prevClientSession = this.clientSessions.get(client.id);
|
|
const routingDecisionToStore = sessionRoutingDecision ?? prevClientSession?.lastRoutingDecision;
|
|
|
|
this.clientSessions.set(client.id, {
|
|
conversationId,
|
|
cleanup,
|
|
assistantText: '',
|
|
toolCalls: [],
|
|
pendingToolCalls: new Map(),
|
|
scope,
|
|
lastRoutingDecision: routingDecisionToStore,
|
|
});
|
|
|
|
// Track channel connection
|
|
this.agentService.addChannel(conversationId, `websocket:${client.id}`, scope);
|
|
|
|
// Send session info so the client knows the model/provider (M4-008: include routing decision)
|
|
// Include agentName when a named agent config is active (M5-001)
|
|
{
|
|
const agentSession = this.agentService.getSession(conversationId, scope);
|
|
if (agentSession) {
|
|
const piSession = agentSession.piSession;
|
|
client.emit('session:info', {
|
|
conversationId,
|
|
provider: agentSession.provider,
|
|
modelId: agentSession.modelId,
|
|
thinkingLevel: piSession.thinkingLevel,
|
|
availableThinkingLevels: piSession.getAvailableThinkingLevels(),
|
|
...(agentSession.agentName ? { agentName: agentSession.agentName } : {}),
|
|
...(routingDecisionToStore ? { routingDecision: routingDecisionToStore } : {}),
|
|
});
|
|
}
|
|
}
|
|
|
|
// Send acknowledgment
|
|
client.emit('message:ack', {
|
|
conversationId,
|
|
messageId: uuid(),
|
|
...(correlationId
|
|
? {
|
|
correlationId,
|
|
discordMessageId: discordIngress?.messageId,
|
|
discordUserId: discordIngress?.userId,
|
|
}
|
|
: {}),
|
|
});
|
|
|
|
// Dispatch to agent
|
|
try {
|
|
await this.agentService.prompt(conversationId, data.content, scope);
|
|
} catch (err) {
|
|
this.logger.error(
|
|
`Agent prompt failed for client=${client.id}, conversation=${conversationId}`,
|
|
err instanceof Error ? err.stack : String(err),
|
|
);
|
|
client.emit('error', {
|
|
conversationId,
|
|
error: 'The agent failed to process your message. Please try again.',
|
|
});
|
|
}
|
|
}
|
|
|
|
@SubscribeMessage('set:thinking')
|
|
handleSetThinking(
|
|
@ConnectedSocket() client: Socket,
|
|
@MessageBody() data: SetThinkingPayload,
|
|
): void {
|
|
const scope = this.getClientScope(client);
|
|
if (!scope) {
|
|
client.emit('error', {
|
|
conversationId: data.conversationId,
|
|
error: 'Authenticated user scope is required.',
|
|
});
|
|
return;
|
|
}
|
|
|
|
const session = this.agentService.getSession(data.conversationId, scope);
|
|
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(),
|
|
...(session.agentName ? { agentName: session.agentName } : {}),
|
|
});
|
|
}
|
|
|
|
@SubscribeMessage('abort')
|
|
async handleAbort(
|
|
@ConnectedSocket() client: Socket,
|
|
@MessageBody() data: AbortPayload,
|
|
): Promise<void> {
|
|
const conversationId = data.conversationId;
|
|
this.logger.log(`Abort requested by ${client.id} for conversation ${conversationId}`);
|
|
|
|
const scope = this.getClientScope(client);
|
|
if (!scope) {
|
|
client.emit('error', { conversationId, error: 'Authenticated user scope is required.' });
|
|
return;
|
|
}
|
|
|
|
const session = this.agentService.getSession(conversationId, scope);
|
|
if (!session) {
|
|
client.emit('error', {
|
|
conversationId,
|
|
error: 'No active session to abort.',
|
|
});
|
|
return;
|
|
}
|
|
|
|
try {
|
|
await session.piSession.abort();
|
|
this.logger.log(`Agent session ${conversationId} aborted successfully`);
|
|
} catch (err) {
|
|
this.logger.error(
|
|
`Failed to abort session ${conversationId}`,
|
|
err instanceof Error ? err.stack : String(err),
|
|
);
|
|
client.emit('error', {
|
|
conversationId,
|
|
error: 'Failed to abort the agent operation.',
|
|
});
|
|
}
|
|
}
|
|
|
|
@SubscribeMessage('command:execute')
|
|
async handleCommandExecute(
|
|
@ConnectedSocket() client: Socket,
|
|
@MessageBody() payload: SlashCommandPayload,
|
|
): Promise<void> {
|
|
const scope = this.getClientScope(client);
|
|
if (!scope) {
|
|
client.emit('command:result', {
|
|
command: payload.command,
|
|
conversationId: payload.conversationId,
|
|
success: false,
|
|
message: 'Authenticated user scope is required.',
|
|
});
|
|
return;
|
|
}
|
|
|
|
const result = await this.commandExecutor.execute(payload, scope);
|
|
client.emit('command:result', result);
|
|
}
|
|
|
|
@SubscribeMessage('command:approve')
|
|
async handleCommandApproval(
|
|
@ConnectedSocket() client: Socket,
|
|
@MessageBody() payload: SlashCommandPayload,
|
|
): Promise<void> {
|
|
const scope = this.getClientScope(client);
|
|
const approval = scope ? await this.commandExecutor.createApproval(payload, scope) : null;
|
|
const result: SlashCommandApprovalResultPayload = approval
|
|
? {
|
|
command: payload.command,
|
|
conversationId: payload.conversationId,
|
|
success: true,
|
|
approvalId: approval.approvalId,
|
|
expiresAt: approval.expiresAt,
|
|
}
|
|
: {
|
|
command: payload.command,
|
|
conversationId: payload.conversationId,
|
|
success: false,
|
|
message: 'Not authorized to approve this command.',
|
|
};
|
|
client.emit('command:approval', result);
|
|
}
|
|
|
|
broadcastReload(payload: SystemReloadPayload): void {
|
|
this.server.emit('system:reload', payload);
|
|
this.logger.log('Broadcasted system:reload to all connected clients');
|
|
}
|
|
|
|
/**
|
|
* Set a per-conversation model override (M4-007 / M5-002).
|
|
* When set, the routing engine is bypassed and the specified model is used.
|
|
* Pass null to clear the override and resume automatic routing.
|
|
* M5-005: Emits session:info to clients subscribed to this conversation when a model is set.
|
|
* M5-007: Records a model switch in session metrics.
|
|
*/
|
|
setModelOverride(
|
|
conversationId: string,
|
|
modelName: string | null,
|
|
scope: ActorTenantScope,
|
|
): void {
|
|
const key = this.modelOverrideKey(conversationId, scope);
|
|
if (modelName) {
|
|
modelOverrides.set(key, modelName);
|
|
this.logger.log(`Model override set: conversation=${conversationId} model="${modelName}"`);
|
|
|
|
// M5-002: Update the live session's modelId so session:info reflects the new model immediately
|
|
this.agentService.updateSessionModel(conversationId, modelName, scope);
|
|
|
|
// M5-005: Broadcast session:info to all clients subscribed to this conversation
|
|
this.broadcastSessionInfo(conversationId, scope);
|
|
} else {
|
|
modelOverrides.delete(key);
|
|
this.logger.log(`Model override cleared: conversation=${conversationId}`);
|
|
}
|
|
}
|
|
|
|
/**
|
|
* Return the active model override for a conversation, or undefined if none.
|
|
*/
|
|
getModelOverride(conversationId: string, scope: ActorTenantScope): string | undefined {
|
|
return modelOverrides.get(this.modelOverrideKey(conversationId, scope));
|
|
}
|
|
|
|
/**
|
|
* M5-005: Broadcast session:info to all clients currently subscribed to a conversation.
|
|
* Called on model or agent switch to ensure the TUI TopBar updates immediately.
|
|
*/
|
|
broadcastSessionInfo(
|
|
conversationId: string,
|
|
scope: ActorTenantScope,
|
|
extra?: { agentName?: string; routingDecision?: RoutingDecisionInfo },
|
|
): void {
|
|
const agentSession = this.agentService.getSession(conversationId, scope);
|
|
if (!agentSession) return;
|
|
|
|
const piSession = agentSession.piSession;
|
|
const resolvedAgentName = extra?.agentName ?? agentSession.agentName;
|
|
const payload = {
|
|
conversationId,
|
|
provider: agentSession.provider,
|
|
modelId: agentSession.modelId,
|
|
thinkingLevel: piSession.thinkingLevel,
|
|
availableThinkingLevels: piSession.getAvailableThinkingLevels(),
|
|
...(resolvedAgentName ? { agentName: resolvedAgentName } : {}),
|
|
...(extra?.routingDecision ? { routingDecision: extra.routingDecision } : {}),
|
|
};
|
|
|
|
// Emit to all clients currently subscribed to this conversation
|
|
for (const [clientId, session] of this.clientSessions) {
|
|
if (session.conversationId === conversationId && this.scopesEqual(session.scope, scope)) {
|
|
const socket = this.server.sockets.sockets.get(clientId);
|
|
if (socket?.connected) {
|
|
socket.emit('session:info', payload);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
/**
|
|
* Ensure a conversation record exists in the DB.
|
|
* Creates it if absent — safe to call concurrently since a duplicate insert
|
|
* would fail on the PK constraint and be caught here.
|
|
*/
|
|
@SubscribeMessage('discord:approve')
|
|
async handleDiscordApproval(
|
|
@ConnectedSocket() client: Socket,
|
|
@MessageBody() envelope: DiscordIngressEnvelope,
|
|
): Promise<void> {
|
|
if (!client.data.discordService) return;
|
|
const ingress = this.resolveDiscordIngress(client, envelope, 'approve');
|
|
const actionParts = ingress?.content.match(/^\/approve\s+([^\s]+)\s+([^\s]+)$/i);
|
|
const tenantId = process.env['DISCORD_SERVICE_TENANT_ID']?.trim();
|
|
if (!ingress || !actionParts || !tenantId || !this.commandAuthorization) return;
|
|
const binding = resolveDiscordInteractionBinding(
|
|
parseDiscordInteractionBindings(process.env['DISCORD_INTERACTION_BINDINGS']),
|
|
ingress.guildId,
|
|
ingress.channelId,
|
|
ingress.userId,
|
|
'approve',
|
|
);
|
|
const actorId = binding && resolveDiscordInteractionActorId(binding, ingress.userId);
|
|
const agentName = process.env['MOSAIC_AGENT_NAME']?.trim();
|
|
if (!actorId || !agentName || binding.instanceId !== agentName) {
|
|
this.logger.warn(
|
|
`Rejected Discord approval without a matching runtime agent from ${client.id}`,
|
|
);
|
|
client.emit('discord:approval', {
|
|
correlationId: ingress.correlationId,
|
|
success: false,
|
|
approvalId: undefined,
|
|
expiresAt: undefined,
|
|
});
|
|
return;
|
|
}
|
|
const providerId = actionParts[1]!;
|
|
const sessionId = actionParts[2]!;
|
|
const approval = await this.commandAuthorization.createRuntimeTerminationApproval({
|
|
providerId,
|
|
sessionId,
|
|
actorId,
|
|
tenantId,
|
|
channelId: ingress.channelId,
|
|
correlationId: this.discordRuntimeActionCorrelation(
|
|
binding.instanceId,
|
|
ingress,
|
|
providerId,
|
|
sessionId,
|
|
),
|
|
agentName,
|
|
});
|
|
client.emit('discord:approval', {
|
|
correlationId: ingress.correlationId,
|
|
success: approval !== null,
|
|
approvalId: approval?.approvalId,
|
|
expiresAt: approval?.expiresAt,
|
|
});
|
|
}
|
|
|
|
@SubscribeMessage('discord:stop')
|
|
async handleDiscordStop(
|
|
@ConnectedSocket() client: Socket,
|
|
@MessageBody() envelope: DiscordIngressEnvelope,
|
|
): Promise<void> {
|
|
if (!client.data.discordService) return;
|
|
const ingress = this.resolveDiscordIngress(client, envelope, 'stop');
|
|
const actionParts = ingress?.content.match(/^\/stop\s+([^\s]+)\s+([^\s]+)\s+([^\s]+)$/i);
|
|
const tenantId = process.env['DISCORD_SERVICE_TENANT_ID']?.trim();
|
|
if (!ingress || !actionParts || !tenantId || !this.runtimeRegistry) return;
|
|
const binding = resolveDiscordInteractionBinding(
|
|
parseDiscordInteractionBindings(process.env['DISCORD_INTERACTION_BINDINGS']),
|
|
ingress.guildId,
|
|
ingress.channelId,
|
|
ingress.userId,
|
|
'stop',
|
|
);
|
|
const actorId = binding && resolveDiscordInteractionActorId(binding, ingress.userId);
|
|
if (!actorId) return;
|
|
const providerId = actionParts[1]!;
|
|
const sessionId = actionParts[2]!;
|
|
|
|
try {
|
|
// RuntimeProviderService consumes the durable approval exactly once using the
|
|
// provisioned approving-admin identity, never the Discord service account.
|
|
await this.runtimeRegistry.terminate(providerId, sessionId, actionParts[3]!, {
|
|
actorScope: {
|
|
userId: actorId,
|
|
tenantId,
|
|
},
|
|
channelId: ingress.channelId,
|
|
correlationId: this.discordRuntimeActionCorrelation(
|
|
binding.instanceId,
|
|
ingress,
|
|
providerId,
|
|
sessionId,
|
|
),
|
|
});
|
|
client.emit('discord:stop', { correlationId: ingress.correlationId, success: true });
|
|
} catch {
|
|
client.emit('discord:stop', { correlationId: ingress.correlationId, success: false });
|
|
}
|
|
}
|
|
|
|
/**
|
|
* Correlates the immutable termination target rather than either Discord message.
|
|
* Approval and stop are distinct ingress events, but must consume the same seven-field action.
|
|
*/
|
|
private discordRuntimeActionCorrelation(
|
|
instanceId: string,
|
|
ingress: DiscordIngressPayload,
|
|
providerId: string,
|
|
sessionId: string,
|
|
): string {
|
|
const target = [
|
|
instanceId,
|
|
ingress.guildId,
|
|
ingress.channelId,
|
|
ingress.conversationId,
|
|
providerId,
|
|
sessionId,
|
|
];
|
|
return `discord-action:v1:${createHash('sha256').update(JSON.stringify(target)).digest('hex')}`;
|
|
}
|
|
|
|
private resolveDiscordIngress(
|
|
client: Socket,
|
|
envelope: DiscordIngressEnvelope,
|
|
operation: 'send' | 'approve' | 'stop' = 'send',
|
|
): DiscordIngressPayload | null {
|
|
const payload = verifyDiscordIngressEnvelope(
|
|
envelope,
|
|
process.env['DISCORD_SERVICE_TOKEN'] ?? '',
|
|
{
|
|
guildIds: this.readDiscordAllowlist('DISCORD_ALLOWED_GUILD_IDS'),
|
|
channelIds: this.readDiscordAllowlist('DISCORD_ALLOWED_CHANNEL_IDS'),
|
|
userIds: this.readDiscordAllowlist('DISCORD_ALLOWED_USER_IDS'),
|
|
},
|
|
);
|
|
if (!payload) {
|
|
this.logger.warn(`Rejected invalid Discord ingress envelope from ${client.id}`);
|
|
return null;
|
|
}
|
|
try {
|
|
const binding = resolveDiscordInteractionBinding(
|
|
parseDiscordInteractionBindings(process.env['DISCORD_INTERACTION_BINDINGS']),
|
|
payload.guildId,
|
|
payload.channelId,
|
|
payload.userId,
|
|
operation,
|
|
);
|
|
if (!binding) {
|
|
this.logger.warn(`Rejected unpaired Discord ingress from ${client.id}`);
|
|
return null;
|
|
}
|
|
} catch {
|
|
this.logger.warn(
|
|
`Rejected Discord ingress without valid binding configuration from ${client.id}`,
|
|
);
|
|
return null;
|
|
}
|
|
if (!this.discordReplayProtector.claim(payload.messageId)) {
|
|
this.logger.warn(
|
|
`Rejected replayed Discord message=${payload.messageId} correlation=${payload.correlationId}`,
|
|
);
|
|
return null;
|
|
}
|
|
return payload;
|
|
}
|
|
|
|
private readDiscordAllowlist(name: string): string[] {
|
|
return (process.env[name] ?? '')
|
|
.split(',')
|
|
.map((id: string): string => id.trim())
|
|
.filter((id: string): boolean => id.length > 0);
|
|
}
|
|
|
|
private async ensureConversation(conversationId: string, userId: string): Promise<void> {
|
|
try {
|
|
const existing = await this.brain.conversations.findById(conversationId, userId);
|
|
if (!existing) {
|
|
await this.brain.conversations.create({
|
|
id: conversationId,
|
|
userId,
|
|
});
|
|
}
|
|
} catch (err) {
|
|
this.logger.error(
|
|
`Failed to ensure conversation record for conversation=${conversationId}`,
|
|
err instanceof Error ? err.stack : String(err),
|
|
);
|
|
}
|
|
}
|
|
|
|
/**
|
|
* M5-004: Bind the agent sessionId to the conversation record in the DB.
|
|
* Updates the sessionId column so future resumes can reuse the session.
|
|
*/
|
|
private async bindSessionToConversation(
|
|
conversationId: string,
|
|
userId: string,
|
|
sessionId: string,
|
|
): Promise<void> {
|
|
try {
|
|
await this.brain.conversations.update(conversationId, userId, { sessionId });
|
|
} catch (err) {
|
|
this.logger.error(
|
|
`Failed to bind sessionId=${sessionId} to conversation=${conversationId}`,
|
|
err instanceof Error ? err.stack : String(err),
|
|
);
|
|
}
|
|
}
|
|
|
|
/**
|
|
* M5-004: Retrieve the sessionId bound to a conversation, if any.
|
|
* Returns undefined when the conversation does not exist or has no bound session.
|
|
*/
|
|
private async getConversationSessionId(
|
|
conversationId: string,
|
|
userId: string,
|
|
): Promise<string | undefined> {
|
|
try {
|
|
const conv = await this.brain.conversations.findById(conversationId, userId);
|
|
return conv?.sessionId ?? undefined;
|
|
} catch (err) {
|
|
this.logger.error(
|
|
`Failed to get sessionId for conversation=${conversationId}`,
|
|
err instanceof Error ? err.stack : String(err),
|
|
);
|
|
return undefined;
|
|
}
|
|
}
|
|
|
|
/**
|
|
* Load prior conversation messages from DB for context injection on session resume (M1-004).
|
|
* Returns an empty array when no history exists, the conversation is not owned by the user,
|
|
* or userId is not provided.
|
|
*/
|
|
private async loadConversationHistory(
|
|
conversationId: string,
|
|
userId: string | undefined,
|
|
): Promise<ConversationHistoryMessage[]> {
|
|
if (!userId) return [];
|
|
|
|
try {
|
|
const messages = await this.brain.conversations.findMessages(conversationId, userId);
|
|
if (messages.length === 0) return [];
|
|
|
|
return messages.map((msg) => ({
|
|
role: msg.role as 'user' | 'assistant' | 'system',
|
|
content: msg.content,
|
|
createdAt: msg.createdAt,
|
|
}));
|
|
} catch (err) {
|
|
this.logger.error(
|
|
`Failed to load conversation history for conversation=${conversationId}`,
|
|
err instanceof Error ? err.stack : String(err),
|
|
);
|
|
return [];
|
|
}
|
|
}
|
|
|
|
private appendAndFlushRedactedEgress(
|
|
client: Socket,
|
|
conversationId: string,
|
|
eventName: 'agent:text' | 'agent:thinking',
|
|
buffers: Map<string, string>,
|
|
delta: string,
|
|
): void {
|
|
const key = this.egressKey(client, eventName);
|
|
if (this.overflowedEgress.has(key)) return;
|
|
|
|
const buffered = `${buffers.get(client.id) ?? ''}${delta}`;
|
|
if (buffered.length > MAX_REDACTION_BUFFER_LENGTH) {
|
|
buffers.delete(client.id);
|
|
this.overflowedEgress.add(key);
|
|
client.emit(eventName, { conversationId, text: '[REDACTED_STREAM_OVERFLOW]' });
|
|
return;
|
|
}
|
|
|
|
buffers.set(client.id, buffered);
|
|
this.flushRedactedEgress(client, conversationId, eventName, buffers, false);
|
|
}
|
|
|
|
/**
|
|
* Holds any suffix that could become a secret, email, or phone number after a
|
|
* later stream chunk. This avoids relying on downstream redaction after data
|
|
* has already reached the socket.
|
|
*/
|
|
private flushRedactedEgress(
|
|
client: Socket,
|
|
conversationId: string,
|
|
eventName: 'agent:text' | 'agent:thinking',
|
|
buffers: Map<string, string>,
|
|
final: boolean,
|
|
): void {
|
|
const key = this.egressKey(client, eventName);
|
|
if (this.overflowedEgress.has(key)) {
|
|
if (final) this.overflowedEgress.delete(key);
|
|
return;
|
|
}
|
|
|
|
const buffered = buffers.get(client.id) ?? '';
|
|
const releaseLength = final ? buffered.length : this.safeRedactionPrefixLength(buffered);
|
|
const released = buffered.slice(0, releaseLength);
|
|
const pending = buffered.slice(releaseLength);
|
|
|
|
if (pending) {
|
|
buffers.set(client.id, pending);
|
|
} else {
|
|
buffers.delete(client.id);
|
|
}
|
|
|
|
if (released) {
|
|
client.emit(eventName, {
|
|
conversationId,
|
|
text: redactSensitiveContent(released).content,
|
|
});
|
|
}
|
|
}
|
|
|
|
private safeRedactionPrefixLength(content: string): number {
|
|
let retainedFrom = content.length;
|
|
|
|
// Retain the current token because it may become a split secret or email.
|
|
const token = /(?:^|\s)(\S*)$/.exec(content);
|
|
if (token) {
|
|
const matched = token[0] ?? '';
|
|
const trailingToken = token[1] ?? '';
|
|
retainedFrom = token.index + matched.length - trailingToken.length;
|
|
}
|
|
|
|
// The secret classifier accepts whitespace around ':' and '=', so preserve
|
|
// a pending label until its value and delimiter are both complete.
|
|
const pendingSecretLabel =
|
|
/(?:^|[^A-Za-z0-9_])((?:api[_-]?key|token|password|secret|bearer|authorization)\s*)$/i.exec(
|
|
content,
|
|
);
|
|
if (pendingSecretLabel) {
|
|
const label = pendingSecretLabel[1] ?? '';
|
|
retainedFrom = Math.min(
|
|
retainedFrom,
|
|
pendingSecretLabel.index + pendingSecretLabel[0].length - label.length,
|
|
);
|
|
}
|
|
|
|
const secretLabel = /(?:api[_-]?key|token|password|secret|authorization)\s*[:=]\s*$/i.exec(
|
|
content,
|
|
);
|
|
if (secretLabel) {
|
|
retainedFrom = Math.min(retainedFrom, secretLabel.index);
|
|
}
|
|
|
|
// Phone numbers can contain whitespace and punctuation; preserve the full
|
|
// trailing numeric candidate until a non-phone character establishes a boundary.
|
|
const phone = /(?:^|[^A-Za-z0-9_])(\+?\d[\d(). -]*)$/.exec(content);
|
|
if (phone) {
|
|
const matched = phone[0] ?? '';
|
|
const trailingPhoneCandidate = phone[1] ?? '';
|
|
retainedFrom = Math.min(
|
|
retainedFrom,
|
|
phone.index + matched.length - trailingPhoneCandidate.length,
|
|
);
|
|
}
|
|
|
|
const privateKeyStart = content.lastIndexOf('-----BEGIN');
|
|
if (privateKeyStart >= 0) {
|
|
const privateKey = content.slice(privateKeyStart);
|
|
if (/-----END(?: [A-Z]+)* KEY-----/.test(privateKey)) {
|
|
// Release the complete block in one pass so the full-block classifier can redact it.
|
|
retainedFrom = content.length;
|
|
} else {
|
|
retainedFrom = Math.min(retainedFrom, privateKeyStart);
|
|
}
|
|
}
|
|
|
|
return retainedFrom;
|
|
}
|
|
|
|
private egressKey(client: Socket, eventName: 'agent:text' | 'agent:thinking'): string {
|
|
return `${client.id}:${eventName}`;
|
|
}
|
|
|
|
private relayEvent(client: Socket, conversationId: string, event: AgentSessionEvent): void {
|
|
if (!client.connected) {
|
|
this.logger.warn(
|
|
`Dropping event ${event.type} for disconnected client=${client.id}, conversation=${conversationId}`,
|
|
);
|
|
return;
|
|
}
|
|
|
|
switch (event.type) {
|
|
case 'agent_start': {
|
|
// Reset accumulation buffers for the new turn
|
|
const cs = this.clientSessions.get(client.id);
|
|
if (cs) {
|
|
cs.assistantText = '';
|
|
cs.toolCalls = [];
|
|
cs.pendingToolCalls.clear();
|
|
}
|
|
this.textEgressBuffers.set(client.id, '');
|
|
this.thinkingEgressBuffers.set(client.id, '');
|
|
this.overflowedEgress.delete(this.egressKey(client, 'agent:text'));
|
|
this.overflowedEgress.delete(this.egressKey(client, 'agent:thinking'));
|
|
client.emit('agent:start', { conversationId });
|
|
break;
|
|
}
|
|
|
|
case 'agent_end': {
|
|
// Gather usage stats from the Pi session
|
|
const activeClientSession = this.clientSessions.get(client.id);
|
|
const agentSession = activeClientSession
|
|
? this.agentService.getSession(conversationId, activeClientSession.scope)
|
|
: undefined;
|
|
const piSession = agentSession?.piSession;
|
|
const stats = piSession?.getSessionStats();
|
|
const contextUsage = piSession?.getContextUsage();
|
|
|
|
const usagePayload = 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;
|
|
|
|
this.flushRedactedEgress(
|
|
client,
|
|
conversationId,
|
|
'agent:text',
|
|
this.textEgressBuffers,
|
|
true,
|
|
);
|
|
this.flushRedactedEgress(
|
|
client,
|
|
conversationId,
|
|
'agent:thinking',
|
|
this.thinkingEgressBuffers,
|
|
true,
|
|
);
|
|
client.emit('agent:end', {
|
|
conversationId,
|
|
usage: usagePayload,
|
|
});
|
|
|
|
// M5-007: Accumulate token usage in session metrics
|
|
if (stats?.tokens) {
|
|
this.agentService.recordTokenUsage(conversationId, {
|
|
input: stats.tokens.input ?? 0,
|
|
output: stats.tokens.output ?? 0,
|
|
cacheRead: stats.tokens.cacheRead ?? 0,
|
|
cacheWrite: stats.tokens.cacheWrite ?? 0,
|
|
total: stats.tokens.total ?? 0,
|
|
});
|
|
}
|
|
|
|
// Persist the assistant message with metadata
|
|
const cs = this.clientSessions.get(client.id);
|
|
const userId = (client.data.user as { id: string } | undefined)?.id;
|
|
if (cs && userId && cs.assistantText.trim().length > 0) {
|
|
const metadata: Record<string, unknown> = {
|
|
timestamp: new Date().toISOString(),
|
|
model: agentSession?.modelId ?? 'unknown',
|
|
provider: agentSession?.provider ?? 'unknown',
|
|
toolCalls: cs.toolCalls,
|
|
};
|
|
|
|
if (stats?.tokens) {
|
|
metadata['tokenUsage'] = {
|
|
input: stats.tokens.input,
|
|
output: stats.tokens.output,
|
|
cacheRead: stats.tokens.cacheRead,
|
|
cacheWrite: stats.tokens.cacheWrite,
|
|
total: stats.tokens.total,
|
|
};
|
|
}
|
|
|
|
this.brain.conversations
|
|
.addMessage(
|
|
{
|
|
conversationId,
|
|
role: 'assistant',
|
|
content: redactSensitiveContent(cs.assistantText).content,
|
|
metadata: {
|
|
...metadata,
|
|
classifications: redactSensitiveContent(cs.assistantText).classifications,
|
|
},
|
|
},
|
|
userId,
|
|
)
|
|
.catch((err: unknown) => {
|
|
this.logger.error(
|
|
`Failed to persist assistant message for conversation=${conversationId}`,
|
|
err instanceof Error ? err.stack : String(err),
|
|
);
|
|
});
|
|
|
|
// Reset accumulation
|
|
cs.assistantText = '';
|
|
cs.toolCalls = [];
|
|
cs.pendingToolCalls.clear();
|
|
}
|
|
break;
|
|
}
|
|
|
|
case 'message_update': {
|
|
const assistantEvent = event.assistantMessageEvent;
|
|
if (assistantEvent.type === 'text_delta') {
|
|
// Keep raw stream material in memory only; persist and emit only redacted text.
|
|
const cs = this.clientSessions.get(client.id);
|
|
if (cs) {
|
|
cs.assistantText += assistantEvent.delta;
|
|
}
|
|
this.appendAndFlushRedactedEgress(
|
|
client,
|
|
conversationId,
|
|
'agent:text',
|
|
this.textEgressBuffers,
|
|
assistantEvent.delta,
|
|
);
|
|
} else if (assistantEvent.type === 'thinking_delta') {
|
|
this.appendAndFlushRedactedEgress(
|
|
client,
|
|
conversationId,
|
|
'agent:thinking',
|
|
this.thinkingEgressBuffers,
|
|
assistantEvent.delta,
|
|
);
|
|
}
|
|
break;
|
|
}
|
|
|
|
case 'tool_execution_start': {
|
|
// Track pending tool call for later recording
|
|
const cs = this.clientSessions.get(client.id);
|
|
if (cs) {
|
|
cs.pendingToolCalls.set(event.toolCallId, {
|
|
toolName: event.toolName,
|
|
args: event.args,
|
|
});
|
|
}
|
|
client.emit('agent:tool:start', {
|
|
conversationId,
|
|
toolCallId: event.toolCallId,
|
|
toolName: event.toolName,
|
|
});
|
|
break;
|
|
}
|
|
|
|
case 'tool_execution_end': {
|
|
// Finalise tool call record
|
|
const cs = this.clientSessions.get(client.id);
|
|
if (cs) {
|
|
const pending = cs.pendingToolCalls.get(event.toolCallId);
|
|
cs.toolCalls.push({
|
|
toolCallId: event.toolCallId,
|
|
toolName: event.toolName,
|
|
args: pending?.args ?? null,
|
|
isError: event.isError,
|
|
});
|
|
cs.pendingToolCalls.delete(event.toolCallId);
|
|
}
|
|
client.emit('agent:tool:end', {
|
|
conversationId,
|
|
toolCallId: event.toolCallId,
|
|
toolName: event.toolName,
|
|
isError: event.isError,
|
|
});
|
|
break;
|
|
}
|
|
}
|
|
}
|
|
}
|