11import axios from "axios"
22import { vi , describe , it , expect , beforeEach } from "vitest"
3- import { LMStudioClient , LLM , LLMInstanceInfo } from "@lmstudio/sdk" // LLMInfo is a type
3+ import { LMStudioClient , LLM , LLMInstanceInfo , LLMInfo } from "@lmstudio/sdk"
44import { getLMStudioModels , parseLMStudioModel } from "../lmstudio"
55import { ModelInfo , lMStudioDefaultModelInfo } from "@roo-code/types" // ModelInfo is a type
66
@@ -11,12 +11,16 @@ const mockedAxios = axios as any
1111// Mock @lmstudio /sdk
1212const mockGetModelInfo = vi . fn ( )
1313const mockListLoaded = vi . fn ( )
14+ const mockListDownloadedModels = vi . fn ( )
1415vi . mock ( "@lmstudio/sdk" , ( ) => {
1516 return {
1617 LMStudioClient : vi . fn ( ) . mockImplementation ( ( ) => ( {
1718 llm : {
1819 listLoaded : mockListLoaded ,
1920 } ,
21+ system : {
22+ listDownloadedModels : mockListDownloadedModels ,
23+ } ,
2024 } ) ) ,
2125 }
2226} )
@@ -28,6 +32,7 @@ describe("LMStudio Fetcher", () => {
2832 MockedLMStudioClientConstructor . mockClear ( )
2933 mockListLoaded . mockClear ( )
3034 mockGetModelInfo . mockClear ( )
35+ mockListDownloadedModels . mockClear ( )
3136 } )
3237
3338 describe ( "parseLMStudioModel" , ( ) => {
@@ -88,8 +93,40 @@ describe("LMStudio Fetcher", () => {
8893 trainedForToolUse : false , // Added
8994 }
9095
91- it ( "should fetch and parse models successfully" , async ( ) => {
96+ it ( "should fetch downloaded models using system.listDownloadedModels" , async ( ) => {
97+ const mockLLMInfo : LLMInfo = {
98+ type : "llm" as const ,
99+ modelKey : "mistralai/devstral-small-2505" ,
100+ format : "safetensors" ,
101+ displayName : "Devstral Small 2505" ,
102+ path : "mistralai/devstral-small-2505" ,
103+ sizeBytes : 13277565112 ,
104+ architecture : "mistral" ,
105+ vision : false ,
106+ trainedForToolUse : false ,
107+ maxContextLength : 131072 ,
108+ }
109+
110+ mockedAxios . get . mockResolvedValueOnce ( { data : { status : "ok" } } )
111+ mockListDownloadedModels . mockResolvedValueOnce ( [ mockLLMInfo ] )
112+
113+ const result = await getLMStudioModels ( baseUrl )
114+
115+ expect ( mockedAxios . get ) . toHaveBeenCalledTimes ( 1 )
116+ expect ( mockedAxios . get ) . toHaveBeenCalledWith ( `${ baseUrl } /v1/models` )
117+ expect ( MockedLMStudioClientConstructor ) . toHaveBeenCalledTimes ( 1 )
118+ expect ( MockedLMStudioClientConstructor ) . toHaveBeenCalledWith ( { baseUrl : lmsUrl } )
119+ expect ( mockListDownloadedModels ) . toHaveBeenCalledTimes ( 1 )
120+ expect ( mockListDownloadedModels ) . toHaveBeenCalledWith ( "llm" )
121+ expect ( mockListLoaded ) . not . toHaveBeenCalled ( )
122+
123+ const expectedParsedModel = parseLMStudioModel ( mockLLMInfo )
124+ expect ( result ) . toEqual ( { [ mockLLMInfo . path ] : expectedParsedModel } )
125+ } )
126+
127+ it ( "should fall back to listLoaded when listDownloadedModels fails" , async ( ) => {
92128 mockedAxios . get . mockResolvedValueOnce ( { data : { status : "ok" } } )
129+ mockListDownloadedModels . mockRejectedValueOnce ( new Error ( "Method not available" ) )
93130 mockListLoaded . mockResolvedValueOnce ( [ { getModelInfo : mockGetModelInfo } ] )
94131 mockGetModelInfo . mockResolvedValueOnce ( mockRawModel )
95132
@@ -99,6 +136,7 @@ describe("LMStudio Fetcher", () => {
99136 expect ( mockedAxios . get ) . toHaveBeenCalledWith ( `${ baseUrl } /v1/models` )
100137 expect ( MockedLMStudioClientConstructor ) . toHaveBeenCalledTimes ( 1 )
101138 expect ( MockedLMStudioClientConstructor ) . toHaveBeenCalledWith ( { baseUrl : lmsUrl } )
139+ expect ( mockListDownloadedModels ) . toHaveBeenCalledTimes ( 1 )
102140 expect ( mockListLoaded ) . toHaveBeenCalledTimes ( 1 )
103141
104142 const expectedParsedModel = parseLMStudioModel ( mockRawModel )
0 commit comments