feat: ✨ add clear method to AzureTokenResolver and update usages

microsoft · Dec 19, 2024 · 19236d2 · 19236d2
1 parent 44dd8a4
commit 19236d2
Show file tree

Hide file tree

Showing 3 changed files with 15 additions and 2 deletions.
diff --git 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
@@ -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
@@ -103,6 +103,7 @@ export function isAzureTokenExpired(token: AuthenticationToken) {
 }
 
 export interface AzureTokenResolver {
+    clear(): void
     token(
         credentialsType: AzureCredentialsType,
         options?: CancellationOptions