From 52ec2ec265bdde684c691d535334c7a21de1ab5c Mon Sep 17 00:00:00 2001
From: nai-degen <khoners@protonmail.com>
Date: Mon, 23 Oct 2023 22:23:06 -0500
Subject: [PATCH] fixes blank AWS responses due to reqs sometimes using wrong
 handler

---
 src/proxy/middleware/request/add-key.ts  | 4 ----
 src/proxy/middleware/request/rewrite.ts  | 6 ++++++
 src/proxy/middleware/response/index.ts   | 9 ++++-----
 src/shared/key-management/aws/checker.ts | 2 +-
 4 files changed, 11 insertions(+), 10 deletions(-)

diff --git a/src/proxy/middleware/request/add-key.ts b/src/proxy/middleware/request/add-key.ts
index c73ce90..507c292 100644
--- a/src/proxy/middleware/request/add-key.ts
+++ b/src/proxy/middleware/request/add-key.ts
@@ -34,10 +34,6 @@ export const addKey: ProxyRequestMiddleware = (proxyReq, req) => {
     throw new Error("You must specify a model with your request.");
   }
 
-  // TODO: use separate middleware to deal with stream flags
-  req.isStreaming = req.body.stream === true || req.body.stream === "true";
-  req.body.stream = req.isStreaming;
-
   if (req.inboundApi === req.outboundApi) {
     assignedKey = keyPool.get(req.body.model);
   } else {
diff --git a/src/proxy/middleware/request/rewrite.ts b/src/proxy/middleware/request/rewrite.ts
index 203e51a..3b62ff2 100644
--- a/src/proxy/middleware/request/rewrite.ts
+++ b/src/proxy/middleware/request/rewrite.ts
@@ -14,6 +14,12 @@ export const createOnProxyReqHandler = ({
   pipeline,
 }: RewriterOptions): ProxyReqCallback => {
   return (proxyReq, req, res, options) => {
+    // The streaming flag must be set before any other middleware runs, because
+    // it may influence which other middleware a particular API pipeline wants
+    // to run.
+    req.isStreaming = req.body.stream === true || req.body.stream === "true";
+    req.body.stream = req.isStreaming;
+
     try {
       for (const validator of beforeRewrite) {
         validator(proxyReq, req, res, options);
diff --git a/src/proxy/middleware/response/index.ts b/src/proxy/middleware/response/index.ts
index 1858712..1f1c7b8 100644
--- a/src/proxy/middleware/response/index.ts
+++ b/src/proxy/middleware/response/index.ts
@@ -3,7 +3,6 @@ import { Request, Response } from "express";
 import * as http from "http";
 import util from "util";
 import zlib from "zlib";
-import { logger } from "../../../logger";
 import { enqueue, trackWaitTime } from "../../queue";
 import { HttpError } from "../../../shared/errors";
 import { keyPool } from "../../../shared/key-management";
@@ -188,7 +187,7 @@ export const decodeResponseBody: RawResponseBodyHandler = async (
           body = await decoder(body);
         } else {
           const errorMessage = `Proxy received response with unsupported content-encoding: ${contentEncoding}`;
-          logger.warn({ contentEncoding, key: req.key?.hash }, errorMessage);
+          req.log.warn({ contentEncoding, key: req.key?.hash }, errorMessage);
           writeErrorResponse(req, res, 500, {
             error: errorMessage,
             contentEncoding,
@@ -205,7 +204,7 @@ export const decodeResponseBody: RawResponseBodyHandler = async (
         return resolve(body.toString());
       } catch (error: any) {
         const errorMessage = `Proxy received response with invalid JSON: ${error.message}`;
-        logger.warn({ error: error.stack, key: req.key?.hash }, errorMessage);
+        req.log.warn({ error: error.stack, key: req.key?.hash }, errorMessage);
         writeErrorResponse(req, res, 500, { error: errorMessage });
         return reject(errorMessage);
       }
@@ -251,7 +250,7 @@ const handleUpstreamErrors: ProxyResHandlerWithBody = async (
     // Likely Bad Gateway or Gateway Timeout from upstream's reverse proxy
     const hash = req.key?.hash;
     const statusMessage = proxyRes.statusMessage || "Unknown error";
-    logger.warn({ statusCode, statusMessage, key: hash }, parseError.message);
+    req.log.warn({ statusCode, statusMessage, key: hash }, parseError.message);
 
     const errorObject = {
       statusCode,
@@ -268,7 +267,7 @@ const handleUpstreamErrors: ProxyResHandlerWithBody = async (
     errorPayload.error?.type ||
     getAwsErrorType(proxyRes.headers["x-amzn-errortype"]);
 
-  logger.warn(
+  req.log.warn(
     { statusCode, type: errorType, errorPayload, key: req.key?.hash },
     `Received error response from upstream. (${proxyRes.statusMessage})`
   );
diff --git a/src/shared/key-management/aws/checker.ts b/src/shared/key-management/aws/checker.ts
index ec6edf3..4640735 100644
--- a/src/shared/key-management/aws/checker.ts
+++ b/src/shared/key-management/aws/checker.ts
@@ -182,7 +182,7 @@ export class AwsKeyChecker extends KeyCheckerBase<AwsBedrockKey> {
     }
 
     this.log.debug(
-      { key: key.hash, errorType, data, status },
+      { key: key.hash, errorType, data, status, model },
       "Liveness test complete."
     );
   }