fix(cwl): Hide LiveTail CodeLenses and display info message when closing session

packages/core/src/awsService/cloudWatchLogs/activation.ts

@@ -34,7 +34,7 @@ export async function activate(context: vscode.ExtensionContext, configuration:
 
     const documentProvider = new LogDataDocumentProvider(registry)
     const liveTailDocumentProvider = new LiveTailDocumentProvider()
-
+    const liveTailCodeLensProvider = new LiveTailCodeLensProvider(liveTailRegistry)
     context.subscriptions.push(
         vscode.languages.registerCodeLensProvider(
             {
@@ -55,7 +55,7 @@ export async function activate(context: vscode.ExtensionContext, configuration:
                 language: 'log',
                 scheme: cloudwatchLogsLiveTailScheme,
             },
-            new LiveTailCodeLensProvider()
+            liveTailCodeLensProvider
         )
     )
 
@@ -121,11 +121,11 @@ export async function activate(context: vscode.ExtensionContext, configuration:
                     ? { regionName: node.regionCode, groupName: node.logGroup.logGroupName! }
                     : undefined
             const source = node ? (logGroupInfo ? 'ExplorerLogGroupNode' : 'ExplorerServiceNode') : 'Command'
-            await tailLogGroup(liveTailRegistry, source, logGroupInfo)
+            await tailLogGroup(liveTailRegistry, source, liveTailCodeLensProvider, logGroupInfo)
         }),
 
         Commands.register('aws.cwl.stopTailingLogGroup', async (document: vscode.TextDocument, source: string) => {
-            closeSession(document.uri, liveTailRegistry, source)
+            closeSession(document.uri, liveTailRegistry, source, liveTailCodeLensProvider)
         }),
 
         Commands.register('aws.cwl.clearDocument', async (document: vscode.TextDocument) => {

packages/core/src/awsService/cloudWatchLogs/commands/tailLogGroup.ts

@@ -16,10 +16,12 @@ import {
 } from '@aws-sdk/client-cloudwatch-logs'
 import { getLogger, globals, ToolkitError } from '../../../shared'
 import { uriToKey } from '../cloudWatchLogsUtils'
+import { LiveTailCodeLensProvider } from '../document/liveTailCodeLensProvider'
 
 export async function tailLogGroup(
     registry: LiveTailSessionRegistry,
     source: string,
+    codeLensProvider: LiveTailCodeLensProvider,
     logData?: { regionName: string; groupName: string }
 ): Promise<void> {
     await telemetry.cwlLiveTail_Start.run(async (span) => {
@@ -55,9 +57,11 @@ export async function tailLogGroup(
         const document = await prepareDocument(session)
 
         hideShowStatusBarItemsOnActiveEditor(session, document)
-        registerTabChangeCallback(session, registry, document)
+        registerTabChangeCallback(session, registry, document, codeLensProvider)
 
         const stream = await session.startLiveTailSession()
+        getLogger().info(`LiveTail session started: ${uriToKey(session.uri)}`)
+
         span.record({
             source: source,
             result: 'Succeeded',
@@ -69,14 +73,21 @@ export async function tailLogGroup(
     })
 }
 
-export function closeSession(sessionUri: vscode.Uri, registry: LiveTailSessionRegistry, source: string) {
+export function closeSession(
+    sessionUri: vscode.Uri,
+    registry: LiveTailSessionRegistry,
+    source: string,
+    codeLensProvider: LiveTailCodeLensProvider
+) {
     telemetry.cwlLiveTail_Stop.run((span) => {
         const session = registry.get(uriToKey(sessionUri))
         if (session === undefined) {
-            throw new ToolkitError(`No LiveTail session found for URI: ${sessionUri.toString()}`)
+            throw new ToolkitError(`No LiveTail session found for URI: ${uriToKey(sessionUri)}`)
         }
         session.stopLiveTailSession()
         registry.delete(uriToKey(sessionUri))
+        void vscode.window.showInformationMessage(`Stopped LiveTail session: ${uriToKey(sessionUri)}`)
+        codeLensProvider.refresh()
         span.record({
             result: 'Succeeded',
             source: source,
@@ -130,7 +141,7 @@ async function handleSessionStream(
             //AbortSignal interrupts the LiveTail stream, causing error to be thrown here.
             //Can assume that stopLiveTailSession() has already been called - AbortSignal is only
             //exposed through that method.
-            getLogger().info(`Session stopped: ${uriToKey(session.uri)}`)
+            getLogger().info(`LiveTail session stopped: ${uriToKey(session.uri)}`)
         } else {
             //Unexpected exception.
             session.stopLiveTailSession()
@@ -233,12 +244,13 @@ function hideShowStatusBarItemsOnActiveEditor(session: LiveTailSession, document
 function registerTabChangeCallback(
     session: LiveTailSession,
     registry: LiveTailSessionRegistry,
-    document: vscode.TextDocument
+    document: vscode.TextDocument,
+    codeLensProvider: LiveTailCodeLensProvider
 ) {
     vscode.window.tabGroups.onDidChangeTabs((tabEvent) => {
         const isOpen = isLiveTailSessionOpenInAnyTab(session)
         if (!isOpen) {
-            closeSession(session.uri, registry, 'ClosedEditors')
+            closeSession(session.uri, registry, 'ClosedEditors', codeLensProvider)
             void clearDocument(document)
         }
     })

packages/core/src/awsService/cloudWatchLogs/document/liveTailCodeLensProvider.ts

@@ -5,16 +5,22 @@
 
 import * as vscode from 'vscode'
 import { cloudwatchLogsLiveTailScheme } from '../../../shared/constants'
+import { LiveTailSessionRegistry } from '../registry/liveTailSessionRegistry'
+import { uriToKey } from '../cloudWatchLogsUtils'
 
 export class LiveTailCodeLensProvider implements vscode.CodeLensProvider {
-    onDidChangeCodeLenses?: vscode.Event<void> | undefined
+    private _onDidChangeCodeLenses: vscode.EventEmitter<void> = new vscode.EventEmitter<void>()
+    public readonly onDidChangeCodeLenses: vscode.Event<void> = this._onDidChangeCodeLenses.event
 
-    provideCodeLenses(
+    public constructor(private readonly registry: LiveTailSessionRegistry) {}
+
+    public provideCodeLenses(
         document: vscode.TextDocument,
         token: vscode.CancellationToken
     ): vscode.ProviderResult<vscode.CodeLens[]> {
         const uri = document.uri
-        if (uri.scheme !== cloudwatchLogsLiveTailScheme) {
+        //if registry does not contain session, it is assumed to have been stopped, thus, hide lenses.
+        if (uri.scheme !== cloudwatchLogsLiveTailScheme || !this.registry.has(uriToKey(uri))) {
             return []
         }
         const codeLenses: vscode.CodeLens[] = []
@@ -23,6 +29,10 @@ export class LiveTailCodeLensProvider implements vscode.CodeLensProvider {
         return codeLenses
     }
 
+    public refresh() {
+        this._onDidChangeCodeLenses.fire()
+    }
+
     private buildClearDocumentCodeLens(document: vscode.TextDocument): vscode.CodeLens {
         const range = this.getBottomOfDocumentRange(document)
         const command: vscode.Command = {

packages/core/src/awsService/cloudWatchLogs/registry/liveTailSession.ts

@@ -10,8 +10,8 @@ import {
     StartLiveTailResponseStream,
 } from '@aws-sdk/client-cloudwatch-logs'
 import { LogStreamFilterResponse } from '../wizard/liveTailLogStreamSubmenu'
-import { CloudWatchLogsSettings, uriToKey } from '../cloudWatchLogsUtils'
-import { getLogger, globals, Settings, ToolkitError } from '../../../shared'
+import { CloudWatchLogsSettings } from '../cloudWatchLogsUtils'
+import { globals, Settings, ToolkitError } from '../../../shared'
 import { createLiveTailURIFromArgs } from './liveTailSessionRegistry'
 import { getUserAgent } from '../../../shared/telemetry/util'
 import { convertToTimeString } from '../../../shared/datetime'
@@ -98,7 +98,6 @@ export class LiveTailSession {
         this.statusBarUpdateTimer = globals.clock.setInterval(() => {
             this.updateStatusBarItemText()
         }, 500)
-        getLogger().info(`LiveTail session started: ${uriToKey(this.uri)}`)
         return commandOutput.responseStream
     }
 

packages/core/src/test/awsService/cloudWatchLogs/commands/tailLogGroup.test.ts

@@ -21,6 +21,7 @@ import { getTestWindow } from '../../../shared/vscode/window'
 import { CloudWatchLogsSettings, uriToKey } from '../../../../awsService/cloudWatchLogs/cloudWatchLogsUtils'
 import { installFakeClock } from '../../../testUtil'
 import { DefaultAwsContext, ToolkitError } from '../../../../shared'
+import { LiveTailCodeLensProvider } from '../../../../awsService/cloudWatchLogs/document/liveTailCodeLensProvider'
 
 describe('TailLogGroup', function () {
     const testLogGroup = 'test-log-group'
@@ -32,6 +33,7 @@ describe('TailLogGroup', function () {
 
     let sandbox: sinon.SinonSandbox
     let registry: LiveTailSessionRegistry
+    let codeLensProvider: LiveTailCodeLensProvider
     let startLiveTailSessionSpy: sinon.SinonSpy
     let stopLiveTailSessionSpy: sinon.SinonSpy
     let cloudwatchSettingsSpy: sinon.SinonSpy
@@ -47,6 +49,7 @@ describe('TailLogGroup', function () {
         clock.reset()
         sandbox = sinon.createSandbox()
         registry = new LiveTailSessionRegistry()
+        codeLensProvider = new LiveTailCodeLensProvider(registry)
     })
 
     after(function () {
@@ -94,7 +97,7 @@ describe('TailLogGroup', function () {
         cloudwatchSettingsSpy = sandbox.stub(CloudWatchLogsSettings.prototype, 'get').callsFake(() => {
             return 1
         })
-        await tailLogGroup(registry, testSource, {
+        await tailLogGroup(registry, testSource, codeLensProvider, {
             groupName: testLogGroup,
             regionName: testRegion,
         })
@@ -132,7 +135,7 @@ describe('TailLogGroup', function () {
             return getTestWizardResponse()
         })
         await assert.rejects(async () => {
-            await tailLogGroup(registry, testSource, {
+            await tailLogGroup(registry, testSource, codeLensProvider, {
                 groupName: testLogGroup,
                 regionName: testRegion,
             })
@@ -153,7 +156,7 @@ describe('TailLogGroup', function () {
         })
         registry.set(uriToKey(session.uri), session)
 
-        closeSession(session.uri, registry, testSource)
+        closeSession(session.uri, registry, testSource, codeLensProvider)
         assert.strictEqual(0, registry.size)
         assert.strictEqual(true, stopLiveTailSessionSpy.calledOnce)
         assert.strictEqual(0, clock.countTimers())