From f8b5031c52f3d91089355cc881de926a1b73f516 Mon Sep 17 00:00:00 2001
From: Kira Pilot <kirapilot@Kiras-MacBook-Pro.local>
Date: Tue, 20 Jun 2023 20:28:26 -0400
Subject: [PATCH] feat: add VS code notifications for autostop

---
 src/remote.ts             | 40 +++++++++++++++++++++++++++++++++++++++
 src/storage.ts            |  2 ++
 src/workspacesProvider.ts | 11 ++++++++---
 3 files changed, 50 insertions(+), 3 deletions(-)

diff --git a/src/remote.ts b/src/remote.ts
index 5d2f1134..aa647d0c 100644
--- a/src/remote.ts
+++ b/src/remote.ts
@@ -3,6 +3,7 @@ import {
   getBuildInfo,
   getTemplate,
   getWorkspace,
+  getWorkspaces,
   getWorkspaceBuildLogs,
   getWorkspaceByOwnerAndName,
   startWorkspace,
@@ -126,6 +127,44 @@ export class Remote {
       this.registerLabelFormatter(remoteAuthority, this.storage.workspace.owner_name, this.storage.workspace.name),
     )
 
+    const notifyWorkspacesEligibleForAutostop = () => {
+      const eligibleWorkspaces = this.storage.ownedWorkspaces?.filter((workspace: Workspace) => {
+        if (workspace.latest_build.transition !== "start" || !workspace.latest_build.deadline) {
+          return false
+        }
+
+        const hoursMilli = 1000 * 60 * 60
+        // return workspaces with a deadline that is in 1 hr or less
+        return Math.abs(new Date().getTime() - new Date(workspace.latest_build.deadline).getTime()) <= hoursMilli
+      })
+
+      eligibleWorkspaces?.forEach((workspace: Workspace) => {
+        if (this.storage.workspaceIdsEligibleForAutostop?.includes(workspace.id)) {
+          // don't message the user; we've already messaged
+          return
+        }
+        // we display individual notifications for each workspace as VS Code
+        // intentionally strips new lines from the message text
+        // https://github.com/Microsoft/vscode/issues/48900
+        this.vscodeProposed.window.showInformationMessage(`${workspace.name} is scheduled to shut down in 1 hour.`)
+        this.storage.workspaceIdsEligibleForAutostop?.push(workspace.id)
+      })
+    }
+
+    let errorCount = 0
+    const fetchWorkspacesInterval = setInterval(async () => {
+      try {
+        const workspacesResult = await getWorkspaces({ q: "owner:me" })
+        this.storage.ownedWorkspaces = workspacesResult.workspaces
+        notifyWorkspacesEligibleForAutostop()
+      } catch (error) {
+        if (errorCount === 3) {
+          clearInterval(fetchWorkspacesInterval)
+        }
+        errorCount++
+      }
+    }, 1000 * 5)
+
     let buildComplete: undefined | (() => void)
     if (this.storage.workspace.latest_build.status === "stopped") {
       this.vscodeProposed.window.withProgress(
@@ -427,6 +466,7 @@ export class Remote {
     return {
       dispose: () => {
         eventSource.close()
+        clearInterval(fetchWorkspacesInterval)
         disposables.forEach((d) => d.dispose())
       },
     }
diff --git a/src/storage.ts b/src/storage.ts
index 588f3408..d1aab723 100644
--- a/src/storage.ts
+++ b/src/storage.ts
@@ -14,6 +14,8 @@ import * as vscode from "vscode"
 
 export class Storage {
   public workspace?: Workspace
+  public ownedWorkspaces?: Workspace[] = []
+  public workspaceIdsEligibleForAutostop?: string[] = []
 
   constructor(
     private readonly output: vscode.OutputChannel,
diff --git a/src/workspacesProvider.ts b/src/workspacesProvider.ts
index 115d3f80..5353298f 100644
--- a/src/workspacesProvider.ts
+++ b/src/workspacesProvider.ts
@@ -20,15 +20,20 @@ export class WorkspaceProvider implements vscode.TreeDataProvider<vscode.TreeIte
       .then((workspaces) => {
         const workspacesTreeItem: WorkspaceTreeItem[] = []
         workspaces.workspaces.forEach((workspace) => {
-          const showMetadata = this.getWorkspacesQuery === WorkspaceQuery.Mine
-          if (showMetadata) {
+          const isOwnedWorkspace = this.getWorkspacesQuery === WorkspaceQuery.Mine
+          if (isOwnedWorkspace) {
+            // update ownedWorkspaces list in storage such that we can display to the user
+            // notifications about their own workspaces
+            this.storage.ownedWorkspaces?.push(workspace)
+
+            // Show metadata for workspaces owned by the user
             const agents = extractAgents(workspace)
             agents.forEach((agent) => this.monitorMetadata(agent.id)) // monitor metadata for all agents
           }
           const treeItem = new WorkspaceTreeItem(
             workspace,
             this.getWorkspacesQuery === WorkspaceQuery.All,
-            showMetadata,
+            isOwnedWorkspace,
           )
           workspacesTreeItem.push(treeItem)
         })