feat(cloud): upload workflow code through cloud storage API (#938)

Co-authored-by: Proactive Runtime Bot <agent@agent-relay.com>

Author: khaliqgant

packages/cloud/src/workflows.test.ts

@@ -425,6 +425,74 @@ describe('runWorkflow code sync', () => {
     });
     expect((runBodies[0] as { paths?: unknown }).paths).toBeUndefined();
   });
+
+  it('uploads code through the cloud API when prepare returns cloud-api storage', async () => {
+    await writeFile('README.md', 'cloud-api\n');
+    const workflowPath = path.join(tmpRoot, 'workflow.yaml');
+    await writeFile(
+      workflowPath,
+      ['version: "1.0"', 'name: cloud-api', 'swarm:', '  pattern: dag', 'agents: []', 'workflows: []'].join(
+        '\n'
+      )
+    );
+
+    const runBodies: unknown[] = [];
+    const uploadPaths: string[] = [];
+    authorizedApiFetchMock.mockImplementation(async (_auth, requestPath, init) => {
+      if (requestPath === '/api/v1/workflows/prepare') {
+        return {
+          auth: { accessToken: 'token' },
+          response: new Response(
+            JSON.stringify({
+              runId: 'run-1',
+              s3Credentials: {
+                ...s3Credentials,
+                backend: 'cloud-api',
+                cloudApiUrl: 'https://agentrelay.com/cloud',
+                cloudApiAccessToken: 'token',
+              },
+              s3CodeKey: 'code.tar.gz',
+              workflowStorage: { backend: 'cloud-api' },
+            }),
+            { status: 200, headers: { 'Content-Type': 'application/json' } }
+          ),
+        };
+      }
+      if (requestPath === '/api/v1/workflows/runs/run-1/storage/code.tar.gz') {
+        uploadPaths.push(requestPath);
+        expect(init?.method).toBe('PUT');
+        expect(init?.headers).toMatchObject({ 'content-type': 'application/gzip' });
+        expect(init?.body).toBeInstanceOf(Buffer);
+        return {
+          auth: { accessToken: 'token' },
+          response: new Response(JSON.stringify({ ok: true }), {
+            status: 200,
+            headers: { 'Content-Type': 'application/json' },
+          }),
+        };
+      }
+      if (requestPath === '/api/v1/workflows/run') {
+        runBodies.push(JSON.parse(String(init?.body)));
+        return {
+          auth: { accessToken: 'token' },
+          response: new Response(JSON.stringify({ runId: 'run-1', status: 'pending' }), {
+            status: 200,
+            headers: { 'Content-Type': 'application/json' },
+          }),
+        };
+      }
+      throw new Error(`unexpected request: ${requestPath}`);
+    });
+
+    await runWorkflow(workflowPath);
+
+    expect(s3SendMock).not.toHaveBeenCalled();
+    expect(uploadPaths).toEqual(['/api/v1/workflows/runs/run-1/storage/code.tar.gz']);
+    expect(runBodies[0]).toMatchObject({
+      runId: 'run-1',
+      s3CodeKey: 'code.tar.gz',
+    });
+  });
 });
 
 describe('workflow schedules', () => {

packages/cloud/src/workflows.ts

@@ -25,17 +25,23 @@ type ResolvedWorkflowInput = {
 };
 
 type S3Credentials = {
+  backend?: 's3' | 'cloud-api';
   accessKeyId: string;
   secretAccessKey: string;
   sessionToken: string;
   bucket: string;
   prefix: string;
+  cloudApiUrl?: string;
+  cloudApiAccessToken?: string;
 };
 
 type PrepareWorkflowResponse = {
   runId: string;
   s3Credentials: S3Credentials;
   s3CodeKey: string;
+  workflowStorage?: {
+    backend?: 's3' | 'cloud-api';
+  };
 };
 
 type WorkflowPathDefinition = {
@@ -525,7 +531,40 @@ export async function runWorkflow(
     const prepared = prepPayload;
     console.error(`  Prepared in ${((Date.now() - t0) / 1000).toFixed(1)}s`);
 
-    const s3Client = createScopedS3Client(prepared.s3Credentials);
+    let s3Client: S3Client | null = null;
+    const uploadCodeObject = async (objectKey: string, tarball: Buffer) => {
+      if (isCloudApiWorkflowStorage(prepared)) {
+        const { response, auth: uploadAuth } = await authorizedApiFetch(
+          auth,
+          workflowStorageObjectPath(prepared.runId, objectKey),
+          {
+            method: 'PUT',
+            headers: {
+              'content-type': 'application/gzip',
+              accept: 'application/json',
+            },
+            body: tarball as unknown as BodyInit,
+          }
+        );
+        auth = uploadAuth;
+        const payload = await readJsonResponse(response);
+        if (!response.ok) {
+          throw new Error(`Workflow storage upload failed: ${describeResponseError(response, payload)}`);
+        }
+        return;
+      }
+
+      s3Client ??= createScopedS3Client(prepared.s3Credentials);
+      const key = scopedCodeKey(prepared.s3Credentials.prefix, objectKey);
+      await s3Client.send(
+        new PutObjectCommand({
+          Bucket: prepared.s3Credentials.bucket,
+          Key: key,
+          Body: tarball,
+          ContentType: 'application/gzip',
+        })
+      );
+    };
     requestBody.runId = prepared.runId;
 
     const declaredPaths = parseWorkflowPaths(input.workflow, input.fileType);
@@ -552,15 +591,7 @@ export async function runWorkflow(
         );
 
         const t2 = Date.now();
-        const key = scopedCodeKey(prepared.s3Credentials.prefix, s3CodeKey);
-        await s3Client.send(
-          new PutObjectCommand({
-            Bucket: prepared.s3Credentials.bucket,
-            Key: key,
-            Body: tarball,
-            ContentType: 'application/gzip',
-          })
-        );
+        await uploadCodeObject(s3CodeKey, tarball);
         console.error(`  ${pathDef.name}: uploaded in ${((Date.now() - t2) / 1000).toFixed(1)}s`);
 
         const repo = parseGitHubRemoteForPath(absolutePath);
@@ -595,16 +626,8 @@ export async function runWorkflow(
       );
 
       const t2 = Date.now();
-      console.error('Uploading to S3...');
-      const key = scopedCodeKey(prepared.s3Credentials.prefix, prepared.s3CodeKey);
-      await s3Client.send(
-        new PutObjectCommand({
-          Bucket: prepared.s3Credentials.bucket,
-          Key: key,
-          Body: tarball,
-          ContentType: 'application/gzip',
-        })
-      );
+      console.error('Uploading to workflow storage...');
+      await uploadCodeObject(prepared.s3CodeKey, tarball);
       console.error(`  Uploaded in ${((Date.now() - t2) / 1000).toFixed(1)}s`);
 
       requestBody.s3CodeKey = prepared.s3CodeKey;
@@ -981,10 +1004,21 @@ function isPrepareWorkflowResponse(payload: unknown): payload is PrepareWorkflow
     typeof creds.secretAccessKey === 'string' &&
     typeof creds.sessionToken === 'string' &&
     typeof creds.bucket === 'string' &&
-    typeof creds.prefix === 'string'
+    typeof creds.prefix === 'string' &&
+    (creds.backend === undefined || creds.backend === 's3' || creds.backend === 'cloud-api')
   );
 }
 
+function isCloudApiWorkflowStorage(prepared: PrepareWorkflowResponse): boolean {
+  return prepared.workflowStorage?.backend === 'cloud-api' || prepared.s3Credentials.backend === 'cloud-api';
+}
+
+function workflowStorageObjectPath(runId: string, objectKey: string): string {
+  const encodedRunId = encodeURIComponent(runId);
+  const encodedKey = objectKey.split('/').map(encodeURIComponent).join('/');
+  return `/api/v1/workflows/runs/${encodedRunId}/storage/${encodedKey}`;
+}
+
 function createScopedS3Client(s3Credentials: S3Credentials): S3Client {
   return new S3Client({
     region: process.env.AWS_REGION ?? process.env.AWS_DEFAULT_REGION ?? 'us-east-1',