@@ -66,7 +66,7 @@ import { DefaultAmazonQAppInitContext } from '../../../amazonq/apps/initContext'
6666import globals from '../../../shared/extensionGlobals'
6767import { MynahIconsType , MynahUIDataModel , QuickActionCommand } from '@aws/mynah-ui'
6868import { LspClient } from '../../../amazonq/lsp/lspClient'
69- import { ContextCommandItem , ContextCommandItemType } from '../../../amazonq/lsp/types'
69+ import { AdditionalContextPrompt , ContextCommandItem , ContextCommandItemType } from '../../../amazonq/lsp/types'
7070import { workspaceCommand } from '../../../amazonq/webview/ui/tabs/constants'
7171import fs from '../../../shared/fs/fs'
7272import { FeatureConfigProvider , Features } from '../../../shared/featureConfig'
@@ -588,8 +588,9 @@ export class ChatController {
588588 return
589589 }
590590
591- // TODO: Fix for multiple workspace setup
592- const projectRoot = session . relativePathToWorkspaceRoot . get ( message . filePath )
591+ // Check if clicked file is in a different workspace root
592+ const projectRoot =
593+ session . relativePathToWorkspaceRoot . get ( message . filePath ) || workspace . workspaceFolders ?. [ 0 ] ?. uri . fsPath
593594 if ( ! projectRoot ) {
594595 return
595596 }
@@ -877,6 +878,34 @@ export class ChatController {
877878 this . messenger . sendStaticTextResponse ( responseType , triggerID , tabID )
878879 }
879880
881+ /**
882+ * @returns A Uri array of prompt files in each workspace root's .amazonq/rules directory
883+ */
884+ private async collectWorkspaceRules ( ) : Promise < vscode . Uri [ ] > {
885+ const rulesFiles : vscode . Uri [ ] = [ ]
886+
887+ if ( ! vscode . workspace . workspaceFolders ) {
888+ return rulesFiles
889+ }
890+
891+ for ( const folder of vscode . workspace . workspaceFolders ) {
892+ const rulesPath = vscode . Uri . joinPath ( folder . uri , '.amazonq' , 'rules' )
893+ const folderExists = await fs . exists ( rulesPath )
894+
895+ if ( folderExists ) {
896+ const entries = await fs . readdir ( rulesPath )
897+
898+ for ( const [ name , type ] of entries ) {
899+ if ( type === vscode . FileType . File && name . endsWith ( promptFileExtension ) ) {
900+ rulesFiles . push ( vscode . Uri . joinPath ( rulesPath , name ) )
901+ }
902+ }
903+ }
904+ }
905+
906+ return rulesFiles
907+ }
908+
880909 private async resolveContextCommandPayload (
881910 triggerPayload : TriggerPayload ,
882911 session : ChatSession
@@ -885,14 +914,13 @@ export class ChatController {
885914 const relativePaths : string [ ] = [ ]
886915
887916 // Check for workspace rules to add to context
888- const workspaceRules = await vscode . workspace . findFiles ( `.amazonq/rules/* ${ promptFileExtension } ` )
917+ const workspaceRules = await this . collectWorkspaceRules ( )
889918 if ( workspaceRules . length > 0 ) {
890919 contextCommands . push (
891920 ...workspaceRules . map ( ( rule ) => {
892921 const workspaceFolderPath = vscode . workspace . getWorkspaceFolder ( rule ) ?. uri ?. path || ''
893922 return {
894923 workspaceFolder : workspaceFolderPath ,
895- // todo: add 'prompt' type to LSP model
896924 type : 'file' as ContextCommandItemType ,
897925 relativePath : path . relative ( workspaceFolderPath , rule . path ) ,
898926 }
@@ -904,7 +932,6 @@ export class ChatController {
904932 // Add context commands added by user to context
905933 if ( triggerPayload . context !== undefined && triggerPayload . context . length > 0 ) {
906934 for ( const context of triggerPayload . context ) {
907- // todo: add handling of 'prompt' type (dependent on LSP changes)
908935 if ( typeof context !== 'string' && context . route && context . route . length === 2 ) {
909936 contextCommands . push ( {
910937 workspaceFolder : context . route [ 0 ] || '' ,
@@ -930,52 +957,59 @@ export class ChatController {
930957 )
931958 session . relativePathToWorkspaceRoot . set ( relativePath , contextCommand . workspaceFolder )
932959 }
933- const prompts = await LspClient . instance . getContextCommandPrompt ( contextCommands )
934- if ( prompts . length === 0 ) {
935- return [ ]
960+ let prompts : AdditionalContextPrompt [ ] = [ ]
961+ try {
962+ prompts = await LspClient . instance . getContextCommandPrompt ( contextCommands )
963+ } catch ( e ) {
964+ // todo: handle @workspace used before indexing is ready
965+ getLogger ( ) . verbose ( `Could not get context command prompts: ${ e } )
936966 }
937967
938968 let currentContextLength = 0
939969 triggerPayload . additionalContents = [ ]
940- triggerPayload . additionalContextLengths = this . telemetryHelper . getContextLengths ( prompts )
941- triggerPayload . truncatedAdditionalContextLengths = {
970+ const emptyLengths = {
942971 fileContextLength : 0 ,
943972 promptContextLength : 0 ,
944973 ruleContextLength : 0 ,
945974 }
946- for ( const prompt of prompts . slice ( 0 , 20 ) ) {
947- // Todo: add mechanism for sorting/prioritization of additional context
948- const entry = {
949- name : prompt . name . substring ( 0 , aditionalContentNameLimit ) ,
950- description : prompt . description . substring ( 0 , aditionalContentNameLimit ) ,
951- innerContext : prompt . content . substring ( 0 , additionalContentInnerContextLimit ) ,
952- }
953- // make sure the relevantDocument + additionalContext
954- // combined does not exceed 40k characters before generating the request payload.
955- // Do truncation and make sure triggerPayload.documentReferences is up-to-date after truncation
956- // TODO: Use a context length indicator
957- if ( currentContextLength + entry . innerContext . length > contextMaxLength ) {
958- getLogger ( ) . warn ( `Selected context exceeds context size limit: ${ entry . description } )
959- break
960- }
975+ triggerPayload . additionalContextLengths = emptyLengths
976+ triggerPayload . truncatedAdditionalContextLengths = emptyLengths
977+
978+ if ( Array . isArray ( prompts ) && prompts . length > 0 ) {
979+ triggerPayload . additionalContextLengths = this . telemetryHelper . getContextLengths ( prompts )
980+ for ( const prompt of prompts . slice ( 0 , 20 ) ) {
981+ const entry = {
982+ name : prompt . name . substring ( 0 , aditionalContentNameLimit ) ,
983+ description : prompt . description . substring ( 0 , aditionalContentNameLimit ) ,
984+ innerContext : prompt . content . substring ( 0 , additionalContentInnerContextLimit ) ,
985+ }
986+ // make sure the relevantDocument + additionalContext
987+ // combined does not exceed 40k characters before generating the request payload.
988+ // Do truncation and make sure triggerPayload.documentReferences is up-to-date after truncation
989+ // TODO: Use a context length indicator
990+ if ( currentContextLength + entry . innerContext . length > contextMaxLength ) {
991+ getLogger ( ) . warn ( `Selected context exceeds context size limit: ${ entry . description } )
992+ break
993+ }
961994
962- const contextType = this . telemetryHelper . getContextType ( prompt )
963- if ( contextType === 'rule' ) {
964- triggerPayload . truncatedAdditionalContextLengths . ruleContextLength += entry . innerContext . length
965- } else if ( contextType === 'prompt' ) {
966- triggerPayload . truncatedAdditionalContextLengths . promptContextLength += entry . innerContext . length
967- } else if ( contextType === 'file' ) {
968- triggerPayload . truncatedAdditionalContextLengths . fileContextLength += entry . innerContext . length
969- }
995+ const contextType = this . telemetryHelper . getContextType ( prompt )
996+ if ( contextType === 'rule' ) {
997+ triggerPayload . truncatedAdditionalContextLengths . ruleContextLength += entry . innerContext . length
998+ } else if ( contextType === 'prompt' ) {
999+ triggerPayload . truncatedAdditionalContextLengths . promptContextLength += entry . innerContext . length
1000+ } else if ( contextType === 'file' ) {
1001+ triggerPayload . truncatedAdditionalContextLengths . fileContextLength += entry . innerContext . length
1002+ }
9701003
971- triggerPayload . additionalContents . push ( entry )
972- currentContextLength += entry . innerContext . length
973- let relativePath = path . relative ( workspaceFolder , prompt . filePath )
974- // Handle user prompts outside the workspace
975- if ( prompt . filePath . startsWith ( getUserPromptsDirectory ( ) ) ) {
976- relativePath = path . basename ( prompt . filePath )
1004+ triggerPayload . additionalContents . push ( entry )
1005+ currentContextLength += entry . innerContext . length
1006+ let relativePath = path . relative ( workspaceFolder , prompt . filePath )
1007+ // Handle user prompts outside the workspace
1008+ if ( prompt . filePath . startsWith ( getUserPromptsDirectory ( ) ) ) {
1009+ relativePath = path . basename ( prompt . filePath )
1010+ }
1011+ relativePaths . push ( relativePath )
9771012 }
978- relativePaths . push ( relativePath )
9791013 }
9801014 getLogger ( ) . info ( `Retrieved chunks of additional context count: ${ triggerPayload . additionalContents . length } )
9811015
0 commit comments