fix(device-agent): unbrick stuck installs (auto-update + legacy session)

fix(device-agent): unbrick stuck installs (auto-update + legacy session)

Author: Marfuen

apps/api/src/device-agent/device-agent.controller.spec.ts

@@ -8,6 +8,8 @@ import { PermissionGuard } from '../auth/permission.guard';
 import type { AuthContext as AuthContextType } from '../auth/types';
 import { Readable } from 'stream';
 
+jest.mock('@db', () => ({ db: {} }));
+
 jest.mock('../auth/auth.server', () => ({
   auth: { api: { getSession: jest.fn() } },
 }));
@@ -26,6 +28,8 @@ describe('DeviceAgentController', () => {
   const mockService = {
     downloadMacAgent: jest.fn(),
     downloadWindowsAgent: jest.fn(),
+    getUpdateFile: jest.fn(),
+    headUpdateFile: jest.fn(),
   };
 
   const mockGuard = { canActivate: jest.fn().mockReturnValue(true) };
@@ -42,6 +46,7 @@ describe('DeviceAgentController', () => {
 
   const mockRes = {
     set: jest.fn(),
+    redirect: jest.fn(),
   };
 
   beforeEach(async () => {
@@ -155,4 +160,97 @@ describe('DeviceAgentController', () => {
       ).rejects.toThrow('Agent not found');
     });
   });
+
+  describe('getUpdateFile', () => {
+    it('streams a yml manifest with cache headers', async () => {
+      const mockStream = Readable.from(Buffer.from('version: 1.0.5'));
+      mockService.getUpdateFile.mockResolvedValue({
+        kind: 'stream',
+        stream: mockStream,
+        contentType: 'text/yaml',
+        contentLength: 14,
+      });
+
+      const result = await controller.getUpdateFile(
+        'latest-mac.yml',
+        mockRes as never,
+      );
+
+      expect(result).toBeInstanceOf(StreamableFile);
+      expect(mockRes.set).toHaveBeenCalledWith(
+        expect.objectContaining({
+          'Content-Type': 'text/yaml',
+          'Cache-Control': 'public, max-age=300',
+          'Content-Length': '14',
+        }),
+      );
+      expect(mockRes.redirect).not.toHaveBeenCalled();
+    });
+
+    it('issues a 302 redirect to the presigned URL for binaries', async () => {
+      mockService.getUpdateFile.mockResolvedValue({
+        kind: 'redirect',
+        url: 'https://s3.example.com/signed-zip',
+      });
+
+      const result = await controller.getUpdateFile(
+        'CompAI-Device-Agent-1.0.5-arm64.zip',
+        mockRes as never,
+      );
+
+      expect(mockRes.redirect).toHaveBeenCalledWith(
+        302,
+        'https://s3.example.com/signed-zip',
+      );
+      expect(mockRes.set).toHaveBeenCalledWith(
+        expect.objectContaining({
+          'Cache-Control': 'no-store',
+        }),
+      );
+      expect(result).toBeUndefined();
+    });
+  });
+
+  describe('headUpdateFile', () => {
+    it('returns metadata headers for a yml manifest', async () => {
+      mockService.headUpdateFile.mockResolvedValue({
+        kind: 'stream',
+        contentType: 'text/yaml',
+        contentLength: 859,
+      });
+
+      const result = await controller.headUpdateFile(
+        'latest-mac.yml',
+        mockRes as never,
+      );
+
+      expect(mockRes.set).toHaveBeenCalledWith(
+        expect.objectContaining({
+          'Content-Type': 'text/yaml',
+          'Cache-Control': 'public, max-age=300',
+          'Content-Length': '859',
+        }),
+      );
+      expect(result).toBe('');
+      expect(mockRes.redirect).not.toHaveBeenCalled();
+    });
+
+    it('issues a 302 redirect for HEAD on binaries', async () => {
+      mockService.headUpdateFile.mockResolvedValue({
+        kind: 'redirect',
+        url: 'https://s3.example.com/signed-head',
+      });
+
+      const result = await controller.headUpdateFile(
+        'CompAI-Device-Agent-1.0.5-arm64.zip',
+        mockRes as never,
+      );
+
+      expect(mockRes.redirect).toHaveBeenCalledWith(
+        302,
+        'https://s3.example.com/signed-head',
+      );
+      expect(result).toBeUndefined();
+    });
+  });
 });

apps/api/src/device-agent/device-agent.controller.ts

@@ -70,6 +70,12 @@ export class DeviceAgentController {
   ) {
     const result = await this.deviceAgentService.getUpdateFile({ filename });
 
+    if (result.kind === 'redirect') {
+      res.set({ 'Cache-Control': 'no-store' });
+      res.redirect(302, result.url);
+      return;
+    }
+
     res.set({
       'Content-Type': result.contentType,
       'Cache-Control': 'public, max-age=300',
@@ -90,6 +96,12 @@ export class DeviceAgentController {
   ) {
     const result = await this.deviceAgentService.headUpdateFile({ filename });
 
+    if (result.kind === 'redirect') {
+      res.set({ 'Cache-Control': 'no-store' });
+      res.redirect(302, result.url);
+      return;
+    }
+
     res.set({
       'Content-Type': result.contentType,
       'Cache-Control': 'public, max-age=300',

apps/api/src/device-agent/device-agent.service.spec.ts

@@ -5,10 +5,28 @@ import {
 import { Readable } from 'stream';
 
 const mockSend = jest.fn();
+const mockGetSignedUrl = jest.fn();
+
+class MockGetObjectCommand {
+  constructor(public readonly input: unknown) {
+    Object.assign(this, input as object);
+  }
+}
+
+class MockHeadObjectCommand {
+  constructor(public readonly input: unknown) {
+    Object.assign(this, input as object);
+  }
+}
 
 jest.mock('@aws-sdk/client-s3', () => ({
   S3Client: jest.fn().mockImplementation(() => ({ send: mockSend })),
-  GetObjectCommand: jest.fn().mockImplementation((input) => input),
+  GetObjectCommand: MockGetObjectCommand,
+  HeadObjectCommand: MockHeadObjectCommand,
+}));
+
+jest.mock('@aws-sdk/s3-request-presigner', () => ({
+  getSignedUrl: (...args: unknown[]) => mockGetSignedUrl(...args),
 }));
 
 import { DeviceAgentService } from './device-agent.service';
@@ -92,6 +110,123 @@ describe('DeviceAgentService', () => {
     });
   });
 
+  describe('getUpdateFile', () => {
+    it('streams .yml manifests directly from S3', async () => {
+      const mockStream = new Readable({ read() {} });
+      mockSend.mockResolvedValue({
+        Body: mockStream,
+        ContentLength: 859,
+      });
+
+      const result = await service.getUpdateFile({ filename: 'latest-mac.yml' });
+
+      expect(result).toEqual({
+        kind: 'stream',
+        stream: mockStream,
+        contentType: 'text/yaml',
+        contentLength: 859,
+      });
+      expect(mockGetSignedUrl).not.toHaveBeenCalled();
+    });
+
+    it('redirects binary downloads to a presigned S3 URL signed for GET', async () => {
+      mockGetSignedUrl.mockResolvedValue('https://s3.example.com/signed-zip-url');
+
+      const result = await service.getUpdateFile({
+        filename: 'CompAI-Device-Agent-1.0.5-arm64.zip',
+      });
+
+      expect(result).toEqual({
+        kind: 'redirect',
+        url: 'https://s3.example.com/signed-zip-url',
+      });
+      expect(mockSend).not.toHaveBeenCalled();
+      expect(mockGetSignedUrl).toHaveBeenCalledTimes(1);
+      const [, command] = mockGetSignedUrl.mock.calls[0];
+      expect(command).toBeInstanceOf(MockGetObjectCommand);
+      expect(command).toMatchObject({
+        Bucket: 'test-bucket',
+        Key: 'device-agent/production/updates/CompAI-Device-Agent-1.0.5-arm64.zip',
+      });
+    });
+
+    it.each([
+      'CompAI-Device-Agent-1.0.5-arm64.zip',
+      'CompAI-Device-Agent-1.0.5-setup.exe',
+      'CompAI-Device-Agent-1.0.5-arm64.dmg',
+      'CompAI-Device-Agent-1.0.5-x86_64.AppImage',
+      'CompAI-Device-Agent-1.0.5-arm64.zip.blockmap',
+    ])('redirects binary file %s', async (filename) => {
+      mockGetSignedUrl.mockResolvedValue('https://s3.example.com/signed');
+
+      const result = await service.getUpdateFile({ filename });
+
+      expect(result).toEqual({
+        kind: 'redirect',
+        url: 'https://s3.example.com/signed',
+      });
+    });
+
+    it('throws NotFoundException for invalid filenames', async () => {
+      await expect(
+        service.getUpdateFile({ filename: '../etc/passwd' }),
+      ).rejects.toThrow(NotFoundException);
+      await expect(
+        service.getUpdateFile({ filename: 'foo.txt' }),
+      ).rejects.toThrow(NotFoundException);
+    });
+
+    it('throws NotFoundException when S3 returns NoSuchKey for a yml manifest', async () => {
+      const error = new Error('Not found');
+      error.name = 'NoSuchKey';
+      mockSend.mockRejectedValue(error);
+
+      await expect(
+        service.getUpdateFile({ filename: 'latest-mac.yml' }),
+      ).rejects.toThrow(NotFoundException);
+    });
+  });
+
+  describe('headUpdateFile', () => {
+    it('returns metadata for .yml manifests', async () => {
+      mockSend.mockResolvedValue({ ContentLength: 859 });
+
+      const result = await service.headUpdateFile({
+        filename: 'latest-mac.yml',
+      });
+
+      expect(result).toEqual({
+        kind: 'stream',
+        contentType: 'text/yaml',
+        contentLength: 859,
+      });
+      expect(mockGetSignedUrl).not.toHaveBeenCalled();
+    });
+
+    it('redirects binary HEAD requests to a URL signed with HeadObjectCommand', async () => {
+      mockGetSignedUrl.mockResolvedValue('https://s3.example.com/signed-head');
+
+      const result = await service.headUpdateFile({
+        filename: 'CompAI-Device-Agent-1.0.5-arm64.zip',
+      });
+
+      expect(result).toEqual({
+        kind: 'redirect',
+        url: 'https://s3.example.com/signed-head',
+      });
+      expect(mockSend).not.toHaveBeenCalled();
+      expect(mockGetSignedUrl).toHaveBeenCalledTimes(1);
+      const [, command] = mockGetSignedUrl.mock.calls[0];
+      // S3 signs each HTTP method separately; a GET-signed URL would be
+      // rejected when used with a HEAD request.
+      expect(command).toBeInstanceOf(MockHeadObjectCommand);
+      expect(command).toMatchObject({
+        Bucket: 'test-bucket',
+        Key: 'device-agent/production/updates/CompAI-Device-Agent-1.0.5-arm64.zip',
+      });
+    });
+  });
+
   describe('downloadWindowsAgent', () => {
     it('should return stream, filename, and contentType on success', async () => {
       const mockStream = new Readable({ read() {} });