Address PR review comments

Author: rculbertson

modal-go/client.go

@@ -76,9 +76,12 @@ var clientProfile Profile
 // client is the default Modal client that talks to the control plane.
 var client pb.ModalClientClient
 
-// clients is a map of server URL => client.
-// The us-east client talks to the control plane; all other clients talk to input planes.
-var clients = map[string]pb.ModalClientClient{}
+// clients is a map of server URL to input-plane client.
+var inputPlaneClients = map[string]pb.ModalClientClient{}
+
+// authToken is the auth token received from the control plane on the first request, and sent with all
+// subsequent requests to both the control plane and the input plane.
+var authToken string
 
 func init() {
 	defaultConfig, _ = readConfigFile()
@@ -113,17 +116,19 @@ func InitializeClient(options ClientOptions) error {
 	return err
 }
 
-// getOrCreateClient returns a client for the given server URL, creating it if it doesn't exist.
-func getOrCreateClient(serverURL string) (pb.ModalClientClient, error) {
-	if client, ok := clients[serverURL]; ok {
+// getOrCreateInputPlaneClient returns a client for the given server URL, creating it if it doesn't exist.
+func getOrCreateInputPlaneClient(serverURL string) (pb.ModalClientClient, error) {
+	if client, ok := inputPlaneClients[serverURL]; ok {
 		return client, nil
 	}
 
-	_, client, err := newClient(nil)
+	profile := clientProfile
+	profile.ServerURL = serverURL
+	_, client, err := newClient(profile)
 	if err != nil {
 		return nil, err
 	}
-	clients[serverURL] = client
+	inputPlaneClients[serverURL] = client
 	return client, nil
 }
 
@@ -150,6 +155,7 @@ func newClient(profile Profile) (*grpc.ClientConn, pb.ModalClientClient, error)
 			grpc.MaxCallSendMsgSize(maxMessageSize),
 		),
 		grpc.WithChainUnaryInterceptor(
+			authTokenInterceptor(),
 			retryInterceptor(),
 			timeoutInterceptor(),
 		),
@@ -176,6 +182,37 @@ func clientContext(ctx context.Context) (context.Context, error) {
 	), nil
 }
 
+// authTokenInterceptor handles sending and receiving the "x-modal-auth-token" header.
+// We receive an auth token from the control plane on our first request. We then include that auth token in every
+// subsequent request to both the control plane and the input plane.
+func authTokenInterceptor() grpc.UnaryClientInterceptor {
+	return func(
+		ctx context.Context,
+		method string,
+		req, reply any,
+		cc *grpc.ClientConn,
+		inv grpc.UnaryInvoker,
+		opts ...grpc.CallOption,
+	) error {
+		var headers, trailers metadata.MD
+		// Add authToken to outgoing context if it's set
+		if authToken != "" {
+			ctx = metadata.AppendToOutgoingContext(ctx, "x-modal-auth-token", authToken)
+		}
+		opts = append(opts, grpc.Header(&headers), grpc.Trailer(&trailers))
+		err := inv(ctx, method, req, reply, cc, opts...)
+		// If we're talking to the control plane, and no auth token was sent, it will return one.
+		// The python server returns it in the trailers, the worker returns it in the headers.
+		if val, ok := headers["x-modal-auth-token"]; ok {
+			authToken = val[0]
+		} else if val, ok := trailers["x-modal-auth-token"]; ok {
+			authToken = val[0]
+		}
+
+		return err
+	}
+}
+
 func timeoutInterceptor() grpc.UnaryClientInterceptor {
 	return func(
 		ctx context.Context,

modal-go/function.go

@@ -10,8 +10,6 @@ import (
 	"encoding/base64"
 	"errors"
 	"fmt"
-	"google.golang.org/grpc"
-	"google.golang.org/grpc/metadata"
 	"net/http"
 	"time"
 
@@ -53,13 +51,12 @@ func FunctionLookup(ctx context.Context, appName string, name string, options *L
 		return nil, err
 	}
 
-	var header, trailer metadata.MD
 	resp, err := client.FunctionGet(ctx, pb.FunctionGetRequest_builder{
 		AppName:         appName,
 		ObjectTag:       name,
 		Namespace:       pb.DeploymentNamespace_DEPLOYMENT_NAMESPACE_WORKSPACE,
 		EnvironmentName: environmentName(options.Environment),
-	}.Build(), grpc.Header(&header), grpc.Trailer(&trailer))
+	}.Build())
 
 	if status, ok := status.FromError(err); ok && status.Code() == codes.NotFound {
 		return nil, NotFoundError{fmt.Sprintf("function '%s/%s' not found", appName, name)}
@@ -68,15 +65,6 @@ func FunctionLookup(ctx context.Context, appName string, name string, options *L
 		return nil, err
 	}
 
-	// Attach x-modal-auth-token to all future requests.
-	authTokenArray := header.Get("x-modal-auth-token")
-	if len(authTokenArray) == 0 {
-		authTokenArray = trailer.Get("x-modal-auth-token")
-	}
-	if len(authTokenArray) > 0 {
-		authToken := authTokenArray[0]
-		ctx = metadata.AppendToOutgoingContext(ctx, "x-modal-auth-token", authToken)
-	}
 	var inputPlaneUrl *string
 	if meta := resp.GetHandleMetadata(); meta != nil {
 		if url := meta.GetInputPlaneUrl(); url != "" {

modal-go/invocation.go

@@ -16,9 +16,6 @@ type invocation interface {
 	retry(retryCount uint32) error
 }
 
-// getOutput is a function type that takes a timeout and returns a FunctionGetOutputsItem or nil, and an error.
-type getOutput func(timeout time.Duration) (*pb.FunctionGetOutputsItem, error)
-
 // controlPlaneInvocation implements the invocation interface.
 type controlPlaneInvocation struct {
 	FunctionCallId  string
@@ -118,7 +115,7 @@ func createInputPlaneInvocation(ctx context.Context, inputPlaneURL string, funct
 		Idx:   0,
 		Input: input,
 	}.Build()
-	client, err := getOrCreateClient(inputPlaneURL)
+	client, err := getOrCreateInputPlaneClient(inputPlaneURL)
 	if err != nil {
 		return nil, err
 	}
@@ -171,7 +168,13 @@ func (i *inputPlaneInvocation) retry(retryCount uint32) error {
 	return nil
 }
 
-// pollFunctionOutput repeatedly fetches the output for a given function call, waiting until a result is available or a timeout occurs.
+// getOutput is a function type that takes a timeout and returns a FunctionGetOutputsItem or nil, and an error.
+// Used by `pollForOutputs` to fetch from either the control plane or the input plane, depending on the implementation.
+type getOutput func(timeout time.Duration) (*pb.FunctionGetOutputsItem, error)
+
+// pollFunctionOutput repeatedly tries to fetch an output using the provided `getOutput` function, and the specified
+// timeout value. We use a timeout value of 55 seconds if the caller does not specify a timeout value, or if the
+// specified timeout value is greater than 55 seconds.
 func pollFunctionOutput(ctx context.Context, getOutput getOutput, timeout *time.Duration) (any, error) {
 	startTime := time.Now()
 	pollTimeout := outputsTimeout

modal-js/src/client.ts

@@ -16,7 +16,7 @@ import { getProfile, type Profile } from "./config";
 const defaultProfile = getProfile(process.env["MODAL_PROFILE"]);
 
 // CLIENT VERSION: Behaves like this Python SDK version
-const clientVersion = "1.0.0"
+const clientVersion = "1.0.0";
 
 let modalAuthToken: string | undefined;
 
@@ -45,46 +45,29 @@ function authMiddleware(profile: Profile): ClientMiddleware {
       options.metadata.set("x-modal-auth-token", modalAuthToken);
     }
 
+    // We receive an auth token from the control plane on our first request. We then include that auth token in every
+    // subsequent request to both the control plane and the input plane. The python server returns it in the trailers,
+    // the worker returns it in the headers.
     const prevOnHeader = options.onHeader;
     options.onHeader = (header) => {
       const token = header.get("x-modal-auth-token");
       if (token) {
         modalAuthToken = token;
       }
-      if (prevOnHeader) {
-        prevOnHeader(header);
-      }
+      prevOnHeader?.(header);
     };
     const prevOnTrailer = options.onTrailer;
     options.onTrailer = (trailer) => {
       const token = trailer.get("x-modal-auth-token");
       if (token) {
         modalAuthToken = token;
       }
-      if (prevOnTrailer) {
-        prevOnTrailer(trailer);
-      }
+      prevOnTrailer?.(trailer);
     };
     return yield* call.next(call.request, options);
   };
 }
 
-/** gRPC client middleware to add auth token to request. */
-function inputPlaneAuthMiddleware(): ClientMiddleware {
-  return async function* inputPlaneAuthMiddleware<Request, Response>(
-    call: ClientMiddlewareCall<Request, Response>,
-    options: CallOptions,
-  ) {
-    options.metadata ??= new Metadata();
-    options.metadata.set("x-modal-client-type", String(ClientType.CLIENT_TYPE_LIBMODAL));
-    options.metadata.set("x-modal-client-version", clientVersion);
-    if (modalAuthToken) {
-      options.metadata.set("x-modal-auth-token", modalAuthToken);
-    }
-    return yield* call.next(call.request, options);
-  };
-}
-
 type TimeoutOptions = {
   /** Timeout for this call, interpreted as a duration in milliseconds */
   timeout?: number;
@@ -248,35 +231,27 @@ const retryMiddleware: ClientMiddleware<RetryOptions> =
 const inputPlaneClients: Record<string, ReturnType<typeof createClient>> = {};
 
 /** Returns a client for the given server URL, creating it if it doesn't exist. */
-export const getOrCreateClient = (
+export const getOrCreateInputPlaneClient = (
   serverUrl: string,
 ): ReturnType<typeof createClient> => {
   const client = inputPlaneClients[serverUrl];
   if (client) {
     return client;
   }
-  const channel = createClientChannel(serverUrl);
-  const newClient = createClientFactory()
-    .use(inputPlaneAuthMiddleware())
-    .use(retryMiddleware)
-    .use(timeoutMiddleware)
-    .create(ModalClientDefinition, channel);
+  const profile = { ...clientProfile, serverUrl };
+  const newClient = createClient(profile);
   inputPlaneClients[serverUrl] = newClient;
   return newClient;
 };
 
-function createClientChannel(serverUrl: string) {
+function createClient(profile: Profile) {
   // Channels don't do anything until you send a request on them.
   // Ref: https://github.com/modal-labs/modal-client/blob/main/modal/_utils/grpc_utils.py
-  return createChannel(serverUrl, undefined, {
+  const channel = createChannel(profile.serverUrl, undefined, {
     "grpc.max_receive_message_length": 100 * 1024 * 1024,
     "grpc.max_send_message_length": 100 * 1024 * 1024,
     "grpc-node.flow_control_window": 64 * 1024 * 1024,
   });
-}
-
-function createClient(profile: Profile) {
-  const channel = createClientChannel(profile.serverUrl);
   return createClientFactory()
     .use(authMiddleware(profile))
     .use(retryMiddleware)

modal-js/src/invocation.ts

@@ -11,7 +11,7 @@ import {
   GenericResult_GenericStatus,
   ModalClientClient,
 } from "../proto/modal_proto/api";
-import { client, getOrCreateClient } from "./client";
+import { client, getOrCreateInputPlaneClient } from "./client";
 import { FunctionTimeoutError, InternalFailure, RemoteError } from "./errors";
 import { loads } from "./pickle";
 
@@ -29,14 +29,6 @@ export interface Invocation {
   retry(retryCount: number): Promise<void>;
 }
 
-/**
- * Signature of a function that fetches a single output. Used by `pollForOutputs` to fetch from either
- * the control plane or the input plane, depending on the implementation.
- */
-type GetOutput = (
-  timeoutMillis: number,
-) => Promise<FunctionGetOutputsItem | undefined>;
-
 /**
  * Implementation of Invocation which sends inputs to the control plane.
  */
@@ -97,19 +89,15 @@ export class ControlPlaneInvocation implements Invocation {
   async #getOutput(
     timeoutMillis: number,
   ): Promise<FunctionGetOutputsItem | undefined> {
-    try {
-      const response = await client.functionGetOutputs({
-        functionCallId: this.functionCallId,
-        maxValues: 1,
-        timeout: timeoutMillis / 1000, // Backend needs seconds
-        lastEntryId: "0-0",
-        clearOnSuccess: true,
-        requestedAt: timeNowSeconds(),
-      });
-      return response.outputs ? response.outputs[0] : undefined;
-    } catch (err) {
-      throw new Error(`FunctionGetOutputs failed: ${err}`);
-    }
+    const response = await client.functionGetOutputs({
+      functionCallId: this.functionCallId,
+      maxValues: 1,
+      timeout: timeoutMillis / 1000, // Backend needs seconds
+      lastEntryId: "0-0",
+      clearOnSuccess: true,
+      requestedAt: timeNowSeconds(),
+    });
+    return response.outputs ? response.outputs[0] : undefined;
   }
 
   async retry(retryCount: number): Promise<void> {
@@ -160,12 +148,12 @@ export class InputPlaneInvocation implements Invocation {
   ) {
     const functionPutInputsItem = {
       idx: 0,
-      input: input,
+      input,
     };
-    const client = getOrCreateClient(inputPlaneUrl);
+    const client = getOrCreateInputPlaneClient(inputPlaneUrl);
     // Single input sync invocation
     const attemptStartResponse = await client.attemptStart({
-      functionId: functionId,
+      functionId,
       input: functionPutInputsItem,
     });
     return new InputPlaneInvocation(
@@ -212,6 +200,19 @@ function timeNowSeconds() {
   return Date.now() / 1e3;
 }
 
+/**
+ * Signature of a function that fetches a single output using the given timeout. Used by `pollForOutputs` to fetch
+ * from either the control plane or the input plane, depending on the implementation.
+ */
+type GetOutput = (
+  timeoutMillis: number,
+) => Promise<FunctionGetOutputsItem | undefined>;
+
+/***
+ * Repeatedly tries to fetch an output using the provided `getOutput` function, and the specified timeout value.
+ * We use a timeout value of 55 seconds if the caller does not specify a timeout value, or if the specified timeout
+ * value is greater than 55 seconds.
+ */
 async function pollFunctionOutput(
   getOutput: GetOutput,
   timeout?: number, // in milliseconds

test-support/libmodal_test_support.py

@@ -19,10 +19,12 @@ def sleep(t: int) -> None:
 def bytelength(buf: bytes) -> int:
     return len(buf)
 
+
 @app.function(min_containers=1, experimental_options={"input_plane_region": "us-west"})
 def input_plane(s: str) -> str:
     return "output: " + s
 
+
 @app.cls(min_containers=1)
 class EchoCls:
     @modal.method()