@@ -137,3 +137,47 @@ test('uses chat model from context', async () => {
137
137
expect ( altModel . calls . length ) . toBe ( 1 ) ;
138
138
expect ( baseModel . calls . length ) . toBe ( 0 ) ;
139
139
} ) ;
140
+
141
+ test ( 'adds unique and returns references' , async ( ) => {
142
+ const baseModel = createMockChatModel ( { delay : 0 , seed : 3 } ) ;
143
+ const refFunc = vitest . fn ( ) ;
144
+
145
+ const addReferenceMiddlewareMock = vitest . fn (
146
+ async ( context : RequestContext , next : ( ) => Promise < MessageResponse > ) => {
147
+ context . references . addReferences ( [ { title : 'Test' , url : 'test.com' } ] ) ;
148
+ return await next ( ) ;
149
+ } ,
150
+ ) ;
151
+
152
+ const getReferenceMiddlewareMock = vitest . fn (
153
+ async ( context : RequestContext , next : ( ) => Promise < MessageResponse > ) => {
154
+ const response = await next ( ) ;
155
+ refFunc ( context . references . getReferences ( ) ) ;
156
+ return response ;
157
+ } ,
158
+ ) ;
159
+
160
+ const app = createApp ( {
161
+ chatModel : baseModel ,
162
+ plugins : [
163
+ {
164
+ name : 'add' ,
165
+ middleware : addReferenceMiddlewareMock ,
166
+ } ,
167
+ {
168
+ name : 'get' ,
169
+ middleware : getReferenceMiddlewareMock ,
170
+ } ,
171
+ ] ,
172
+ } ) ;
173
+
174
+ expect ( refFunc ) . not . toBeCalled ( ) ;
175
+ await app . processMessages ( messages ) ;
176
+ expect ( addReferenceMiddlewareMock ) . toBeCalledTimes ( 1 ) ;
177
+ expect ( getReferenceMiddlewareMock ) . toBeCalledTimes ( 1 ) ;
178
+ expect ( refFunc ) . toBeCalledWith ( [ { title : 'Test' , url : 'test.com' } ] ) ;
179
+ await app . processMessages ( messages ) ;
180
+ expect ( addReferenceMiddlewareMock ) . toBeCalledTimes ( 2 ) ;
181
+ expect ( getReferenceMiddlewareMock ) . toBeCalledTimes ( 2 ) ;
182
+ expect ( refFunc ) . toBeCalledWith ( [ { title : 'Test' , url : 'test.com' } ] ) ;
183
+ } ) ;
0 commit comments