diff --git a/.changeset/hot-cheetahs-call.md b/.changeset/hot-cheetahs-call.md new file mode 100644 index 0000000..7aad9ad --- /dev/null +++ b/.changeset/hot-cheetahs-call.md @@ -0,0 +1,5 @@ +--- +'@callstack/byorg-core': patch +--- + +Fixed missing references diff --git a/packages/core/src/__tests__/application.test.ts b/packages/core/src/__tests__/application.test.ts index 64916b2..4301ee7 100644 --- a/packages/core/src/__tests__/application.test.ts +++ b/packages/core/src/__tests__/application.test.ts @@ -137,3 +137,47 @@ test('uses chat model from context', async () => { expect(altModel.calls.length).toBe(1); expect(baseModel.calls.length).toBe(0); }); + +test('adds unique and returns references', async () => { + const baseModel = createMockChatModel({ delay: 0, seed: 3 }); + const refFunc = vitest.fn(); + + const addReferenceMiddlewareMock = vitest.fn( + async (context: RequestContext, next: () => Promise) => { + context.references.addReferences([{ title: 'Test', url: 'test.com' }]); + return await next(); + }, + ); + + const getReferenceMiddlewareMock = vitest.fn( + async (context: RequestContext, next: () => Promise) => { + const response = await next(); + refFunc(context.references.getReferences()); + return response; + }, + ); + + const app = createApp({ + chatModel: baseModel, + plugins: [ + { + name: 'add', + middleware: addReferenceMiddlewareMock, + }, + { + name: 'get', + middleware: getReferenceMiddlewareMock, + }, + ], + }); + + expect(refFunc).not.toBeCalled(); + await app.processMessages(messages); + expect(addReferenceMiddlewareMock).toBeCalledTimes(1); + expect(getReferenceMiddlewareMock).toBeCalledTimes(1); + expect(refFunc).toBeCalledWith([{ title: 'Test', url: 'test.com' }]); + await app.processMessages(messages); + expect(addReferenceMiddlewareMock).toBeCalledTimes(2); + expect(getReferenceMiddlewareMock).toBeCalledTimes(2); + expect(refFunc).toBeCalledWith([{ title: 'Test', url: 'test.com' }]); +}); diff --git a/packages/core/src/references.ts b/packages/core/src/references.ts index 4fc0913..b5aad33 100644 --- a/packages/core/src/references.ts +++ b/packages/core/src/references.ts @@ -5,8 +5,8 @@ export const getReferenceStorage = (): ReferenceStorage => { const getKey = (ref: DocumentReference) => ref.url; return { - addReferences(references: DocumentReference[]): void { - for (const ref of references) { + addReferences(referencesToAdd: DocumentReference[]): void { + for (const ref of referencesToAdd) { const refKey = getKey(ref); if (!referencesKeys.has(refKey)) { referencesKeys.add(refKey);