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,13 @@ 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 { resolveModelInputForRouting } from "../model-selection.js" ;
49+ import {
50+ type RequestContext ,
51+ getWorkspaceConfigContext ,
52+ parseTraceParent ,
53+ requestContextStorage ,
54+ } from "../request-context.js" ;
4055import { serverRequestManager } from "../server-request-manager.js" ;
4156import { getRequestHeader } from "../server-utils.js" ;
4257import { startSessionWithPolicy } from "../session-initialization.js" ;
@@ -219,12 +234,15 @@ export function handleChatWebSocket(
219234 ws : WebSocket ,
220235 req : IncomingMessage ,
221236 context : WebServerContext ,
237+ workspaceConfig ?: WorkspaceConfigRequestContext ,
222238) {
223239 const {
224240 createAgent,
225241 createBackgroundAgent,
226242 getRegisteredModel,
227243 defaultApprovalMode,
244+ defaultProvider,
245+ defaultModelId,
228246 acquireSse,
229247 releaseSse,
230248 } = context ;
@@ -245,6 +263,17 @@ export function handleChatWebSocket(
245263 ) ;
246264 const slimFromQuery = parseBoolean ( url . searchParams . get ( "slim" ) ) ;
247265 const clientHeaderFromQuery = url . searchParams . get ( "client" ) ?. trim ( ) ;
266+ const websocketRequestContext : RequestContext | undefined = workspaceConfig
267+ ? {
268+ requestId : randomUUID ( ) ,
269+ traceId : parseTraceParent ( req . headers . traceparent ) . traceId ,
270+ spanId : randomBytes ( 8 ) . toString ( "hex" ) ,
271+ startTime : performance . now ( ) ,
272+ method : req . method || "GET" ,
273+ url : url . pathname ,
274+ workspaceConfig,
275+ }
276+ : undefined ;
248277
249278 const sendErrorAndClose = ( message : string ) => {
250279 try {
@@ -271,7 +300,7 @@ export function handleChatWebSocket(
271300 return validatePayload < ChatRequestInput > ( parsed , ChatRequestSchema ) ;
272301 } ;
273302
274- ws . on ( "message" , async ( data ) => {
303+ const handleMessage = async ( data : RawData ) => {
275304 if ( requestHandled ) {
276305 try {
277306 const size = getRawDataSize ( data ) ;
@@ -450,7 +479,55 @@ export function handleChatWebSocket(
450479 }
451480 }
452481
453- const registeredModel = await getRegisteredModel ( chatReq . model ) ;
482+ const routingSelection = selectIntelligentRouterModel ( {
483+ req,
484+ requestedModel : resolveModelInputForRouting (
485+ chatReq . model ,
486+ defaultProvider ,
487+ defaultModelId ,
488+ ) ,
489+ body : chatReq ,
490+ } ) ;
491+ let registeredModel : RegisteredModel | undefined ;
492+ let lastModelError : unknown ;
493+ let selectedModelError : unknown ;
494+ let selectedModelPolicyViolation : ReturnType <
495+ typeof evaluateModelPolicy
496+ > | null = null ;
497+ for ( const [
498+ index ,
499+ modelInput ,
500+ ] of routingSelection . modelInputs . entries ( ) ) {
501+ try {
502+ const candidateModel = await getRegisteredModel ( modelInput ) ;
503+ const violation = evaluateModelPolicy (
504+ workspaceConfig ?. config ?? getWorkspaceConfigContext ( ) ?. config ,
505+ {
506+ provider : candidateModel . provider ,
507+ modelId : candidateModel . id ,
508+ } ,
509+ ) ;
510+ if ( violation ) {
511+ if ( index === 0 ) selectedModelPolicyViolation = violation ;
512+ continue ;
513+ }
514+ registeredModel = candidateModel ;
515+ break ;
516+ } catch ( error ) {
517+ lastModelError = error ;
518+ if ( index === 0 ) selectedModelError = error ;
519+ }
520+ }
521+ if ( ! registeredModel && selectedModelPolicyViolation ) {
522+ sendErrorAndClose ( selectedModelPolicyViolation . message ) ;
523+ if ( sseLease && releaseSse ) {
524+ releaseSse ( sseLease ) ;
525+ sseLease = null ;
526+ }
527+ return ;
528+ }
529+ if ( ! registeredModel ) throw selectedModelError ?? lastModelError ;
530+ const routeStartedAt = Date . now ( ) ;
454531
455532 const headerApproval = ( ( ) => {
456533 const headerMode = normalizeApprovalMode (
@@ -529,6 +606,10 @@ export function handleChatWebSocket(
529606 : { } ) ,
530607 } ,
531608 ) ;
609+ applyOracleConsultationDirective (
610+ agent ,
611+ routingSelection . decision . oracleConsultation ,
612+ ) ;
532613
533614 const historyMessages = incomingMessages . slice ( 0 , - 1 ) ;
534615 const hydratedHistory = convertComposerMessagesToApp (
@@ -844,6 +925,13 @@ export function handleChatWebSocket(
844925 sessionId : sessionManager . getSessionId ( ) ,
845926 subject,
846927 } ) ;
928+ recordIntelligentRouterChatMetric ( {
929+ taskType : routingSelection . taskType ,
930+ provider : registeredModel . provider ,
931+ model : registeredModel . id ,
932+ startedAt : routeStartedAt ,
933+ message : event . message ,
934+ } ) ;
847935 automaticMemoryExtraction . schedule ( sessionManager . getSessionFile ( ) ) ;
848936 }
849937 sessionManager . updateSnapshot (
@@ -962,5 +1050,14 @@ export function handleChatWebSocket(
9621050 sseLease = null ;
9631051 }
9641052 }
1053+ } ;
1054+
1055+ ws . on ( "message" , ( data ) => {
1056+ if ( websocketRequestContext ) {
1057+ return requestContextStorage . run ( websocketRequestContext , ( ) =>
1058+ handleMessage ( data ) ,
1059+ ) ;
1060+ }
1061+ return handleMessage ( data ) ;
9651062 } ) ;
9661063}
0 commit comments