fixes Force Key Recheck admin function for azure/aws

src/admin/web/manage.ts

@@ -6,7 +6,7 @@ import { HttpError } from "../../shared/errors";
 import * as userStore from "../../shared/users/user-store";
 import { parseSort, sortBy, paginate } from "../../shared/utils";
 import { keyPool } from "../../shared/key-management";
-import { MODEL_FAMILIES } from "../../shared/models";
+import { LLMService, MODEL_FAMILIES } from "../../shared/models";
 import { getTokenCostUsd, prettyTokens } from "../../shared/stats";
 import {
   User,
@@ -196,13 +196,14 @@ router.post("/maintenance", (req, res) => {
   let flash = { type: "", message: "" };
   switch (action) {
     case "recheck": {
-      keyPool.recheck("openai");
-      keyPool.recheck("anthropic");
-      const size = keyPool
+      const checkable: LLMService[] = ["openai", "anthropic", "aws", "azure"];
+      checkable.forEach((s) => keyPool.recheck(s));
+      const keyCount = keyPool
         .list()
-        .filter((k) => k.service !== "google-ai").length;
+        .filter((k) => checkable.includes(k.service)).length;
+
       flash.type = "success";
-      flash.message = `Scheduled recheck of ${size} keys for OpenAI and Anthropic.`;
+      flash.message = `Scheduled recheck of ${keyCount} keys.`;
       break;
     }
     case "resetQuotas": {

src/proxy/middleware/response/streaming/sse-message-transformer.ts

@@ -12,14 +12,12 @@ import {
   StreamingCompletionTransformer,
 } from "./index";
 
-const genlog = logger.child({ module: "sse-transformer" });
-
 type SSEMessageTransformerOptions = TransformOptions & {
   requestedModel: string;
   requestId: string;
   inputFormat: APIFormat;
   inputApiVersion?: string;
-  logger?: typeof logger;
+  logger: typeof logger;
 };
 
 /**
@@ -37,7 +35,7 @@ export class SSEMessageTransformer extends Transform {
 
   constructor(options: SSEMessageTransformerOptions) {
     super({ ...options, readableObjectMode: true });
-    this.log = options.logger?.child({ module: "sse-transformer" }) ?? genlog;
+    this.log = options.logger?.child({ module: "sse-transformer" });
     this.lastPosition = 0;
     this.msgCount = 0;
     this.transformFn = getTransformer(

src/proxy/middleware/response/streaming/sse-stream-adapter.ts

@@ -136,9 +136,13 @@ export class SSEStreamAdapter extends Transform {
       }
       callback();
     } catch (error) {
-      error.lastEvent = data?.toString();
-      this.emit("error", error);
+      error.lastEvent = data?.toString() ?? "[SSEStreamAdapter] no data";
       callback(error);
     }
   }
+
+  _flush(callback: (err?: Error | null) => void) {
+    this.log.debug("SSEStreamAdapter flushing");
+    callback();
+  }
 }

src/shared/key-management/aws/provider.ts

@@ -192,6 +192,7 @@ export class AwsBedrockKeyProvider implements KeyProvider<AwsBedrockKey> {
     this.keys.forEach(({ hash }) =>
       this.update(hash, { lastChecked: 0, isDisabled: false })
     );
+    this.checker?.scheduleNextCheck();
   }
 
   /**

src/shared/key-management/azure/provider.ts

@@ -33,7 +33,7 @@ const RATE_LIMIT_LOCKOUT = 4000;
  * to be used again. This is to prevent the queue from flooding a key with too
  * many requests while we wait to learn whether previous ones succeeded.
  */
-const KEY_REUSE_DELAY = 250;
+const KEY_REUSE_DELAY = 500;
 
 export class AzureOpenAIKeyProvider implements KeyProvider<AzureOpenAIKey> {
   readonly service = "azure";
@@ -194,6 +194,7 @@ export class AzureOpenAIKeyProvider implements KeyProvider<AzureOpenAIKey> {
     this.keys.forEach(({ hash }) =>
       this.update(hash, { lastChecked: 0, isDisabled: false })
     );
+    this.checker?.scheduleNextCheck();
   }
 
   /**