Support connecting Grammarly account on github.dev (#246)

* Use IndexDB to persist data in localStorage

* restart on error

* Use vscode.env.asExternalUri() to get a redirect URI in web extension

* add changeset

* use a netlify function to handle redirect using state

Co-authored-by: Rahul Kadyan <rahul.kadyan@grammarly.com>

.changeset/happy-coats-report.md

@@ -0,0 +1,6 @@
+---
+'grammarly': minor
+'grammarly-languageserver': patch
+---
+
+Support for connected Grammarly account in web extension (https://github.dev and https://vscode.dev)

.gitignore

@@ -2,3 +2,7 @@ node_modules/
 *.vsix
 dist/
 .tmp/
+.vscode-test-web/
+
+# Local Netlify folder
+.netlify

.vscode/settings.json

@@ -9,17 +9,18 @@
   "editor.formatOnSave": true,
   "cSpell.words": [
     "AUTOCORRECT",
-    "Emogenie",
-    "Grammarly",
     "capi",
     "docid",
+    "Emogenie",
     "errored",
     "freews",
     "gnar",
+    "Grammarly",
     "heatmap",
     "inversify",
     "languageclient",
     "minicard",
+    "reqid",
     "subalerts",
     "textdocument"
   ],

extension/src/GrammarlyClient.ts

@@ -16,15 +16,13 @@ import {
 import { Registerable } from './interfaces'
 
 export class GrammarlyClient implements Registerable {
-  public client: GrammarlyLanguageClient
+  public client!: GrammarlyLanguageClient
   private session?: Disposable
   private callbacks = new Set<() => unknown>()
   private isReady = false
   private selectors: DocumentFilter[] = []
 
-  constructor(private readonly context: ExtensionContext) {
-    this.client = this.createClient()
-  }
+  constructor(private readonly context: ExtensionContext) {}
 
   public onReady(fn: () => unknown): Disposable {
     this.callbacks.add(fn)
@@ -64,11 +62,15 @@ export class GrammarlyClient implements Registerable {
       {
         id: 'client_BaDkMgx4X19X9UxxYRCXZo',
         name: 'Grammarly',
+        outputChannel: window.createOutputChannel('Grammarly'),
         documentSelector: this.selectors
           .map((selector) =>
             selector.language != null || selector.pattern != null || selector.scheme != null ? (selector as any) : null,
           )
           .filter(<T>(value: T | null): value is T => value != null),
+
+        revealOutputChannelOn: 3,
+        progressOnInitialization: true,
         errorHandler: {
           error(error) {
             window.showErrorMessage(error.message)
@@ -90,9 +92,7 @@ export class GrammarlyClient implements Registerable {
         handleUri: async (uri) => {
           if (uri.path === '/auth/callback') {
             try {
-              await this.client.protocol.handleOAuthCallbackUri(
-                `${uri.scheme}://${uri.authority}${uri.path}?${decodeURIComponent(uri.query)}`,
-              )
+              await this.client.protocol.handleOAuthCallbackUri(uri.toString(true))
             } catch (error) {
               await window.showErrorMessage((error as Error).message)
               return
@@ -148,16 +148,17 @@ export class GrammarlyClient implements Registerable {
         await this.client.protocol.dismissSuggestion(options)
       }),
       commands.registerCommand('grammarly.login', async () => {
-        if (env.appHost !== 'desktop') {
-          await window.showErrorMessage('Connected account is not supported in web extension yet.')
-          return
-        }
+        const internalRedirectUri = Uri.parse(`${env.uriScheme}://znck.grammarly/auth/callback`, true)
+        const externalRedirectUri = await env.asExternalUri(internalRedirectUri)
 
-        const url = await this.client.protocol.getOAuthUrl(
-          Uri.from({ scheme: env.uriScheme, authority: 'znck.grammarly', path: '/auth/callback' }).toString(),
-        )
+        const isExternalURLDifferent = internalRedirectUri.toString(true) === externalRedirectUri.toString(true)
+        const redirectUri = isExternalURLDifferent
+          ? internalRedirectUri.toString(true)
+          : 'https://vscode-extension-grammarly.netlify.app/.netlify/functions/redirect'
+        const url = new URL(await this.client.protocol.getOAuthUrl(redirectUri))
+        url.searchParams.set('state', toBase64URL(externalRedirectUri.toString(true)))
 
-        if (!(await env.openExternal(Uri.parse(url)))) {
+        if (!(await env.openExternal(Uri.parse(url.toString(), true)))) {
           await window.showErrorMessage('Failed to open login page.')
         }
       }),
@@ -186,6 +187,8 @@ export class GrammarlyClient implements Registerable {
           console.error(error)
         }
       })
+    } catch (error) {
+      await window.showErrorMessage(`The extension couldn't be started. See the output channel for details.`)
     } finally {
       statusbar.dispose()
     }
@@ -195,3 +198,8 @@ export class GrammarlyClient implements Registerable {
 function isNode(): boolean {
   return typeof process !== 'undefined' && process.versions?.node != null
 }
+
+function toBase64URL(text: string): string {
+  if (typeof Buffer !== 'undefined') return Buffer.from(text, 'utf-8').toString('base64url')
+  return btoa(text).replace(/=/g, '').replace(/\+/g, '-').replace(/\//g, '_')
+}

extension/src/StatusBarController.ts

@@ -21,12 +21,20 @@ export class StatusBarController {
 
   public register() {
     this.update()
+
+    let isRestarting = false
     return Disposable.from(
       this.#statusbar,
       workspace.onDidCloseTextDocument(() => this.update()),
       window.onDidChangeActiveTextEditor(() => this.update()),
       commands.registerCommand('grammarly.restartServer', async () => {
-        // TODO: Restart?
+        if (isRestarting) return
+        try {
+          isRestarting = true
+          await this.grammarly.start()
+        } finally {
+          isRestarting = false
+        }
       }),
     )
   }
@@ -62,9 +70,13 @@ export class StatusBarController {
       role: status === 'error' ? 'button' : undefined,
     }
 
-    this.#statusbar.tooltip = `Your Grammarly account is ${
-      isUser ? '' : 'not '
-    }used for this file. \nConnection status: ${status}`
+    this.#statusbar.tooltip = [
+      `Your Grammarly account is ${isUser ? '' : 'not '}used for this file.`,
+      `Connection status: ${status}`,
+      status === 'error' ? `Restart now?` : null,
+    ]
+      .filter(Boolean)
+      .join('\n')
     this.#statusbar.command = status === 'error' ? 'grammarly.restartServer' : undefined
     this.#statusbar.show()
   }

package.json

@@ -8,15 +8,18 @@
     "build:wasm": "node scripts/build-wasm.mjs",
     "watch": "rollup --watch -c",
     "test": "jest",
-    "release": "changeset version"
+    "release": "changeset version",
+    "open-in-browser": "vscode-test-web --host 127.0.0.1 --browser firefox --extensionDevelopmentPath=./extension ./fixtures"
   },
   "devDependencies": {
     "@changesets/cli": "^2.22.0",
+    "@netlify/functions": "^1.0.0",
     "@rollup/plugin-alias": "^3.1.9",
     "@rollup/plugin-commonjs": "^22.0.0",
     "@rollup/plugin-node-resolve": "^13.2.1",
     "@rollup/plugin-replace": "^4.0.0",
     "@rollup/plugin-typescript": "^8.3.2",
+    "@vscode/test-web": "^0.0.24",
     "@vuedx/monorepo-tools": "^0.2.2-next-1651055813.0",
     "esbuild": "^0.14.38",
     "husky": "^7.0.4",

packages/grammarly-languageserver/package.json

@@ -69,6 +69,7 @@
   "dependencies": {
     "@grammarly/sdk": "^1.7.4",
     "htmlparser2": "^8.0.1",
+    "idb-keyval": "^6.1.0",
     "inversify": "^6.0.1",
     "reflect-metadata": "^0.1.13",
     "vscode-languageserver": "^7.0.0",

packages/grammarly-languageserver/src/VirtualStorage.ts

@@ -1,3 +1,5 @@
+import { createStore, get, set, keys, del, delMany } from 'idb-keyval'
+
 export class VirtualStorage {
   public items = new Map<string, string>()
 
@@ -25,3 +27,43 @@ export class VirtualStorage {
     this.items.set(key, value)
   }
 }
+
+export class IDBStorage extends VirtualStorage {
+  private store = createStore('grammarly-languageserver', 'localStorage')
+
+  async load(): Promise<void> {
+    for (const key of await keys<string>(this.store)) {
+      const value = await get<string>(key, this.store)
+      if (value != null) {
+        this.items.set(key, value)
+      }
+    }
+  }
+
+  clear(): void {
+    delMany(Array.from(this.items.keys()), this.store)
+    super.clear()
+  }
+
+  getItem(key: string): string | null {
+    get(key, this.store).then((value) => {
+      if (value != null) {
+        this.items.set(key, value)
+      } else {
+        this.items.delete(key)
+      }
+    })
+
+    return super.getItem(key)
+  }
+
+  removeItem(key: string): void {
+    del(key, this.store)
+    this.items.delete(key)
+  }
+
+  setItem(key: string, value: string): void {
+    set(key, value, this.store)
+    this.items.set(key, value)
+  }
+}

packages/grammarly-languageserver/src/createLanguageServer.ts

@@ -22,7 +22,7 @@ export interface Options {
   getConnection(): ReturnType<typeof createConnection>
   createTextDocuments<T>(config: TextDocumentsConfiguration<T>): TextDocuments<T>
   init(clientId: string): Promise<SDK>
-  pathEnvironmentForSDK(clientId: string): void
+  pathEnvironmentForSDK(clientId: string): void | Promise<void>
 }
 
 export function createLanguageServer({
@@ -58,7 +58,7 @@ export function createLanguageServer({
     connection.onInitialize(async (params) => {
       const options = params.initializationOptions as { clientId: string } | undefined
       if (options?.clientId == null) throw new Error('clientId is required')
-      pathEnvironmentForSDK(options.clientId)
+      await pathEnvironmentForSDK(options.clientId)
       const sdk = await init(options.clientId)
 
       container.bind(CLIENT).toConstantValue(params.capabilities)

packages/grammarly-languageserver/src/index.browser.ts

@@ -8,7 +8,7 @@ import {
 } from 'vscode-languageserver/browser'
 import { createLanguageServer } from './createLanguageServer'
 import { DOMParser } from './DOMParser'
-import { VirtualStorage } from './VirtualStorage'
+import { IDBStorage, VirtualStorage } from './VirtualStorage'
 function getConnection() {
   const messageReader = new BrowserMessageReader(self as unknown as Worker)
   const messageWriter = new BrowserMessageWriter(self as unknown as Worker)
@@ -21,8 +21,10 @@ function createTextDocuments<T>(config: TextDocumentsConfiguration<T>): TextDocu
 
 // Polyfill DOMParser as it is not available in worker.
 if (!('DOMParser' in globalThis)) (globalThis as any).DOMParser = DOMParser
-if (!('localStorage' in globalThis)) (globalThis as any).localStorage = new VirtualStorage()
-if (!('sessionStorage' in globalThis)) (globalThis as any).sessionStorage = new VirtualStorage()
+const localStorage = new IDBStorage()
+if (!('localStorage' in globalThis)) (globalThis as any).localStorage = localStorage
+const sessionStorage = new VirtualStorage()
+if (!('sessionStorage' in globalThis)) (globalThis as any).sessionStorage = sessionStorage
 
 export const startLanguageServer = createLanguageServer({
   getConnection,
@@ -31,5 +33,9 @@ export const startLanguageServer = createLanguageServer({
     // @ts-ignore
     return new globalThis.Grammarly.SDK(clientId)
   },
-  pathEnvironmentForSDK() {},
+  async pathEnvironmentForSDK() {
+    if (globalThis.localStorage === localStorage) {
+      await localStorage.load()
+    }
+  },
 })