From 19236d2a960d56874bbdba1a8094a8504dcc5252 Mon Sep 17 00:00:00 2001
From: pelikhan <jhalleux@microsoft.com>
Date: Wed, 18 Dec 2024 18:10:22 -0800
Subject: [PATCH] =?UTF-8?q?feat:=20=E2=9C=A8=20add=20clear=20method=20to?=
 =?UTF-8?q?=20AzureTokenResolver=20and=20update=20usages?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 packages/cli/src/azuretoken.ts | 12 ++++++++++--
 packages/cli/src/nodehost.ts   |  4 ++++
 packages/core/src/host.ts      |  1 +
 3 files changed, 15 insertions(+), 2 deletions(-)

diff --git a/packages/cli/src/azuretoken.ts b/packages/cli/src/azuretoken.ts
index 60e8aae9c..78013ba65 100644
--- a/packages/cli/src/azuretoken.ts
+++ b/packages/cli/src/azuretoken.ts
@@ -108,6 +108,12 @@ class AzureTokenResolverImpl implements AzureTokenResolver {
         return this._error
     }
 
+    clear() {
+        this._token = undefined
+        this._error = undefined
+        this._resolver = undefined
+    }
+
     async token(
         credentialsType: AzureCredentialsType,
         options?: CancellationOptions
@@ -121,23 +127,25 @@ class AzureTokenResolverImpl implements AzureTokenResolver {
         if (!this._resolver) {
             const scope = await runtimeHost.readSecret(this.envName)
             const scopes = scope ? scope.split(",") : this.scopes
-            this._resolver = createAzureToken(
+            const resolver = (this._resolver = createAzureToken(
                 scopes,
                 credentialsType,
                 cancellationToken
             )
                 .then((res) => {
+                    if (this._resolver !== resolver) return undefined
                     this._token = res
                     this._error = undefined
                     this._resolver = undefined
                     return { token: this._token, error: this._error }
                 })
                 .catch((err) => {
+                    if (this._resolver !== resolver) return undefined
                     this._resolver = undefined
                     this._token = undefined
                     this._error = serializeError(err)
                     return { token: this._token, error: this._error }
-                })
+                }))
         }
         return this._resolver
     }
diff --git a/packages/cli/src/nodehost.ts b/packages/cli/src/nodehost.ts
index 3a5dafceb..427e72550 100644
--- a/packages/cli/src/nodehost.ts
+++ b/packages/cli/src/nodehost.ts
@@ -128,6 +128,7 @@ export class NodeHost implements RuntimeHost {
                 source: "default",
                 candidates: DEFAULT_REASONING_SMALL_MODEL_CANDIDATES,
             },
+            long: { model: LARGE_MODEL_ID, source: "default" },
             agent: { model: LARGE_MODEL_ID, source: "default" },
             memory: { model: SMALL_MODEL_ID, source: "default" },
         },
@@ -195,6 +196,9 @@ export class NodeHost implements RuntimeHost {
     }
 
     async readConfig(): Promise<HostConfiguration> {
+        this.azureToken.clear()
+        this.azureServerlessToken.clear()
+
         const config = await resolveGlobalConfiguration(this.dotEnvPath)
         const { envFile, modelAliases } = config
         if (modelAliases)
diff --git a/packages/core/src/host.ts b/packages/core/src/host.ts
index 023f6c159..58064117c 100644
--- a/packages/core/src/host.ts
+++ b/packages/core/src/host.ts
@@ -103,6 +103,7 @@ export function isAzureTokenExpired(token: AuthenticationToken) {
 }
 
 export interface AzureTokenResolver {
+    clear(): void
     token(
         credentialsType: AzureCredentialsType,
         options?: CancellationOptions