@@ -145,38 +145,13 @@ func initLicenseValidation(ctx context.Context, config *config.Config) (*License
145
145
146
146
// InitToolsets initializes and returns the toolset groups
147
147
func InitToolsets (ctx context.Context , config * config.Config ) (* toolsets.ToolsetGroup , error ) {
148
- // Create a toolset group
148
+ // Create a toolset tsg
149
149
tsg := toolsets .NewToolsetGroup (config .ReadOnly )
150
150
151
151
// Initialize license validation if enabled
152
- var licenseInfo * LicenseInfo
153
- var err error
154
152
if config .EnableLicense {
155
- licenseInfo , err = initLicenseValidation (ctx , config )
156
- if err != nil {
157
- slog .Warn ("License validation failed" , "error" , err )
158
- return nil , fmt .Errorf ("failed to fetch license details, error: %w" , err )
159
- }
160
-
161
- // Create a module registry
162
- registry := modules .NewModuleRegistry (config , tsg )
163
-
164
- // Get all modules that are enabled based on configuration
165
- configEnabledModules := registry .GetEnabledModules ()
166
- // Get enabled modules based on configuration and license
167
-
168
- enabledModules := getEnabledModules (configEnabledModules , licenseInfo )
169
- // Register toolsets for enabled modules
170
- for _ , module := range enabledModules {
171
- slog .Info ("registering toolsets for" , "modules" , module .ID ())
172
- if err := module .RegisterToolsets (); err != nil {
173
- return nil , fmt .Errorf ("failed to register toolsets for module %s: %w" , module .ID (), err )
174
- }
175
-
176
- // Enable toolsets for this module
177
- if err := module .EnableToolsets (tsg ); err != nil {
178
- return nil , fmt .Errorf ("failed to enable toolsets for module %s: %w" , module .ID (), err )
179
- }
153
+ if err := initModuleBasedToolsets (ctx , config , tsg ); err != nil {
154
+ return nil , err
180
155
}
181
156
} else {
182
157
// License validation is disabled, use legacy toolset registration
@@ -185,9 +160,57 @@ func InitToolsets(ctx context.Context, config *config.Config) (*toolsets.Toolset
185
160
}
186
161
}
187
162
163
+ // Register all toolsets with the main tracker for tool-to-toolset mapping
164
+ registerAllToolsetsWithTracker (tsg )
165
+
188
166
return tsg , nil
189
167
}
190
168
169
+ // registerAllToolsetsWithTracker registers all toolsets in the group with the main tracker
170
+ // This ensures that findToolGroup() can find which toolset a tool belongs to
171
+ func registerAllToolsetsWithTracker (group * toolsets.ToolsetGroup ) {
172
+ tracker := toolsets .GetMainToolTracker ()
173
+
174
+ // Register each toolset with the tracker
175
+ for _ , toolset := range group .Toolsets {
176
+ if err := tracker .RegisterToolGroup (toolset ); err != nil {
177
+ slog .Warn ("Failed to register toolset with tracker" , "toolset" , toolset .Name , "error" , err )
178
+ }
179
+ }
180
+
181
+ slog .Info ("Registered toolsets with tracker" , "count" , len (group .Toolsets ))
182
+ }
183
+
184
+ // New function that handles the module-based initialization
185
+ func initModuleBasedToolsets (ctx context.Context , config * config.Config , tsg * toolsets.ToolsetGroup ) error {
186
+ // Get license info
187
+ licenseInfo , err := initLicenseValidation (ctx , config )
188
+ if err != nil {
189
+ slog .Warn ("License validation failed" , "error" , err )
190
+ return fmt .Errorf ("failed to fetch license details, error: %w" , err )
191
+ }
192
+
193
+ // Create module registry and get enabled modules
194
+ registry := modules .NewModuleRegistry (config , tsg )
195
+ configEnabledModules := registry .GetEnabledModules ()
196
+ enabledModules := getEnabledModules (configEnabledModules , licenseInfo )
197
+
198
+ // Register and enable toolsets for each module
199
+ for _ , module := range enabledModules {
200
+ slog .Info ("registering toolsets for" , "modules" , module .ID ())
201
+ if err := module .RegisterToolsets (); err != nil {
202
+ return fmt .Errorf ("failed to register toolsets for module %s: %w" , module .ID (), err )
203
+ }
204
+
205
+ if err := module .EnableToolsets (tsg ); err != nil {
206
+ return fmt .Errorf ("failed to enable toolsets for module %s: %w" , module .ID (), err )
207
+ }
208
+ }
209
+
210
+ return nil
211
+ }
212
+
213
+ // initLegacyToolsets initializes toolsets using the legacy approach (without modules)
191
214
func initLegacyToolsets (config * config.Config , tsg * toolsets.ToolsetGroup ) error {
192
215
// Check if specific toolsets are enabled
193
216
if len (config .Toolsets ) == 0 {
0 commit comments