Skip to content

Commit d8d10b8

Browse files
authored
Merge branch 'main' into feat/subtask-links
2 parents 319f712 + 0e208c4 commit d8d10b8

22 files changed

Lines changed: 965 additions & 62 deletions

‎src/api/providers/__tests__/base-provider.spec.ts‎

Lines changed: 158 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -32,6 +32,11 @@ class TestProvider extends BaseProvider {
3232
public testConvertToolsForOpenAI(tools: any[] | undefined, strict?: boolean): any[] | undefined {
3333
return this.convertToolsForOpenAI(tools, strict)
3434
}
35+
36+
// Expose protected method for testing
37+
public testEnsureCatalogLoaded(load: () => Promise<boolean | void>, signal?: AbortSignal): Promise<void> {
38+
return this.ensureCatalogLoaded(load, signal)
39+
}
3540
}
3641

3742
describe("BaseProvider", () => {
@@ -527,4 +532,157 @@ describe("BaseProvider", () => {
527532
expect(result?.[0]).toEqual(tools[0])
528533
})
529534
})
535+
536+
describe("ensureCatalogLoaded", () => {
537+
it("loads once and memoizes a successful load", async () => {
538+
const load = vi.fn().mockResolvedValue(undefined)
539+
540+
await provider.testEnsureCatalogLoaded(load)
541+
await provider.testEnsureCatalogLoaded(load)
542+
543+
expect(load).toHaveBeenCalledTimes(1)
544+
})
545+
546+
it("single-flights concurrent callers", async () => {
547+
let resolveLoad: () => void = () => {}
548+
const load = vi.fn(
549+
() =>
550+
new Promise<void>((resolve) => {
551+
resolveLoad = resolve
552+
}),
553+
)
554+
555+
const first = provider.testEnsureCatalogLoaded(load)
556+
const second = provider.testEnsureCatalogLoaded(load)
557+
558+
expect(load).toHaveBeenCalledTimes(1)
559+
560+
resolveLoad()
561+
await Promise.all([first, second])
562+
})
563+
564+
it("retries after a failed load", async () => {
565+
const load = vi.fn().mockRejectedValueOnce(new Error("catalog unavailable")).mockResolvedValue(undefined)
566+
567+
await expect(provider.testEnsureCatalogLoaded(load)).rejects.toThrow("catalog unavailable")
568+
await expect(provider.testEnsureCatalogLoaded(load)).resolves.toBeUndefined()
569+
570+
expect(load).toHaveBeenCalledTimes(2)
571+
})
572+
573+
it("rejects a pre-aborted caller without starting a load", async () => {
574+
const controller = new AbortController()
575+
controller.abort(new Error("caller gave up"))
576+
const load = vi.fn().mockResolvedValue(undefined)
577+
578+
await expect(provider.testEnsureCatalogLoaded(load, controller.signal)).rejects.toThrow("caller gave up")
579+
expect(load).not.toHaveBeenCalled()
580+
})
581+
582+
it("detaches only the aborted waiter while the shared load still settles", async () => {
583+
let resolveLoad: () => void = () => {}
584+
const load = vi.fn(
585+
() =>
586+
new Promise<void>((resolve) => {
587+
resolveLoad = resolve
588+
}),
589+
)
590+
591+
const controller = new AbortController()
592+
const wait = provider.testEnsureCatalogLoaded(load, controller.signal)
593+
controller.abort(new Error("caller gave up"))
594+
595+
await expect(wait).rejects.toThrow("caller gave up")
596+
597+
resolveLoad()
598+
await new Promise((resolve) => setTimeout(resolve, 0))
599+
600+
// The shared load completed, so a later caller is a memoized no-op.
601+
await expect(provider.testEnsureCatalogLoaded(load)).resolves.toBeUndefined()
602+
expect(load).toHaveBeenCalledTimes(1)
603+
})
604+
605+
it("does not forward the caller signal to the shared load", async () => {
606+
const controller = new AbortController()
607+
const load = vi.fn().mockResolvedValue(undefined)
608+
609+
await provider.testEnsureCatalogLoaded(load, controller.signal)
610+
611+
expect(load).toHaveBeenCalledWith()
612+
})
613+
614+
it("propagates a shared-load failure to a signal-carrying waiter", async () => {
615+
const controller = new AbortController()
616+
const load = vi.fn().mockRejectedValueOnce(new Error("catalog unavailable")).mockResolvedValue(undefined)
617+
618+
// The waiter is watching a signal that never aborts, so the rejection can
619+
// only come from the shared fetch failing -- the branch that must reject
620+
// the waiter rather than hang until the caller's signal fires.
621+
await expect(provider.testEnsureCatalogLoaded(load, controller.signal)).rejects.toThrow(
622+
"catalog unavailable",
623+
)
624+
625+
// The failed load cleared the shared promise, so the next caller retries.
626+
await expect(provider.testEnsureCatalogLoaded(load)).resolves.toBeUndefined()
627+
expect(load).toHaveBeenCalledTimes(2)
628+
})
629+
630+
it("removes the exact abort listener it registered after a successful load", async () => {
631+
const controller = new AbortController()
632+
const addListener = vi.spyOn(controller.signal, "addEventListener")
633+
const removeListener = vi.spyOn(controller.signal, "removeEventListener")
634+
const load = vi.fn().mockResolvedValue(undefined)
635+
636+
await provider.testEnsureCatalogLoaded(load, controller.signal)
637+
638+
// A long-lived signal must not accumulate listeners: the removal has to
639+
// target the same function reference that was registered.
640+
expect(addListener).toHaveBeenCalledWith("abort", expect.any(Function), { once: true })
641+
const registered = addListener.mock.calls[0]?.[1]
642+
expect(removeListener).toHaveBeenCalledWith("abort", registered)
643+
expect(removeListener).toHaveBeenCalledTimes(1)
644+
})
645+
646+
it("removes the exact abort listener it registered after a rejected load", async () => {
647+
const controller = new AbortController()
648+
const addListener = vi.spyOn(controller.signal, "addEventListener")
649+
const removeListener = vi.spyOn(controller.signal, "removeEventListener")
650+
const load = vi.fn().mockRejectedValueOnce(new Error("catalog unavailable"))
651+
652+
await expect(provider.testEnsureCatalogLoaded(load, controller.signal)).rejects.toThrow(
653+
"catalog unavailable",
654+
)
655+
656+
expect(addListener).toHaveBeenCalledWith("abort", expect.any(Function), { once: true })
657+
const registered = addListener.mock.calls[0]?.[1]
658+
expect(removeListener).toHaveBeenCalledWith("abort", registered)
659+
expect(removeListener).toHaveBeenCalledTimes(1)
660+
})
661+
662+
it("does not memoize a load that reports an empty catalog", async () => {
663+
const load = vi.fn().mockResolvedValueOnce(false).mockResolvedValueOnce(true)
664+
665+
await provider.testEnsureCatalogLoaded(load)
666+
// The empty result must stay retryable, so the next caller loads again.
667+
await provider.testEnsureCatalogLoaded(load)
668+
expect(load).toHaveBeenCalledTimes(2)
669+
670+
// The second load produced a catalog, so it is memoized.
671+
await provider.testEnsureCatalogLoaded(load)
672+
expect(load).toHaveBeenCalledTimes(2)
673+
})
674+
675+
it("memoizes a load that returns true or void", async () => {
676+
const returningTrue = vi.fn().mockResolvedValue(true)
677+
await provider.testEnsureCatalogLoaded(returningTrue)
678+
await provider.testEnsureCatalogLoaded(returningTrue)
679+
expect(returningTrue).toHaveBeenCalledTimes(1)
680+
681+
const other = new TestProvider()
682+
const returningVoid = vi.fn().mockResolvedValue(undefined)
683+
await other.testEnsureCatalogLoaded(returningVoid)
684+
await other.testEnsureCatalogLoaded(returningVoid)
685+
expect(returningVoid).toHaveBeenCalledTimes(1)
686+
})
687+
})
530688
})

