@@ -42,6 +42,7 @@ import {
4242 STREAM_LIMIT_EPOCH_KEY ,
4343 StreamLimitExceededError ,
4444} from '@/llm/streamLimits' ;
45+ import { ProviderTextProtectionError } from '@/protection/providerText' ;
4546import { GraphEvents , Providers , StepTypes } from '@/common' ;
4647import * as eagerArgs from '@/tools/eagerEventExecution' ;
4748import { ChatModelStreamHandler } from '@/stream' ;
@@ -664,14 +665,13 @@ describe('eager args divergence (LibreChat#14371)', () => {
664665 expect ( toolExecuteCalls ) . toHaveLength ( 0 ) ;
665666 } ) ;
666667
667- it ( 'does not prestart eager tools when the breaker trips during tool-call handling' , async ( ) => {
668+ it . each ( [
669+ { label : 'stream limit' , trip : new StreamLimitExceededError ( { kind : 'tool_call_args' , limit : 10 , observed : 11 , toolName : 'db_query' } ) } ,
670+ { label : 'blocked' , trip : new ProviderTextProtectionError ( 'blocked' ) } ,
671+ { label : 'timeout' , trip : new ProviderTextProtectionError ( 'timeout' ) } ,
672+ { label : 'cancelled' , trip : new ProviderTextProtectionError ( 'cancelled' ) } ,
673+ ] ) ( 'does not prestart eager tools after an awaited $label trip' , async ( { trip } ) => {
668674 const graph = createGraph ( ) ;
669- const trip = new StreamLimitExceededError ( {
670- kind : 'tool_call_args' ,
671- limit : 10 ,
672- observed : 11 ,
673- toolName : 'db_query' ,
674- } ) ;
675675 const originalDispatch = graph . dispatchRunStep ;
676676 graph . dispatchRunStep = ( async (
677677 stepKey : string ,
@@ -712,6 +712,70 @@ describe('eager args divergence (LibreChat#14371)', () => {
712712 expect ( toolExecuteCalls ) . toHaveLength ( 0 ) ;
713713 } ) ;
714714
715+ it . each ( [ 'blocked' , 'timeout' ] as const ) ( 'stops streamed seals when protection trips during awaited publication (%s)' , async ( code ) => {
716+ const graph = createGraph ( ) ;
717+ const trip = new ProviderTextProtectionError ( code ) ;
718+ const originalDispatch = graph . dispatchRunStep . bind ( graph ) ;
719+ graph . dispatchRunStep = async ( key , details , metadata ) : Promise < string > => {
720+ const id = await originalDispatch ( key , details , metadata ) ;
721+ graph . breakerAbort . abort ( trip ) ;
722+ return id ;
723+ } ;
724+ const { toolExecuteCalls } = installToolExecuteResponder ( ) ;
725+ const chunk = new AIMessageChunk ( { content : '' , tool_call_chunks : [ { index : 0 , id : 'sealed-call' , name : 'db_query' , args : '{"sql":"SELECT 1;"}' } ] , response_metadata : {
726+ [ STREAMED_TOOL_CALL_ADAPTER_METADATA_KEY ] : BEDROCK_CONVERSE_STREAMED_TOOL_CALL_ADAPTER ,
727+ [ STREAMED_TOOL_CALL_SEAL_METADATA_KEY ] : { kind : 'single' , index : 0 } ,
728+ } } ) ;
729+ chunk . tool_calls = [ ] ;
730+ await expect ( new ChatModelStreamHandler ( ) . handle ( GraphEvents . CHAT_MODEL_STREAM , { chunk } , { langgraph_node : 'agent' } , graph ) ) . rejects . toBe ( trip ) ;
731+ expect ( toolExecuteCalls ) . toHaveLength ( 0 ) ;
732+ expect ( graph . eagerEventToolExecutions . size ) . toBe ( 0 ) ;
733+ expect ( graph . eagerEventToolCallChunks . size ) . toBe ( 1 ) ;
734+ } ) ;
735+
736+ it ( 'drops already queued stream events after a protection trip' , async ( ) => {
737+ const graph = createGraph ( ) ;
738+ const trip = new ProviderTextProtectionError ( 'blocked' ) ;
739+ graph . breakerAbort . abort ( trip ) ;
740+ const { toolExecuteCalls } = installToolExecuteResponder ( ) ;
741+ await expect ( new ChatModelStreamHandler ( ) . handle ( GraphEvents . CHAT_MODEL_STREAM , {
742+ chunk : new AIMessageChunk ( { content : '' , tool_calls : [ { id : 'queued-call' , name : 'db_query' , args : { sql : 'SELECT 1;' } , type : 'tool_call' } ] , response_metadata : { finish_reason : 'tool_calls' } } ) ,
743+ } , { langgraph_node : 'agent' } , graph ) ) . rejects . toBe ( trip ) ;
744+ expect ( graph . dispatchRunStep ) . not . toHaveBeenCalled ( ) ;
745+ expect ( toolExecuteCalls ) . toHaveLength ( 0 ) ;
746+ } ) ;
747+
748+ it . each ( [
749+ { signalSource : 'constructor' , shape : 'parsed' } ,
750+ { signalSource : 'constructor' , shape : 'seal' } ,
751+ { signalSource : 'call' , shape : 'parsed' } ,
752+ { signalSource : 'call' , shape : 'seal' } ,
753+ ] ) ( 'honors inherited $signalSource trips after awaited $shape handling' , async ( { signalSource, shape } ) => {
754+ const parent = new AbortController ( ) ;
755+ const graph = createGraph ( signalSource === 'constructor' ? { signal : parent . signal } : { } ) ;
756+ if ( signalSource === 'call' ) graph . config = { ...graph . config , signal : parent . signal } ;
757+ const trip = new ProviderTextProtectionError ( 'blocked' ) ;
758+ const originalDispatch = graph . dispatchRunStep . bind ( graph ) ;
759+ graph . dispatchRunStep = async ( key , details , metadata ) : Promise < string > => {
760+ const id = await originalDispatch ( key , details , metadata ) ;
761+ parent . abort ( trip ) ;
762+ graph . config = { ...graph . config , signal : new AbortController ( ) . signal } ;
763+ return id ;
764+ } ;
765+ const { toolExecuteCalls } = installToolExecuteResponder ( ) ;
766+ const chunk = shape === 'parsed'
767+ ? new AIMessageChunk ( { content : '' , tool_calls : [ { id : 'inherited-call' , name : 'db_query' , args : { sql : 'SELECT 1;' } , type : 'tool_call' } ] , response_metadata : { finish_reason : 'tool_calls' } } )
768+ : new AIMessageChunk ( { content : '' , tool_call_chunks : [ { index : 0 , id : 'inherited-call' , name : 'db_query' , args : '{"sql":"SELECT 1;"}' } ] , response_metadata : {
769+ [ STREAMED_TOOL_CALL_ADAPTER_METADATA_KEY ] : BEDROCK_CONVERSE_STREAMED_TOOL_CALL_ADAPTER ,
770+ [ STREAMED_TOOL_CALL_SEAL_METADATA_KEY ] : { kind : 'single' , index : 0 } ,
771+ } } ) ;
772+ if ( shape === 'seal' ) chunk . tool_calls = [ ] ;
773+ await expect ( new ChatModelStreamHandler ( ) . handle ( GraphEvents . CHAT_MODEL_STREAM , { chunk } , { langgraph_node : 'agent' } , graph ) ) . rejects . toBe ( trip ) ;
774+ expect ( graph . breakerAbort . signal . aborted ) . toBe ( false ) ;
775+ expect ( toolExecuteCalls ) . toHaveLength ( 0 ) ;
776+ expect ( graph . eagerEventToolExecutions . size ) . toBe ( 0 ) ;
777+ } ) ;
778+
715779 it ( 'sends a breaker-composed abort signal with eager prestart requests' , async ( ) => {
716780 const graph = createGraph ( ) ;
717781 const { toolExecuteCalls } = installToolExecuteResponder ( ) ;
0 commit comments