From a254dc3095472ae63cdb302142a8f9a0fd9f87b2 Mon Sep 17 00:00:00 2001
From: Zack <zachary.siegel@tecton.ai>
Date: Mon, 2 Dec 2024 00:11:00 -0500
Subject: [PATCH] Set request index as a header

---
 .../java/ai/tecton/client/transport/TectonHttpClient.java    | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/src/main/java/ai/tecton/client/transport/TectonHttpClient.java b/src/main/java/ai/tecton/client/transport/TectonHttpClient.java
index ed6dbfc0..eb71ccbe 100644
--- a/src/main/java/ai/tecton/client/transport/TectonHttpClient.java
+++ b/src/main/java/ai/tecton/client/transport/TectonHttpClient.java
@@ -138,7 +138,8 @@ public void onResponse(Call call, Response response) {
             try (ResponseBody responseBody = response.body()) {
               // Add response to corresponding index
               parallelCallHandler.set(
-                  (Integer) call.request().tag(), new HttpResponse(response, responseBody));
+                  Integer.parseInt(call.request().header("request-index")),
+                  new HttpResponse(response, responseBody));
             } catch (Exception e) {
               throw new TectonServiceException(e.getMessage());
             } finally {
@@ -183,7 +184,7 @@ public Request buildRequestWithDefaultHeaders(HttpRequest httpRequest, int index
     okhttp3.MediaType mediaType = okhttp3.MediaType.parse(MediaType.APPLICATION_JSON.getName());
     RequestBody requestBody = RequestBody.create(httpRequest.getJsonBody(), mediaType);
     requestBuilder.post(requestBody);
-    requestBuilder.tag(index);
+    requestBuilder.header("request-index", String.valueOf(index));
     return requestBuilder.build();
   }