44 * Mirrors the SSE chat flow but streams AgentEvent payloads over WebSocket.
55 */
66
7+ import { randomBytes , randomUUID } from "node:crypto" ;
78import type { IncomingMessage } from "node:http" ;
89import type { ComposerChatRequest , ComposerMessage } from "@evalops/contracts" ;
910import type { RawData , WebSocket } from "ws" ;
11+ import { applyOracleConsultationDirective } from "../../agent/oracle-consultation-policy.js" ;
1012import { isAssistantMessage } from "../../agent/type-guards.js" ;
1113import type {
1214 Attachment as AgentAttachment ,
@@ -22,7 +24,13 @@ import {
2224} from "../../memory/auto-consolidation.js" ;
2325import { createAutomaticMemoryExtractionCoordinator } from "../../memory/auto-extraction.js" ;
2426import type { RegisteredModel } from "../../models/registry.js" ;
27+ import {
28+ recordIntelligentRouterChatMetric ,
29+ selectIntelligentRouterModel ,
30+ } from "../../services/intelligent-router/recorder.js" ;
2531import { recordAssistantUsageMetric } from "../../services/usage-analytics/recorder.js" ;
32+ import { evaluateModelPolicy } from "../../services/workspace-config/policy.js" ;
33+ import type { WorkspaceConfigRequestContext } from "../../services/workspace-config/types.js" ;
2634import { toSessionModelMetadata } from "../../session/manager.js" ;
2735import { createRuntimeSessionSummaryUpdater } from "../../session/runtime-summary-updater.js" ;
2836import { createLogger } from "../../utils/logger.js" ;
@@ -37,6 +45,12 @@ import { getAuthSubject } from "../authz.js";
3745import { getAgentCircuitBreaker } from "../circuit-breaker.js" ;
3846import { clientToolService } from "../client-tools-service.js" ;
3947import { isHostedSessionManager } from "../hosted-session-manager.js" ;
48+ import {
49+ type RequestContext ,
50+ getWorkspaceConfigContext ,
51+ parseTraceParent ,
52+ requestContextStorage ,
53+ } from "../request-context.js" ;
4054import { serverRequestManager } from "../server-request-manager.js" ;
4155import { getRequestHeader } from "../server-utils.js" ;
4256import { startSessionWithPolicy } from "../session-initialization.js" ;
@@ -219,6 +233,7 @@ export function handleChatWebSocket(
219233 ws : WebSocket ,
220234 req : IncomingMessage ,
221235 context : WebServerContext ,
236+ workspaceConfig ?: WorkspaceConfigRequestContext ,
222237) {
223238 const {
224239 createAgent,
@@ -245,6 +260,17 @@ export function handleChatWebSocket(
245260 ) ;
246261 const slimFromQuery = parseBoolean ( url . searchParams . get ( "slim" ) ) ;
247262 const clientHeaderFromQuery = url . searchParams . get ( "client" ) ?. trim ( ) ;
263+ const websocketRequestContext : RequestContext | undefined = workspaceConfig
264+ ? {
265+ requestId : randomUUID ( ) ,
266+ traceId : parseTraceParent ( req . headers . traceparent ) . traceId ,
267+ spanId : randomBytes ( 8 ) . toString ( "hex" ) ,
268+ startTime : performance . now ( ) ,
269+ method : req . method || "GET" ,
270+ url : url . pathname ,
271+ workspaceConfig,
272+ }
273+ : undefined ;
248274
249275 const sendErrorAndClose = ( message : string ) => {
250276 try {
@@ -271,7 +297,7 @@ export function handleChatWebSocket(
271297 return validatePayload < ChatRequestInput > ( parsed , ChatRequestSchema ) ;
272298 } ;
273299
274- ws . on ( "message" , async ( data ) => {
300+ const handleMessage = async ( data : RawData ) => {
275301 if ( requestHandled ) {
276302 try {
277303 const size = getRawDataSize ( data ) ;
@@ -450,7 +476,51 @@ export function handleChatWebSocket(
450476 }
451477 }
452478
453- const registeredModel = await getRegisteredModel ( chatReq . model ) ;
479+ const routingSelection = selectIntelligentRouterModel ( {
480+ req,
481+ requestedModel : chatReq . model ,
482+ body : chatReq ,
483+ } ) ;
484+ let registeredModel : RegisteredModel | undefined ;
485+ let lastModelError : unknown ;
486+ let selectedModelError : unknown ;
487+ let selectedModelPolicyViolation : ReturnType <
488+ typeof evaluateModelPolicy
489+ > | null = null ;
490+ for ( const [
491+ index ,
492+ modelInput ,
493+ ] of routingSelection . modelInputs . entries ( ) ) {
494+ try {
495+ const candidateModel = await getRegisteredModel ( modelInput ) ;
496+ const violation = evaluateModelPolicy (
497+ workspaceConfig ?. config ?? getWorkspaceConfigContext ( ) ?. config ,
498+ {
499+ provider : candidateModel . provider ,
500+ modelId : candidateModel . id ,
501+ } ,
502+ ) ;
503+ if ( violation ) {
504+ if ( index === 0 ) selectedModelPolicyViolation = violation ;
505+ continue ;
506+ }
507+ registeredModel = candidateModel ;
508+ break ;
509+ } catch ( error ) {
510+ lastModelError = error ;
511+ if ( index === 0 ) selectedModelError = error ;
512+ }
513+ }
514+ if ( ! registeredModel && selectedModelPolicyViolation ) {
515+ sendErrorAndClose ( selectedModelPolicyViolation . message ) ;
516+ if ( sseLease && releaseSse ) {
517+ releaseSse ( sseLease ) ;
518+ sseLease = null ;
519+ }
520+ return ;
521+ }
522+ if ( ! registeredModel ) throw selectedModelError ?? lastModelError ;
523+ const routeStartedAt = Date . now ( ) ;
454524
455525 const headerApproval = ( ( ) => {
456526 const headerMode = normalizeApprovalMode (
@@ -529,6 +599,10 @@ export function handleChatWebSocket(
529599 : { } ) ,
530600 } ,
531601 ) ;
602+ applyOracleConsultationDirective (
603+ agent ,
604+ routingSelection . decision . oracleConsultation ,
605+ ) ;
532606
533607 const historyMessages = incomingMessages . slice ( 0 , - 1 ) ;
534608 const hydratedHistory = convertComposerMessagesToApp (
@@ -844,6 +918,13 @@ export function handleChatWebSocket(
844918 sessionId : sessionManager . getSessionId ( ) ,
845919 subject,
846920 } ) ;
921+ recordIntelligentRouterChatMetric ( {
922+ taskType : routingSelection . taskType ,
923+ provider : registeredModel . provider ,
924+ model : registeredModel . id ,
925+ startedAt : routeStartedAt ,
926+ message : event . message ,
927+ } ) ;
847928 automaticMemoryExtraction . schedule ( sessionManager . getSessionFile ( ) ) ;
848929 }
849930 sessionManager . updateSnapshot (
@@ -962,5 +1043,14 @@ export function handleChatWebSocket(
9621043 sseLease = null ;
9631044 }
9641045 }
1046+ } ;
1047+
1048+ ws . on ( "message" , ( data ) => {
1049+ if ( websocketRequestContext ) {
1050+ return requestContextStorage . run ( websocketRequestContext , ( ) =>
1051+ handleMessage ( data ) ,
1052+ ) ;
1053+ }
1054+ return handleMessage ( data ) ;
9651055 } ) ;
9661056}
0 commit comments