From 12f3c82d3e627f41622ef7e1858a3aa1e5f96da9 Mon Sep 17 00:00:00 2001
From: Alessandro Mauri <alemauri001@gmail.com>
Date: Wed, 20 Dec 2023 19:37:28 +0100
Subject: [PATCH] display tokens as they arrive

---
 src/common.ts     | 18 ++++++++++++++++++
 src/extension.ts  |  3 +++
 src/openai-api.ts |  4 ++++
 3 files changed, 25 insertions(+)

diff --git a/src/common.ts b/src/common.ts
index 5cc1833..c531d15 100644
--- a/src/common.ts
+++ b/src/common.ts
@@ -1,3 +1,4 @@
+import { mkdirSync } from 'fs';
 import * as vscode from 'vscode';
 
 // common data structures and functions
@@ -56,3 +57,20 @@ export async function showPendingStatusBar(
 		);
 	// we already resolve the operation elsewhere
 }
+
+let st_msg: vscode.StatusBarItem | undefined;
+
+export function updateStatusBarMessage(text: string) {
+	if (!st_msg) {
+		st_msg = vscode.window.createStatusBarItem(vscode.StatusBarAlignment.Left, -100);
+	}
+
+	const run_color = new vscode.ThemeColor('statusBarItem.warningBackground');
+	if (text.length > 0) {
+		st_msg.backgroundColor = run_color;
+		st_msg.text = '$(megaphone) ' + text.trim();
+		st_msg.show();
+	} else {
+		st_msg.hide();
+	}
+}
diff --git a/src/extension.ts b/src/extension.ts
index 88a8196..cd816f8 100644
--- a/src/extension.ts
+++ b/src/extension.ts
@@ -18,6 +18,7 @@ import {
 	ResponseData,
 	showMessageWithTimeout,
 	showPendingStatusBar,
+	updateStatusBarMessage,
 } from './common';
 
 // clean up the document
@@ -53,6 +54,8 @@ export function activate(context: vscode.ExtensionContext) {
 		config.update('completionEnabled', false);
 	});
 
+	updateStatusBarMessage('');
+
 	// Register a new provider of inline completions, this does not decide how it is invoked
 	// only what the completion should be
 	// https://github.com/microsoft/vscode-extension-samples/blob/main/inline-completions/src/extension.ts
diff --git a/src/openai-api.ts b/src/openai-api.ts
index b3227be..16178c0 100644
--- a/src/openai-api.ts
+++ b/src/openai-api.ts
@@ -4,6 +4,7 @@ import {
 	ResponseData,
 	showMessageWithTimeout,
 	showPendingStatusBar,
+	updateStatusBarMessage,
 } from './common';
 import { config } from 'process';
 
@@ -193,6 +194,7 @@ export async function openAIMakeRequest(
 			// FIXME: why the choices may be multiple?
 			// TODO: display the multiple choices
 			//console.log(data.choices[0].text);
+			updateStatusBarMessage(data.choices[0].text);
 			ret.content += data.choices[0].text;
 			ret.tokens += data.usage?.completion_tokens || 0;
 		}
@@ -200,6 +202,8 @@ export async function openAIMakeRequest(
 		// stop the timer
 		const timer_end = performance.now();
 		ret.time = (timer_end - timer_start) / 1000.0;
+		// clear the status bar item
+		updateStatusBarMessage('');
 	} catch (e: any) {
 		console.error(e);
 		const err = e as TypeError;