Add branch argument to chalk-ts

Signed-off-by: Kelvin Lu <kelvin@chalk.ai>

src/__test__/client.test.ts

@@ -34,16 +34,18 @@ describe("ChalkClient", () => {
     const client = new ChalkClient({
       activeEnvironment: "a",
       apiServer: "b",
-      clientId: "c",
-      clientSecret: "d",
+      branch: "c",
+      clientId: "d",
+      clientSecret: "e",
       timestampFormat: TimestampFormat.EPOCH_MILLIS,
     });
 
     expect(getConfig(client)).toEqual<ChalkClientConfig>({
       activeEnvironment: "a",
       apiServer: "b",
-      clientId: "c",
-      clientSecret: "d",
+      branch: "c",
+      clientId: "d",
+      clientSecret: "e",
       queryServer: "b",
       timestampFormat: TimestampFormat.EPOCH_MILLIS,
     });
@@ -52,6 +54,7 @@ describe("ChalkClient", () => {
   it("reads from environment variables when set", () => {
     process.env._CHALK_ACTIVE_ENVIRONMENT = "env";
     process.env._CHALK_API_SERVER = "http://localhost:8000";
+    process.env._CHALK_BRANCH = "not_a_real_branch";
     process.env._CHALK_CLIENT_ID = "client_id";
     process.env._CHALK_CLIENT_SECRET = "secret";
     process.env._CHALK_QUERY_SERVER = "http://localhost:1337";
@@ -61,6 +64,7 @@ describe("ChalkClient", () => {
     expect(getConfig(client)).toEqual<ChalkClientConfig>({
       activeEnvironment: "env",
       apiServer: "http://localhost:8000",
+      branch: "not_a_real_branch",
       clientId: "client_id",
       clientSecret: "secret",
       queryServer: "http://localhost:1337",
@@ -84,6 +88,7 @@ describe("ChalkClient", () => {
     expect(getConfig(client)).toEqual<ChalkClientConfig>({
       activeEnvironment: undefined,
       apiServer: DEFAULT_API_SERVER,
+      branch: undefined,
       clientId: "client_id",
       clientSecret: "secret",
       queryServer: DEFAULT_API_SERVER,
@@ -99,6 +104,7 @@ describe("ChalkClient", () => {
     expect(getConfig(client)).toEqual<ChalkClientConfig>({
       activeEnvironment: undefined,
       apiServer: DEFAULT_API_SERVER,
+      branch: undefined,
       clientId: "client_id",
       clientSecret: "secret",
       queryServer: "query server",

src/_client.ts

@@ -56,12 +56,21 @@ export interface ChalkClientOpts {
   queryServer?: string;
 
   /**
-   * The environment that your client will run against. This value will be read from the _CHALK_ACTIVE_ENVIRONMENT environment variable if not set explicitly.
+   * The environment that your client will run against.
+   * This value will be read from the _CHALK_ACTIVE_ENVIRONMENT environment variable if not set explicitly.
    *
    * If not specified and unset by your environment, an error will be thrown on client creation
    */
   activeEnvironment?: string;
 
+  /**
+   * If specified, Chalk will route all requests from this client instance to the relevant branch.
+   * This value will be read from the _CHALK_BRANCH environment variable if not set explicitly.
+   *
+   * Some methods allow you to override this instance-level branch configuration by passing in a `branch` argument.
+   */
+  branch?: string;
+
   /**
    * A custom fetch client that will replace the fetch polyfill used by default.
    *
@@ -106,6 +115,11 @@ function valueWithEnvFallback(
 }
 
 export interface ChalkRequestOptions {
+  /**
+   * If specified, Chalk will route this request to the relevant branch. Overrides the branch passed in to the
+   * client initialization.
+   */
+  branch?: string;
   /**
    * The timeout for the request in milliseconds. If not provided, the client will use the default timeout
    * specified at the client level.
@@ -131,6 +145,7 @@ export class ChalkClient<TFeatureMap = Record<string, ChalkScalar>>
         process.env._CHALK_ACTIVE_ENVIRONMENT ??
         undefined,
       apiServer: resolvedApiServer,
+      branch: opts?.branch ?? process.env._CHALK_BRANCH ?? undefined,
       clientId: valueWithEnvFallback(
         "clientId",
         opts?.clientId,
@@ -157,7 +172,7 @@ export class ChalkClient<TFeatureMap = Record<string, ChalkScalar>>
   async whoami(): Promise<ChalkWhoamiResponse> {
     return this.http.v1_who_am_i({
       baseUrl: this.config.apiServer,
-      headers: this.getDefaultHeaders(),
+      headers: this.getHeaders(),
       credentials: this.credentials,
     });
   }
@@ -168,7 +183,7 @@ export class ChalkClient<TFeatureMap = Record<string, ChalkScalar>>
       pathParams: {
         run_id: runId,
       },
-      headers: this.getDefaultHeaders(),
+      headers: this.getHeaders(),
       credentials: this.credentials,
     });
   }
@@ -181,7 +196,7 @@ export class ChalkClient<TFeatureMap = Record<string, ChalkScalar>>
       body: {
         resolver_fqn: request.resolverFqn,
       },
-      headers: this.getDefaultHeaders(),
+      headers: this.getHeaders(),
       credentials: this.credentials,
     });
   }
@@ -213,7 +228,7 @@ export class ChalkClient<TFeatureMap = Record<string, ChalkScalar>>
         include_meta: !!request.include_meta,
         planner_options: request.plannerOptions,
       },
-      headers: this.getDefaultHeaders(),
+      headers: this.getHeaders(requestOptions),
       credentials: this.credentials,
       timeout: requestOptions?.timeout,
     });
@@ -250,7 +265,7 @@ export class ChalkClient<TFeatureMap = Record<string, ChalkScalar>>
     const rawResult = await this.http.v1_query_feather({
       baseUrl: this.config.queryServer,
       body: requestBuffer.buffer,
-      headers: this.getDefaultHeaders(),
+      headers: this.getHeaders(requestOptions),
       credentials: this.credentials,
       timeout: requestOptions?.timeout,
     });
@@ -286,7 +301,7 @@ export class ChalkClient<TFeatureMap = Record<string, ChalkScalar>>
     const rawResult = await this.http.v1_query_feather({
       baseUrl: this.config.queryServer,
       body: requestBuffer.buffer,
-      headers: this.getDefaultHeaders(),
+      headers: this.getHeaders(requestOptions),
       credentials: this.credentials,
       timeout: requestOptions?.timeout,
     });
@@ -312,7 +327,7 @@ export class ChalkClient<TFeatureMap = Record<string, ChalkScalar>>
         correlation_id: request.correlationId,
         deployment_id: request.previewDeploymentId,
       },
-      headers: this.getDefaultHeaders(),
+      headers: this.getHeaders(),
       credentials: this.credentials,
     });
 
@@ -324,10 +339,17 @@ export class ChalkClient<TFeatureMap = Record<string, ChalkScalar>>
     }
   }
 
-  private getDefaultHeaders(): ChalkHttpHeaders {
-    return {
+  private getHeaders(requestOptions?: ChalkRequestOptions): ChalkHttpHeaders {
+    const headers: ChalkHttpHeaders = {
       "X-Chalk-Env-Id": this.config.activeEnvironment,
       "User-Agent": "chalk-ts v1.17.0",
     };
+
+    const branch = requestOptions?.branch ?? this.config.branch;
+    if (branch != null) {
+      headers["X-Chalk-Branch-Id"] = branch;
+    }
+
+    return headers;
   }
 }

src/_http.ts

@@ -7,6 +7,7 @@ import { ChalkErrorCategory, ChalkErrorCode } from "./_interface";
 
 export interface ChalkHttpHeaders {
   "X-Chalk-Env-Id"?: string;
+  "X-Chalk-Branch-Id"?: string;
   "User-Agent"?: string;
 }
 
@@ -239,6 +240,10 @@ export class ChalkHTTPService {
         headers.set("X-Chalk-Timeout", effectiveTimeout.toString());
       }
 
+      if (callArgs.headers?.["X-Chalk-Branch-Id"] != null) {
+        headers.set("X-Chalk-Branch-Id", callArgs.headers["X-Chalk-Branch-Id"]);
+      }
+
       const body =
         callArgs.body !== undefined
           ? !opts.binaryResponseBody

src/_types.ts

@@ -1,9 +1,9 @@
 import { TimestampFormat } from "./_interface";
 
-
 export interface ChalkClientConfig {
   activeEnvironment: string | undefined;
   apiServer: string;
+  branch: string | undefined;
   clientId: string;
   clientSecret: string;
   queryServer: string;
@@ -24,6 +24,7 @@ export interface ChalkEnvironmentVariables {
   _CHALK_API_SERVER: string;
   _CHALK_QUERY_SERVER: string;
   _CHALK_ACTIVE_ENVIRONMENT: string;
+  _CHALK_BRANCH: string;
 }
 
 /**