*: move to grpc backed communication for the agent

.github/workflows/build.yml

@@ -8,6 +8,16 @@ jobs:
       - uses: actions/setup-node@v3
         with:
           node-version: '16'
+      - uses: bufbuild/buf-setup-action@v1
+        with:
+          github_token: ${{ github.token }}
+      - name: Configure npm for buf registry
+        env:
+          BUF_TOKEN: ${{ secrets.BUF_TOKEN }}
+        run: |
+          npm config set @buf:registry https://buf.build/gen/npm/v1/
+          npm config set //buf.build/gen/npm/v1/:_authToken $BUF_TOKEN
       - run: npm ci
+      - run: npm install @buf/blacksmith_vm-agent.connectrpc_es@latest
       - run: npm run build
-      - run: npm test 
+      - run: npm test

.github/workflows/verify-build.yml

@@ -21,9 +21,21 @@ jobs:
           node-version: '16'
           cache: 'npm'
 
+      - uses: bufbuild/buf-setup-action@v1
+        with:
+          github_token: ${{ github.token }}
+
+      - name: Configure npm for buf registry
+        env:
+          BUF_TOKEN: ${{ secrets.BUF_TOKEN }}
+        run: |
+          npm config set @buf:registry https://buf.build/gen/npm/v1/
+          npm config set //buf.build/gen/npm/v1/:_authToken $BUF_TOKEN
+
       - name: Install dependencies
         run: |
           npm ci
+          npm install @buf/blacksmith_vm-agent.connectrpc_es@latest
 
       - name: Build
         run: npm run build

dist/licenses.txt

@@ -394,6 +394,19 @@ OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
 SOFTWARE.
 
 
+@buf/blacksmith_vm-agent.bufbuild_es
+
+@buf/blacksmith_vm-agent.connectrpc_es
+
+@bufbuild/protobuf
+(Apache-2.0 AND BSD-3-Clause)
+
+@connectrpc/connect
+Apache-2.0
+
+@connectrpc/connect-node
+Apache-2.0
+
 @docker/actions-toolkit
 Apache-2.0
 

package-lock.json

@@ -8,6 +8,9 @@
       "license": "Apache-2.0",
       "dependencies": {
         "@actions/core": "^1.10.1",
+        "@buf/blacksmith_vm-agent.connectrpc_es": "^1.6.1-20241213043610-906584953dd9.2",
+        "@connectrpc/connect": "^1.6.1",
+        "@connectrpc/connect-node": "^1.6.1",
         "@docker/actions-toolkit": "0.37.1",
         "@iarna/toml": "^2.2.5",
         "axios-retry": "^4.5.0",
@@ -1052,6 +1055,52 @@
       "integrity": "sha512-0hYQ8SB4Db5zvZB4axdMHGwEaQjkZzFjQiN9LVYvIFB2nSUHW9tYpxWriPrWDASIxiaXax83REcLxuSdnGPZtw==",
       "dev": true
     },
