@@ -10,6 +10,7 @@ import {
1010import { tool } from "@langchain/core/tools" ;
1111import { HumanMessage , AIMessage } from "@langchain/core/messages" ;
1212import { BaseChatModel } from "@langchain/core/language_models/chat_models" ;
13+ import { FakeListChatModel } from "@langchain/core/utils/testing" ;
1314
1415import { llmToolSelectorMiddleware } from "../llmToolSelector.js" ;
1516import { createAgent } from "../../index.js" ;
@@ -417,3 +418,83 @@ describe("llmToolSelectorMiddleware", () => {
417418 expect ( firstCall [ 1 ] . content ) . toContain ( "Second message" ) ;
418419 } ) ;
419420} ) ;
421+
422+ describe ( "llmToolSelectorMiddleware – streaming isolation" , ( ) => {
423+ const getWeather = tool (
424+ ( { location } : { location : string } ) =>
425+ `Weather in ${ location } : Sunny, 72°F` ,
426+ {
427+ name : "get_weather" ,
428+ description : "Get current weather for a location" ,
429+ schema : z . object ( {
430+ location : z . string ( ) . describe ( "City name" ) ,
431+ } ) ,
432+ }
433+ ) ;
434+
435+ const searchDatabase = tool (
436+ ( { customerId } : { customerId : string } ) =>
437+ `Customer ${ customerId } : Premium account` ,
438+ {
439+ name : "search_database" ,
440+ description : "Look up customer information by customer ID" ,
441+ schema : z . object ( {
442+ customerId : z . string ( ) . describe ( "Customer ID" ) ,
443+ } ) ,
444+ }
445+ ) ;
446+
447+ const calculatePrice = tool (
448+ ( { items, discount } : { items : number ; discount : number } ) =>
449+ `Total: $${ ( items * 29.99 * ( 1 - discount / 100 ) ) . toFixed ( 2 ) } ` ,
450+ {
451+ name : "calculate_price" ,
452+ description : "Calculate pricing with discounts" ,
453+ schema : z . object ( {
454+ items : z . number ( ) ,
455+ discount : z . number ( ) ,
456+ } ) ,
457+ }
458+ ) ;
459+
460+ it ( "does not leak tool-selector output into messages stream" , async ( ) => {
461+ // given
462+ const selectorModel = new FakeListChatModel ( {
463+ responses : [ JSON . stringify ( { tools : [ "get_weather" ] } ) ] ,
464+ } ) ;
465+ const agentModel = new FakeListChatModel ( {
466+ responses : [ "The weather in Seoul is sunny and 72°F." ] ,
467+ } ) ;
468+
469+ const middleware = llmToolSelectorMiddleware ( {
470+ model : selectorModel ,
471+ maxTools : 1 ,
472+ } ) ;
473+
474+ const agent = createAgent ( {
475+ model : agentModel ,
476+ tools : [ getWeather , searchDatabase , calculatePrice ] ,
477+ middleware : [ middleware ] ,
478+ } ) ;
479+
480+ // when
481+ const stream = await agent . stream (
482+ { messages : [ new HumanMessage ( "What's the weather in Seoul?" ) ] } ,
483+ { streamMode : "messages" }
484+ ) ;
485+
486+ const parts : string [ ] = [ ] ;
487+ for await ( const chunk of stream ) {
488+ parts . push ( JSON . stringify ( chunk ) ) ;
489+ }
490+ const serialized = parts . join ( "" ) ;
491+
492+ // then
493+ const hasToolSelectorLeak =
494+ serialized . includes ( '"content":"{\\"tools' ) ||
495+ serialized . includes ( '"content":"tools' ) ||
496+ / " c o n t e n t " : " [ ^ " ] * t o o l s [ ^ " ] * " , " t o o l _ c a l l _ c h u n k s " : \[ \] / . test ( serialized ) ;
497+
498+ expect ( hasToolSelectorLeak ) . toBe ( false ) ;
499+ } ) ;
500+ } ) ;
0 commit comments