@@ -184,7 +184,9 @@ export class Messenger {
184184 chatEvent . assistantResponseEvent . content . length > 0
185185 ) {
186186 message += chatEvent . assistantResponseEvent . content
187- codeBlockLanguage = this . extractCodeBlockLanguage ( message )
187+ if ( codeBlockLanguage === 'plaintext' ) {
188+ codeBlockLanguage = this . extractCodeBlockLanguage ( message )
189+ }
188190 this . dispatcher . sendChatMessage (
189191 new ChatMessage (
190192 {
@@ -341,9 +343,21 @@ export class Messenger {
341343 }
342344
343345 private extractCodeBlockLanguage ( message : string ) : string {
344- const firstLine = message . split ( '\n' ) [ 0 ]
345- const match = firstLine . match ( / ^ ` ` ` ( \w + ) / )
346- return match ? match [ 1 ] : 'plaintext'
346+ // This fulfills both the cases of unit test generation(java, python) and general use case(Non java and Non python) languages.
347+ const codeBlockStart = message . indexOf ( '```' )
348+ if ( codeBlockStart === - 1 ) {
349+ return 'plaintext'
350+ }
351+
352+ const languageStart = codeBlockStart + 3
353+ const languageEnd = message . indexOf ( '\n' , languageStart )
354+
355+ if ( languageEnd === - 1 ) {
356+ return 'plaintext'
357+ }
358+
359+ const language = message . substring ( languageStart , languageEnd ) . trim ( )
360+ return language !== '' ? language : 'plaintext'
347361 }
348362
349363 public sendErrorMessage ( errorMessage : string | undefined , tabID : string , requestID : string | undefined ) {
0 commit comments