@ai-sdk/sandbox-vercel 1.0.71 → 1.0.72

package/src/vercel-network-sandbox-session.ts

@@ -2,9 +2,12 @@ import {
   HarnessCapabilityUnsupportedError,
   type HarnessV1NetworkPolicy,
   type HarnessV1NetworkSandboxSession,
+  type HarnessV1PortEndpoint,
+  type HarnessV1RequestTransformation,
 } from '@ai-sdk/harness';
 import type { Experimental_SandboxSession as SandboxSession } from '@ai-sdk/provider-utils';
-import type { Sandbox, NetworkPolicy } from '@vercel/sandbox';
+import type { Sandbox } from '@vercel/sandbox';
+import { VercelNetworkPolicyManager } from './vercel-network-policy-manager';
 import { VercelSandboxSession } from './vercel-sandbox-session';
 
 const VERCEL_PROVIDER_ID = 'vercel-sandbox';
@@ -24,12 +27,16 @@ export class VercelNetworkSandboxSession
   readonly id: string;
   readonly defaultWorkingDirectory: string;
   private readonly ownsLifecycle: boolean;
+  readonly #networkPolicyManager: VercelNetworkPolicyManager;
 
   constructor(input: { sandbox: Sandbox; ownsLifecycle: boolean }) {
     super(input.sandbox);
     this.ownsLifecycle = input.ownsLifecycle;
     this.id = input.sandbox.name;
     this.defaultWorkingDirectory = input.sandbox.currentSession().cwd;
+    this.#networkPolicyManager = new VercelNetworkPolicyManager({
+      sandbox: input.sandbox,
+    });
   }
 
   get ports(): ReadonlyArray<number> {
@@ -40,10 +47,10 @@ export class VercelNetworkSandboxSession
     return new VercelSandboxSession(this.sandbox);
   }
 
-  getPortUrl = async (options: {
+  getPortEndpoint = async (options: {
     port: number;
     protocol?: 'http' | 'https' | 'ws';
-  }): Promise<string> => {
+  }): Promise<HarnessV1PortEndpoint> => {
     const exposedPorts = this.ports;
     if (!exposedPorts.includes(options.port)) {
       throw new HarnessCapabilityUnsupportedError({
@@ -65,11 +72,33 @@ export class VercelNetworkSandboxSession
         url.protocol = isSecure ? 'wss:' : 'ws:';
         break;
     }
-    return url.toString();
+    return { url: url.toString() };
+  };
+
+  /**
+   * @deprecated Use `getPortEndpoint` instead.
+   */
+  getPortUrl = async (options: {
+    port: number;
+    protocol?: 'http' | 'https' | 'ws';
+  }): Promise<string> => {
+    return (await this.getPortEndpoint(options)).url;
   };
 
   setNetworkPolicy = async (policy: HarnessV1NetworkPolicy): Promise<void> => {
-    await this.sandbox.update({ networkPolicy: toVercelPolicy(policy) });
+    await this.#networkPolicyManager.setNetworkPolicy(policy);
+  };
+
+  setRequestTransformations = async (
+    transformations: ReadonlyArray<HarnessV1RequestTransformation>,
+  ): Promise<void> => {
+    await this.#networkPolicyManager.setRequestTransformations(transformations);
+  };
+
+  addRequestTransformations = async (
+    transformations: ReadonlyArray<HarnessV1RequestTransformation>,
+  ): Promise<void> => {
+    await this.#networkPolicyManager.addRequestTransformations(transformations);
   };
 
   setPorts = async (
@@ -93,40 +122,3 @@ export class VercelNetworkSandboxSession
     await this.sandbox.delete();
   };
 }
-
-export function toVercelPolicy(policy: HarnessV1NetworkPolicy): NetworkPolicy {
-  switch (policy.mode) {
-    case 'allow-all':
-      return 'allow-all';
-    case 'deny-all':
-      return 'deny-all';
-    case 'custom': {
-      const result: Extract<NetworkPolicy, { allow?: unknown }> = {};
-      const { allowedHosts, allowedCIDRs, deniedCIDRs } = policy;
-      if (allowedHosts != null && allowedHosts.length > 0) {
-        result.allow = [...allowedHosts];
-      }
-      if (
-        (allowedCIDRs != null && allowedCIDRs.length > 0) ||
-        (deniedCIDRs != null && deniedCIDRs.length > 0)
-      ) {
-        result.subnets = {
-          ...(allowedCIDRs != null && allowedCIDRs.length > 0
-            ? { allow: [...allowedCIDRs] }
-            : {}),
-          ...(deniedCIDRs != null && deniedCIDRs.length > 0
-            ? { deny: [...deniedCIDRs] }
-            : {}),
-        };
-      }
-      if (result.allow == null && result.subnets == null) {
-        throw new HarnessCapabilityUnsupportedError({
-          harnessId: VERCEL_PROVIDER_ID,
-          message:
-            'Custom network policy requires at least one of allowedHosts or allowedCIDRs to be non-empty.',
-        });
-      }
-      return result;
-    }
-  }
-}