Add test and fixes

Author: rculbertson

modal-js/src/client.ts

@@ -13,6 +13,8 @@ import {
 import { ClientType, ModalClientDefinition } from "../proto/modal_proto/api";
 import { type Profile, profile } from "./config";
 
+let modalAuthToken: string | undefined;
+
 /** gRPC client middleware to add auth token to request. */
 function authMiddleware(profile: Profile): ClientMiddleware {
   return async function* authMiddleware<Request, Response>(
@@ -27,6 +29,30 @@ function authMiddleware(profile: Profile): ClientMiddleware {
     options.metadata.set("x-modal-client-version", "1.0.0"); // CLIENT VERSION: Behaves like this Python SDK version
     options.metadata.set("x-modal-token-id", profile.tokenId);
     options.metadata.set("x-modal-token-secret", profile.tokenSecret);
+    if (modalAuthToken) {
+      options.metadata.set("x-modal-auth-token", modalAuthToken);
+    }
+
+    const prevOnHeader = options.onHeader;
+    options.onHeader = (header) => {
+      const token = header.get("x-modal-auth-token");
+      if (token) {
+        modalAuthToken = token;
+      }
+      if (prevOnHeader) {
+        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);
+      }
+    };
     return yield* call.next(call.request, options);
   };
 }
@@ -223,4 +249,5 @@ const createClient = (serverURL: string) => {
     .create(ModalClientDefinition, channel);
 };
 
+/** The default Modal client that talks to the control plane. */
 export const client = getOrCreateClient(profile.serverUrl);

modal-js/src/function.ts

@@ -2,11 +2,7 @@
 
 import { createHash } from "node:crypto";
 
-import {
-  DataFormat,
-  DeploymentNamespace,
-  FunctionPutInputsItem,
-} from "../proto/modal_proto/api";
+import { DataFormat, DeploymentNamespace } from "../proto/modal_proto/api";
 import type { LookupOptions } from "./app";
 import { client } from "./client";
 import { FunctionCall } from "./function_call";
@@ -18,6 +14,8 @@ import {
   ControlPlaneStrategy,
   GetOutput,
   getOutputsFromControlPlane,
+  InputPlaneStrategy,
+  InvocationStrategy,
   pollForOutputs,
 } from "./invocation_strategy";
 
@@ -28,11 +26,13 @@ const maxObjectSizeBytes = 2 * 1024 * 1024; // 2 MiB
 export class Function_ {
   readonly functionId: string;
   readonly methodName: string | undefined;
+  private readonly inputPlaneUrl: string | undefined;
 
   /** @ignore */
-  constructor(functionId: string, methodName?: string) {
+  constructor(functionId: string, methodName?: string, inputPlaneUrl?: string) {
     this.functionId = functionId;
     this.methodName = methodName;
+    this.inputPlaneUrl = inputPlaneUrl;
   }
 
   static async lookup(
@@ -47,7 +47,11 @@ export class Function_ {
         namespace: DeploymentNamespace.DEPLOYMENT_NAMESPACE_WORKSPACE,
         environmentName: environmentName(options.environment),
       });
-      return new Function_(resp.functionId);
+      return new Function_(
+        resp.functionId,
+        undefined,
+        resp.handleMetadata?.inputPlaneUrl,
+      );
     } catch (err) {
       if (err instanceof ClientError && err.code === Status.NOT_FOUND)
         throw new NotFoundError(`Function '${appName}/${name}' not found`);
@@ -60,8 +64,10 @@ export class Function_ {
     args: any[] = [],
     kwargs: Record<string, any> = {},
   ): Promise<any> {
-    const input = await this.#createInput(args, kwargs);
-    const invocationStrategy = new ControlPlaneStrategy(this.functionId, input);
+    const invocationStrategy = await this.#createInvocationStrategy(
+      args,
+      kwargs,
+    );
     return await invocationStrategy.remote();
   }
 
@@ -70,16 +76,18 @@ export class Function_ {
     args: any[] = [],
     kwargs: Record<string, any> = {},
   ): Promise<FunctionCall> {
-    const input = await this.#createInput(args, kwargs);
-    const invocationStrategy = new ControlPlaneStrategy(this.functionId, input);
+    const invocationStrategy = await this.#createInvocationStrategy(
+      args,
+      kwargs,
+    );
     const functionCallId = await invocationStrategy.spawn();
     return new FunctionCall(functionCallId);
   }
 
-  async #createInput(
+  async #createInvocationStrategy(
     args: any[] = [],
     kwargs: Record<string, any> = {},
-  ): Promise<FunctionPutInputsItem> {
+  ): Promise<InvocationStrategy> {
     const payload = dumps([args, kwargs]);
 
     let argsBlobId: string | undefined = undefined;
@@ -88,7 +96,7 @@ export class Function_ {
     }
 
     // Single input sync invocation
-    return {
+    const input = {
       idx: 0,
       input: {
         args: argsBlobId ? undefined : payload,
@@ -98,6 +106,11 @@ export class Function_ {
         finalInput: false, // This field isn't specified in the Python client, so it defaults to false.
       },
     };
+    if (this.inputPlaneUrl) {
+      return new InputPlaneStrategy(this.functionId, this.inputPlaneUrl, input);
+    }
+
+    return new ControlPlaneStrategy(this.functionId, input);
   }
 }
 

modal-js/src/invocation_strategy.ts

@@ -35,6 +35,10 @@ export interface InvocationStrategy {
   spawn(): Promise<string>;
 }
 
+/**
+ * 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.
+ */
 export type GetOutput = (
   timeoutMillis: number,
 ) => Promise<FunctionGetOutputsItem | undefined>;
@@ -91,19 +95,27 @@ export class ControlPlaneStrategy implements InvocationStrategy {
 export class InputPlaneStrategy implements InvocationStrategy {
   private readonly functionId: string;
   private readonly client: ModalClientClient;
+  private readonly input: FunctionPutInputsItem;
   private attemptToken: string | undefined;
 
-  constructor(functionId: string, inputPlaneUrl: string) {
+  constructor(
+    functionId: string,
+    inputPlaneUrl: string,
+    input: FunctionPutInputsItem,
+  ) {
     this.client = getOrCreateClient(inputPlaneUrl);
     this.functionId = functionId;
+    this.input = input;
   }
 
-  async remote(input: FunctionPutInputsItem): Promise<FunctionGetOutputsItem> {
-    await this.#attemptStart(input);
-    return await pollForOutputs(this.#attemptAwait);
+  async remote(): Promise<FunctionGetOutputsItem> {
+    await this.#attemptStart(this.input);
+    const getOutput: GetOutput = (timeoutMillis: number) =>
+      this.#attemptAwait(timeoutMillis);
+    return await pollForOutputs(getOutput);
   }
 
-  async spawn(_input: FunctionPutInputsItem): Promise<string> {
+  async spawn(): Promise<string> {
     throw new Error("Spawn operations are not supported by the input plane.");
   }
 
@@ -136,7 +148,7 @@ function timeNowSeconds() {
 }
 
 export async function pollForOutputs(
-  outputFetcher: GetOutput,
+  getOutput: GetOutput,
   timeout?: number, // in milliseconds
 ): Promise<any> {
   const startTime = Date.now();
@@ -146,7 +158,7 @@ export async function pollForOutputs(
   }
 
   while (true) {
-    const output = await outputFetcher(pollTimeout);
+    const output = await getOutput(pollTimeout);
     if (output) {
       return await processResult(output.result, output.dataFormat);
     }

modal-js/test/function.test.ts

@@ -34,3 +34,12 @@ test("FunctionNotFound", async () => {
   );
   await expect(promise).rejects.toThrowError(NotFoundError);
 });
+
+test("FunctionCallInputPlane", async () => {
+  const function_ = await Function_.lookup(
+    "libmodal-test-support",
+    "input_plane",
+  );
+  const result = await function_.remote(["hello"]);
+  expect(result).toBe("output: hello");
+});

test-support/libmodal_test_support.py

@@ -34,3 +34,7 @@ class EchoClsParametrized:
     @modal.method()
     def echo_parameter(self) -> str:
         return "output: " + self.name
+
+@app.function(min_containers=1, experimental_options={"input_plane_region": "us-west"})
+def input_plane(s: str) -> str:
+    return "output: " + s