33 * Licensed under the MIT License. See License.txt in the project root for license information.
44 *--------------------------------------------------------------------------------------------*/
55
6+ import { CancellationToken } from '../../../../base/common/cancellation.js' ;
67import { Disposable } from '../../../../base/common/lifecycle.js' ;
7- import { autorun , derived , IObservable , observableFromEvent , ObservableMap , observableValue , transaction } from '../../../../base/common/observable.js' ;
8+ import { derived , IObservable , observableFromEvent , ObservableMap } from '../../../../base/common/observable.js' ;
9+ import { isObject } from '../../../../base/common/types.js' ;
10+ import { URI } from '../../../../base/common/uri.js' ;
11+ import { IInstantiationService } from '../../../../platform/instantiation/common/instantiation.js' ;
812import { ObservableMemento , observableMemento } from '../../../../platform/observable/common/observableMemento.js' ;
913import { IStorageService , StorageScope , StorageTarget } from '../../../../platform/storage/common/storage.js' ;
14+ import { IChatMode } from '../common/chatModes.js' ;
1015import { ChatMode } from '../common/constants.js' ;
11- import { ILanguageModelToolsService , IToolData , ToolSet } from '../common/languageModelToolsService.js' ;
16+ import { ILanguageModelToolsService , IToolAndToolSetEnablementMap , IToolData , ToolSet } from '../common/languageModelToolsService.js' ;
17+ import { PromptFileRewriter } from './promptSyntax/promptFileRewriter.js' ;
1218
1319
1420/**
@@ -23,28 +29,24 @@ type StoredData = {
2329 readonly disabledTools ?: readonly string [ ] ;
2430} ;
2531
26- const storedTools = observableMemento < StoredData > ( {
27- defaultValue : { } ,
28- key : 'chat/selectedTools' ,
29- } ) ;
32+ export enum ToolsScope {
33+ Global ,
34+ Session ,
35+ Mode
36+ }
3037
3138export class ChatSelectedTools extends Disposable {
3239
33- private readonly _selectedTools : ObservableMemento < StoredData > ;
40+ private readonly _selectedTools : ObservableMemento < IToolAndToolSetEnablementMap > ;
3441
35- private readonly _sessionSelectedTools = observableValue < StoredData > ( this , { } ) ;
42+ private readonly _sessionStates = new ObservableMap < string , IToolAndToolSetEnablementMap | undefined > ( ) ;
3643
3744 private readonly _allTools : IObservable < Readonly < IToolData > [ ] > ;
3845
39- /**
40- * All tools and tool sets with their enabled state.
41- */
42- readonly entriesMap = new ObservableMap < ToolSet | IToolData , boolean > ( ) ;
43-
4446 /**
4547 * All enabled tools and tool sets.
4648 */
47- readonly entries : IObservable < ReadonlySet < IToolData | ToolSet > > = this . entriesMap . observable . map ( function ( value ) {
49+ readonly entries : IObservable < ReadonlySet < IToolData | ToolSet > > = this . entriesMap . map ( function ( value ) {
4850 const result = new Set < IToolData | ToolSet > ( ) ;
4951 for ( const [ item , enabled ] of value ) {
5052 if ( enabled ) {
@@ -55,103 +57,139 @@ export class ChatSelectedTools extends Disposable {
5557 } ) ;
5658
5759 constructor (
58- mode : IObservable < ChatMode > ,
60+ private readonly _mode : IObservable < IChatMode > ,
5961 @ILanguageModelToolsService private readonly _toolsService : ILanguageModelToolsService ,
60- @IStorageService storageService : IStorageService ,
62+ @IStorageService _storageService : IStorageService ,
63+ @IInstantiationService private readonly _instantiationService : IInstantiationService ,
6164 ) {
6265 super ( ) ;
6366
64- this . _selectedTools = this . _store . add ( storedTools ( StorageScope . WORKSPACE , StorageTarget . MACHINE , storageService ) ) ;
65-
66- this . _allTools = observableFromEvent ( _toolsService . onDidChangeTools , ( ) => Array . from ( _toolsService . getTools ( ) ) ) ;
67-
68- const disabledDataObs = derived ( r => {
69- const globalData = this . _selectedTools . read ( r ) ;
70- const sessionData = this . _sessionSelectedTools . read ( r ) ;
71-
72- const toolSetIds = new Set < string > ( ) ;
73- const toolIds = new Set < string > ( ) ;
74-
75- for ( const data of [ globalData , sessionData ] ) {
76- if ( data . disabledToolSets ) {
77- for ( const id of data . disabledToolSets ) {
78- toolSetIds . add ( id ) ;
67+ const storedTools = observableMemento < IToolAndToolSetEnablementMap > ( {
68+ defaultValue : new Map ( ) ,
69+ toStorage : ( value ) => {
70+ const data = {
71+ disabledToolSets : [ ] as string [ ] ,
72+ disabledTools : [ ] as string [ ] ,
73+ } ;
74+ for ( const [ item , enabled ] of value ) {
75+ if ( ! enabled ) {
76+ if ( item instanceof ToolSet ) {
77+ data . disabledToolSets . push ( item . id ) ;
78+ } else {
79+ data . disabledTools . push ( item . id ) ;
80+ }
81+ }
82+ }
83+ return JSON . stringify ( data ) ;
84+ } ,
85+ fromStorage : ( value ) => {
86+ const obj = JSON . parse ( value ) as StoredData ;
87+ const map = new Map < IToolData | ToolSet , boolean > ( ) ;
88+ if ( ! obj || ! isObject ( obj ) ) {
89+ return map ;
90+ }
91+ if ( Array . isArray ( obj . disabledToolSets ) ) {
92+ for ( const toolSetId of obj . disabledToolSets ) {
93+ const toolset = this . _toolsService . getToolSet ( toolSetId ) ;
94+ if ( toolset ) {
95+ map . set ( toolset , false ) ;
96+ }
7997 }
8098 }
81- if ( data . disabledTools ) {
82- for ( const id of data . disabledTools ) {
83- toolIds . add ( id ) ;
99+ if ( Array . isArray ( obj . disabledTools ) ) {
100+ for ( const toolId of obj . disabledTools ) {
101+ const tool = this . _toolsService . getTool ( toolId ) ;
102+ if ( tool ) {
103+ map . set ( tool , false ) ;
104+ }
84105 }
85106 }
86- }
87-
88- if ( toolSetIds . size === 0 && toolIds . size === 0 ) {
89- return undefined ;
90- }
91- return { toolSetIds, toolIds } ;
107+ return map ;
108+ } ,
109+ key : 'chat/selectedTools' ,
92110 } ) ;
93111
94- this . _store . add ( autorun ( r => {
95112
96- const tools = this . _allTools . read ( r ) . filter ( t => t . canBeReferencedInPrompt ) ;
97- const toolSets = _toolsService . toolSets . read ( r ) ;
98-
99- const oldItems = new Set ( this . entriesMap . keys ( ) ) ;
100-
101- const disabledData = mode . read ( r ) === ChatMode . Agent
102- ? disabledDataObs . read ( r )
103- : undefined ;
113+ this . _selectedTools = this . _store . add ( storedTools ( StorageScope . WORKSPACE , StorageTarget . MACHINE , _storageService ) ) ;
114+ this . _allTools = observableFromEvent ( _toolsService . onDidChangeTools , ( ) => Array . from ( _toolsService . getTools ( ) ) ) ;
115+ }
104116
105- transaction ( tx => {
117+ /**
118+ * All tools and tool sets with their enabled state.
119+ */
120+ get entriesMap ( ) : IObservable < IToolAndToolSetEnablementMap > {
121+ return derived ( r => {
122+ const currentMode = this . _mode . read ( r ) ;
123+
124+ let currentMap = this . _sessionStates . get ( currentMode . id ) ;
125+ let defaultEnablement = false ;
126+ if ( ! currentMap && currentMode . kind === ChatMode . Agent && currentMode . customTools ) {
127+ currentMap = this . _toolsService . toToolAndToolSetEnablementMap ( new Set ( currentMode . customTools ) ) ;
128+ }
129+ if ( ! currentMap ) {
130+ currentMap = this . _selectedTools . read ( r ) ;
131+ defaultEnablement = true ;
132+ }
106133
107- for ( const tool of tools ) {
108- const enabled = ! disabledData || ! disabledData . toolIds . has ( tool . id ) ;
109- this . entriesMap . set ( tool , enabled , tx ) ;
110- oldItems . delete ( tool ) ;
111- }
134+ // create a complete map of all tools and tool sets
135+ const map = new Map < IToolData | ToolSet , boolean > ( ) ;
136+ const tools = this . _allTools . read ( r ) . filter ( t => t . canBeReferencedInPrompt ) ;
137+ for ( const tool of tools ) {
138+ map . set ( tool , currentMap . get ( tool ) ?? defaultEnablement ) ;
139+ }
140+ const toolSets = this . _toolsService . toolSets . read ( r ) ;
141+ for ( const toolSet of toolSets ) {
142+ map . set ( toolSet , currentMap . get ( toolSet ) ?? defaultEnablement ) ;
143+ }
144+ return map ;
145+ } ) ;
146+ }
112147
113- for ( const toolSet of toolSets ) {
114- const enabled = ! disabledData || ! disabledData . toolSetIds . has ( toolSet . id ) ;
115- this . entriesMap . set ( toolSet , enabled , tx ) ;
116- oldItems . delete ( toolSet ) ;
117- }
148+ get entriesScope ( ) {
149+ const mode = this . _mode . get ( ) ;
150+ if ( this . _sessionStates . has ( mode . id ) ) {
151+ return ToolsScope . Session ;
152+ }
153+ if ( mode . kind === ChatMode . Agent && mode . customTools && mode . uri ) {
154+ return ToolsScope . Mode ;
155+ }
156+ return ToolsScope . Global ;
157+ }
118158
119- for ( const item of oldItems ) {
120- this . entriesMap . delete ( item , tx ) ;
121- }
122- } ) ;
123- } ) ) ;
159+ get currentMode ( ) : IChatMode {
160+ return this . _mode . get ( ) ;
124161 }
125162
126163 resetSessionEnablementState ( ) {
127- this . _sessionSelectedTools . set ( { } , undefined ) ;
164+ const mode = this . _mode . get ( ) ;
165+ this . _sessionStates . delete ( mode . id ) ;
128166 }
129167
130- enable ( toolSets : readonly ToolSet [ ] , tools : readonly IToolData [ ] , sessionOnly : boolean ) : void {
131- const toolIds = new Set ( tools . map ( t => t . id ) ) ;
132- const toolsetIds = new Set ( toolSets . map ( t => t . id ) ) ;
133-
134- const disabledTools = this . _allTools . get ( ) . filter ( tool => ! toolIds . has ( tool . id ) ) ;
135- const disabledToolSets = Array . from ( this . _toolsService . toolSets . get ( ) ) . filter ( toolset => ! toolsetIds . has ( toolset . id ) ) ;
136-
137- this . disable ( disabledToolSets , disabledTools , sessionOnly ) ;
168+ set ( enablementMap : IToolAndToolSetEnablementMap , sessionOnly : boolean ) : void {
169+ const mode = this . _mode . get ( ) ;
170+ if ( sessionOnly ) {
171+ this . _sessionStates . set ( mode . id , enablementMap ) ;
172+ return ;
173+ }
174+ if ( this . _sessionStates . has ( mode . id ) ) {
175+ this . _sessionStates . set ( mode . id , enablementMap ) ;
176+ return ;
177+ }
178+ if ( mode . kind === ChatMode . Agent && mode . customTools && mode . uri ) {
179+ // apply directly to mode.
180+ this . updateCustomModeTools ( mode . uri , enablementMap ) ;
181+ return ;
182+ }
183+ this . _selectedTools . set ( enablementMap , undefined ) ;
138184 }
139185
140- disable ( disabledToolSets : readonly ToolSet [ ] , disableTools : readonly IToolData [ ] , sessionOnly : boolean ) : void {
141-
142- const target = sessionOnly
143- ? this . _sessionSelectedTools
144- : this . _selectedTools ;
145-
146- target . set ( {
147- disabledToolSets : disabledToolSets . map ( t => t . id ) ,
148- disabledTools : disableTools . map ( t => t . id )
149- } , undefined ) ;
186+ async updateCustomModeTools ( uri : URI , enablementMap : IToolAndToolSetEnablementMap ) : Promise < void > {
187+ await this . _instantiationService . createInstance ( PromptFileRewriter ) . openAndRewriteTools ( uri , enablementMap , CancellationToken . None ) ;
150188 }
151189
152190 asEnablementMap ( ) : Map < IToolData , boolean > {
153191 const result = new Map < IToolData , boolean > ( ) ;
154- const map = this . entriesMap ;
192+ const map = this . entriesMap . get ( ) ;
155193
156194 const _set = ( tool : IToolData , enabled : boolean ) => {
157195 // ONLY disable a tool that isn't enabled yet
0 commit comments