+    "node_modules/@buf/blacksmith_vm-agent.bufbuild_es": {
+      "version": "1.10.0-20241213043610-906584953dd9.1",
+      "resolved": "https://buf.build/gen/npm/v1/@buf/blacksmith_vm-agent.bufbuild_es/-/blacksmith_vm-agent.bufbuild_es-1.10.0-20241213043610-906584953dd9.1.tgz",
+      "peerDependencies": {
+        "@bufbuild/protobuf": "^1.10.0"
+      }
+    },
+    "node_modules/@buf/blacksmith_vm-agent.connectrpc_es": {
+      "version": "1.6.1-20241213043610-906584953dd9.2",
+      "resolved": "https://buf.build/gen/npm/v1/@buf/blacksmith_vm-agent.connectrpc_es/-/blacksmith_vm-agent.connectrpc_es-1.6.1-20241213043610-906584953dd9.2.tgz",
+      "dependencies": {
+        "@buf/blacksmith_vm-agent.bufbuild_es": "1.10.0-20241213043610-906584953dd9.1"
+      },
+      "peerDependencies": {
+        "@connectrpc/connect": "^1.6.1"
+      }
+    },
+    "node_modules/@bufbuild/protobuf": {
+      "version": "1.10.0",
+      "resolved": "https://registry.npmjs.org/@bufbuild/protobuf/-/protobuf-1.10.0.tgz",
+      "integrity": "sha512-QDdVFLoN93Zjg36NoQPZfsVH9tZew7wKDKyV5qRdj8ntT4wQCOradQjRaTdwMhWUYsgKsvCINKKm87FdEk96Ag==",
+      "peer": true
+    },
+    "node_modules/@connectrpc/connect": {
+      "version": "1.6.1",
+      "resolved": "https://registry.npmjs.org/@connectrpc/connect/-/connect-1.6.1.tgz",
+      "integrity": "sha512-KchMDNtU4CDTdkyf0qG7ugJ6qHTOR/aI7XebYn3OTCNagaDYWiZUVKgRgwH79yeMkpNgvEUaXSK7wKjaBK9b/Q==",
+      "peerDependencies": {
+        "@bufbuild/protobuf": "^1.10.0"
+      }
+    },
+    "node_modules/@connectrpc/connect-node": {
+      "version": "1.6.1",
+      "resolved": "https://registry.npmjs.org/@connectrpc/connect-node/-/connect-node-1.6.1.tgz",
+      "integrity": "sha512-DxcD1wsF/aX9GegjAtl7VbpiZNjVJozy87VbaFoN6AF0Ln1Q757r5dgV59Gz0wmlk5f17txUsrEr1f2inlnnAg==",
+      "dependencies": {
+        "undici": "^5.28.4"
+      },
+      "engines": {
+        "node": ">=16.0.0"
+      },
+      "peerDependencies": {
+        "@bufbuild/protobuf": "^1.10.0",
+        "@connectrpc/connect": "1.6.1"
+      }
+    },
     "node_modules/@cspotcode/source-map-support": {
       "version": "0.8.1",
       "resolved": "https://registry.npmjs.org/@cspotcode/source-map-support/-/source-map-support-0.8.1.tgz",

package.json

@@ -27,6 +27,9 @@
   "packageManager": "yarn@3.6.3",
   "dependencies": {
     "@actions/core": "^1.10.1",
+    "@buf/blacksmith_vm-agent.connectrpc_es": "^1.6.1-20241213043610-906584953dd9.2",
+    "@connectrpc/connect": "^1.6.1",
+    "@connectrpc/connect-node": "^1.6.1",
     "@docker/actions-toolkit": "0.37.1",
     "@iarna/toml": "^2.2.5",
     "axios-retry": "^4.5.0",

src/__tests__/setup-builder.test.ts

@@ -1,54 +0,0 @@
-import * as reporter from '../reporter';
-import { getStickyDisk } from '../setup_builder';
-import FormData from 'form-data';
-
-jest.mock('../reporter');
-
-describe('getStickyDisk', () => {
-  const mockGet = jest.fn();
-
-  beforeEach(() => {
-    jest.resetAllMocks();
-    process.env.GITHUB_REPO_NAME = 'test-repo';
-    process.env.BLACKSMITH_REGION = 'test-region';
-    process.env.BLACKSMITH_INSTALLATION_MODEL_ID = 'test-model';
-    process.env.VM_ID = 'test-vm';
-
-    (reporter.createBlacksmithAgentClient as jest.Mock).mockResolvedValue({});
-    (reporter.get as jest.Mock).mockImplementation(mockGet);
-    mockGet.mockResolvedValue({
-      data: {
-        expose_id: 'test-expose-id',
-        disk_identifier: 'test-device'
-      }
-    });
-  });
-
-  it('sets both FormData and query parameters correctly', async () => {
-    const appendSpy = jest.spyOn(FormData.prototype, 'append');
-    
-    await getStickyDisk();
-
-    expect(mockGet).toHaveBeenCalledTimes(1);
-    const [, url, formData] = mockGet.mock.calls[0];
-
-    // Verify query parameters
-    expect(url).toContain('stickyDiskKey=test-repo');
-    expect(url).toContain('region=test-region');
-    expect(url).toContain('installationModelID=test-model');
-    expect(url).toContain('vmID=test-vm');
-
-    // Verify FormData is correct type
-    expect(formData instanceof FormData).toBeTruthy();
-    
-    // Verify the headers are set correctly
-    const headers = formData.getHeaders();
-    expect(headers['content-type']).toContain('multipart/form-data');
-
-    // Verify the correct fields were appended
-    expect(appendSpy).toHaveBeenCalledWith('stickyDiskKey', 'test-repo');
-    expect(appendSpy).toHaveBeenCalledWith('region', 'test-region');
-    expect(appendSpy).toHaveBeenCalledWith('installationModelID', 'test-model');
-    expect(appendSpy).toHaveBeenCalledWith('vmID', 'test-vm');
-  });
-}); 

src/reporter.ts

@@ -1,8 +1,11 @@
 import * as core from '@actions/core';
-import axios, {AxiosError, AxiosInstance, AxiosResponse, AxiosStatic } from 'axios';
+import axios, {AxiosError, AxiosInstance, AxiosResponse } from 'axios';
 import axiosRetry from 'axios-retry';
 import {ExportRecordResponse} from '@docker/actions-toolkit/lib/types/buildx/history';
 import FormData from 'form-data';
+import { createClient } from "@connectrpc/connect";
+import { createGrpcTransport } from "@connectrpc/connect-node";
+import { StickyDiskService } from "@buf/blacksmith_vm-agent.connectrpc_es/stickydisk/v1/stickydisk_connect";
 
 // Configure base axios instance for Blacksmith API.
 const createBlacksmithAPIClient = () => {
@@ -31,26 +34,13 @@ const createBlacksmithAPIClient = () => {
   return client;
 };
 
-export async function createBlacksmithAgentClient(): Promise<AxiosInstance> {
+export function createBlacksmithAgentClient() {
-  const stickyDiskMgrUrl = 'http://192.168.127.1:5556';
+  const transport = createGrpcTransport({
-  const client = axios.create({
+    baseUrl: 'http://192.168.127.1:5557',
-    baseURL: stickyDiskMgrUrl,
+    httpVersion: '2',
-    headers: {
-      Authorization: `Bearer ${process.env.BLACKSMITH_STICKYDISK_TOKEN}`,
-      'X-Github-Repo-Name': process.env.GITHUB_REPO_NAME || '',
-    }
-  });
-
-  axiosRetry(client, {
-    retries: 5,
-    retryDelay: axiosRetry.exponentialDelay,
-    retryCondition: (error) => {
-      return axiosRetry.isNetworkOrIdempotentRequestError(error) || 
-             (error.response?.status ? error.response.status >= 500 : false);
-    }
   });
 
-  return client;
+  return createClient(StickyDiskService, transport);
 }
 
 export async function reportBuildPushActionFailure(error?: Error) {
@@ -77,13 +67,15 @@ export async function reportBuildCompleted(exportRes?: ExportRecordResponse, bla
 
   try {
     const agentClient = await createBlacksmithAgentClient();
-    const formData = new FormData();
+    
-    formData.append('shouldCommit', 'true');
+    await agentClient.commitStickyDisk({
-    formData.append('vmID', process.env.VM_ID || '');
+      exposeId: exposeId || '',
-    formData.append('exposeID', exposeId || '');
+      stickyDiskKey: process.env.GITHUB_REPO_NAME || '',
-    formData.append('stickyDiskKey', process.env.GITHUB_REPO_NAME || '');
+      vmId: process.env.VM_ID || '',
-
+      shouldCommit: true,
-    await post(agentClient, '/stickydisks', formData);
+      repoName: process.env.GITHUB_REPO_NAME || '',
+      stickyDiskToken: process.env.BLACKSMITH_STICKYDISK_TOKEN || ''
+    });
 
     // Report success to Blacksmith API
     const requestOptions = {
@@ -126,13 +118,14 @@ export async function reportBuildFailed(dockerBuildId: string | null, dockerBuil
 
   try {
     const blacksmithAgentClient = await createBlacksmithAgentClient();
-    const formData = new FormData();
+    await blacksmithAgentClient.commitStickyDisk({
-    formData.append('shouldCommit', 'false');
+      exposeId: exposeId || '',
-    formData.append('vmID', process.env.VM_ID || '');
+      stickyDiskKey: process.env.GITHUB_REPO_NAME || '',
-    formData.append('exposeID', exposeId || '');
+      vmId: process.env.VM_ID || '',
-    formData.append('stickyDiskKey', process.env.GITHUB_REPO_NAME || '');
+      shouldCommit: false,
-
+      repoName: process.env.GITHUB_REPO_NAME || '',
-    await post(blacksmithAgentClient, '/stickydisks', formData);
+      stickyDiskToken: process.env.BLACKSMITH_STICKYDISK_TOKEN || ''
+    });
 
     // Report failure to Blacksmith API
     const requestOptions = {

src/setup_builder.ts

@@ -147,38 +147,28 @@ async function getDiskSize(device: string): Promise<number> {
 export async function getStickyDisk(options?: {signal?: AbortSignal}): Promise<{expose_id: string; device: string}> {
   const client = await reporter.createBlacksmithAgentClient();
   
-  // Prepare data for both FormData and query params
   const stickyDiskKey = process.env.GITHUB_REPO_NAME || '';
   if (stickyDiskKey === '') {
     throw new Error('GITHUB_REPO_NAME is not set');
   }
-  const region = process.env.BLACKSMITH_REGION || 'eu-central';
-  const installationModelID = process.env.BLACKSMITH_INSTALLATION_MODEL_ID || '';
-  const vmID = process.env.VM_ID || '';
-
-  // Create FormData (for backwards compatibility).
-  // TODO(adityamaru): Remove this once all of our VM agents are reading query params.
-  const formData = new FormData();
-  formData.append('stickyDiskKey', stickyDiskKey);
-  formData.append('region', region);
-  formData.append('installationModelID', installationModelID);
-  formData.append('vmID', vmID);
-
-  // Create query params string.
-  const queryParams = new URLSearchParams({
-    stickyDiskKey,
-    region,
-    installationModelID,
-    vmID
-  }).toString();
-
   core.debug(`Getting sticky disk for ${stickyDiskKey}`);
 
-  // Send request with both FormData and query params
+
-  const response = await reporter.get(client, `/stickydisks?${queryParams}`, formData, options);
+  const response = await client.getStickyDisk({
-  const exposeId = response.data?.expose_id || '';
+    stickyDiskKey: stickyDiskKey,
-  const device = response.data?.disk_identifier || '';
+    region: process.env.BLACKSMITH_REGION || 'eu-central',
-  return {expose_id: exposeId, device: device};
+    installationModelId: process.env.BLACKSMITH_INSTALLATION_MODEL_ID || '',
+    vmId: process.env.VM_ID || '',
+    stickyDiskType: 'dockerfile',
+    repoName: process.env.GITHUB_REPO_NAME || '',
+    stickyDiskToken: process.env.BLACKSMITH_STICKYDISK_TOKEN || ''
+  }, {
+    signal: options?.signal
+  });
+  return {
+    expose_id: response.exposeId || '',
+    device: response.diskIdentifier || ''
+  };
 }
 
 export async function startAndConfigureBuildkitd(parallelism: number, device: string): Promise<string> {