‎src/api/providers/__tests__/lmstudio.spec.ts‎

Lines changed: 42 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -62,9 +62,21 @@ vi.mock("openai", () => {
6262

6363
import type { Anthropic } from "@anthropic-ai/sdk"
6464

65+
import { providerIdentifiers } from "@roo-code/types"
66+
6567
import { LmStudioHandler } from "../lm-studio"
6668
import type { ApiHandlerOptions } from "../../../shared/api"
6769

70+
const { mockGetModels, mockGetModelsFromCache } = vi.hoisted(() => ({
71+
mockGetModels: vi.fn(),
72+
mockGetModelsFromCache: vi.fn(),
73+
}))
74+
75+
vi.mock("../fetchers/modelCache", () => ({
76+
getModels: mockGetModels,
77+
getModelsFromCache: mockGetModelsFromCache,
78+
}))
79+
6880
describe("LmStudioHandler", () => {
6981
let handler: LmStudioHandler
7082
let mockOptions: ApiHandlerOptions
@@ -77,6 +89,9 @@ describe("LmStudioHandler", () => {
7789
}
7890
handler = new LmStudioHandler(mockOptions)
7991
mockCreate.mockClear()
92+
mockGetModels.mockReset()
93+
mockGetModelsFromCache.mockReset()
94+
mockGetModelsFromCache.mockReturnValue(undefined)
8095
})
8196

8297
describe("constructor", () => {
@@ -258,4 +273,31 @@ describe("LmStudioHandler", () => {
258273
expect(modelInfo.info.contextWindow).toBe(128_000)
259274
})
260275
})
276+
277+
describe("ensureModelFetched", () => {
278+
it("returns default info before the pre-warm and cached info after", async () => {
279+
// The pre-warm populates the shared cache that getModel() reads, so the
280+
// cache starts empty and is filled by the mocked getModels() call.
281+
let cacheContents: Record<string, { maxTokens: number; contextWindow: number }> | undefined
282+
mockGetModelsFromCache.mockImplementation(() => cacheContents)
283+
mockGetModels.mockImplementation(async () => {
284+
cacheContents = { "local-model": { maxTokens: 4096, contextWindow: 32_768 } }
285+
return cacheContents
286+
})
287+
288+
expect(handler.getModel().info.contextWindow).toBe(128_000)
289+
290+
await handler.ensureModelFetched()
291+
292+
expect(mockGetModels).toHaveBeenCalledWith(
293+
{
294+
provider: providerIdentifiers.lmstudio,
295+
baseUrl: "http://localhost:1234",
296+
},
297+
handler,
298+
)
299+
expect(handler.getModel().info.contextWindow).toBe(32_768)
300+
expect(handler.getModel().info.maxTokens).toBe(4096)
301+
})
302+
})
261303
})

0 commit comments

Comments
 (0)