add size limit on workspace persistence

pgayvallet · pgayvallet · commit 6beb802913b0 · 2026-05-28T22:48:53.000+02:00
diff --git a/x-pack/platform/plugins/shared/agent_builder/server/services/execution/filesystem/capacity_limited_fs.test.ts b/x-pack/platform/plugins/shared/agent_builder/server/services/execution/filesystem/capacity_limited_fs.test.ts
@@ -0,0 +1,111 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+import { InMemoryFs } from 'just-bash';
+import { CapacityLimitedFs } from './capacity_limited_fs';
+
+const make = (cap: number) => {
+  const inner = new InMemoryFs();
+  return { inner, fs: new CapacityLimitedFs(inner, cap) };
+};
+
+describe('CapacityLimitedFs', () => {
+  describe('writeFile', () => {
+    it('allows writes under the cap', async () => {
+      const { fs } = make(100);
+      await fs.writeFile('/a.txt', 'x'.repeat(50));
+      await fs.writeFile('/b.txt', 'y'.repeat(40));
+      expect(await fs.readFile('/a.txt')).toBe('x'.repeat(50));
+      expect(await fs.readFile('/b.txt')).toBe('y'.repeat(40));
+    });
+
+    it('rejects a write that would exceed the cap', async () => {
+      const { fs } = make(100);
+      await fs.writeFile('/a.txt', 'x'.repeat(60));
+      await expect(fs.writeFile('/b.txt', 'y'.repeat(50))).rejects.toThrow(/ENOSPC/);
+    });
+
+    it('allows overwriting an existing file when the new content fits', async () => {
+      const { fs } = make(100);
+      await fs.writeFile('/a.txt', 'x'.repeat(60));
+      // 60 - 60 + 80 = 80 ≤ 100
+      await fs.writeFile('/a.txt', 'y'.repeat(80));
+      expect((await fs.readFile('/a.txt')).length).toBe(80);
+    });
+
+    it('rejects an overwrite that would still exceed the cap', async () => {
+      const { fs } = make(100);
+      await fs.writeFile('/a.txt', 'x'.repeat(60));
+      await fs.writeFile('/b.txt', 'y'.repeat(30));
+      // current 90, overwrite /a from 60→120 = projected 150 > 100
+      await expect(fs.writeFile('/a.txt', 'x'.repeat(120))).rejects.toThrow(/ENOSPC/);
+    });
+
+    it('reflects rm() freeing capacity', async () => {
+      const { fs } = make(100);
+      await fs.writeFile('/a.txt', 'x'.repeat(90));
+      await expect(fs.writeFile('/b.txt', 'y'.repeat(20))).rejects.toThrow(/ENOSPC/);
+      await fs.rm('/a.txt');
+      await fs.writeFile('/b.txt', 'y'.repeat(20));
+      expect((await fs.readFile('/b.txt')).length).toBe(20);
+    });
+  });
+
+  describe('appendFile', () => {
+    it('rejects an append that would exceed the cap', async () => {
+      const { fs } = make(100);
+      await fs.writeFile('/a.txt', 'x'.repeat(80));
+      await expect(fs.appendFile('/a.txt', 'y'.repeat(30))).rejects.toThrow(/ENOSPC/);
+    });
+
+    it('allows an append that fits', async () => {
+      const { fs } = make(100);
+      await fs.writeFile('/a.txt', 'x'.repeat(50));
+      await fs.appendFile('/a.txt', 'y'.repeat(40));
+      expect((await fs.readFile('/a.txt')).length).toBe(90);
+    });
+  });
+
+  describe('cp', () => {
+    it('rejects a copy that would exceed the cap', async () => {
+      const { fs } = make(100);
+      await fs.writeFile('/src.txt', 'x'.repeat(60));
+      await fs.writeFile('/other.txt', 'y'.repeat(30));
+      // current 90, cp adds 60 (no existing dest) = 150 > 100
+      await expect(fs.cp('/src.txt', '/dst.txt')).rejects.toThrow(/ENOSPC/);
+    });
+
+    it('allows a copy that fits', async () => {
+      const { fs } = make(100);
+      await fs.writeFile('/src.txt', 'x'.repeat(30));
+      await fs.cp('/src.txt', '/dst.txt');
+      expect(await fs.readFile('/dst.txt')).toBe('x'.repeat(30));
+    });
+  });
+
+  describe('mv', () => {
+    it('passes through — net-neutral in size', async () => {
+      const { fs } = make(100);
+      await fs.writeFile('/a.txt', 'x'.repeat(60));
+      await fs.writeFile('/b.txt', 'y'.repeat(35));
+      // total 95; mv doesn't change total, so allow.
+      await fs.mv('/a.txt', '/c.txt');
+      expect(await fs.exists('/a.txt')).toBe(false);
+      expect(await fs.readFile('/c.txt')).toBe('x'.repeat(60));
+    });
+  });
+
+  describe('reads pass through', () => {
+    it('readFile / stat / readdir delegate to inner', async () => {
+      const { fs } = make(100);
+      await fs.writeFile('/a.txt', 'hello');
+      expect(await fs.readFile('/a.txt')).toBe('hello');
+      expect((await fs.stat('/a.txt')).size).toBe(5);
+      expect(await fs.readdir('/')).toContain('a.txt');
+    });
+  });
+});
diff --git a/x-pack/platform/plugins/shared/agent_builder/server/services/execution/filesystem/capacity_limited_fs.ts b/x-pack/platform/plugins/shared/agent_builder/server/services/execution/filesystem/capacity_limited_fs.ts
@@ -0,0 +1,178 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+import type { ByteString, IFileSystem, FsStat } from 'just-bash';
+
+const encoder = new TextEncoder();
+
+const byteLengthOf = (content: string | Uint8Array): number =>
+  typeof content === 'string' ? encoder.encode(content).length : content.byteLength;
+
+const enospc = (op: string, path: string, limit: number) =>
+  new Error(
+    `ENOSPC: workspace capacity exceeded — ${op} '${path}' would exceed the ${limit}-byte limit`
+  );
+
+/**
+ * Caps the total content size of an underlying writable `IFileSystem`. Writes
+ * that would push the aggregate above `maxBytes` are rejected with `ENOSPC`.
+ * Reads and non-growing operations (rm, mkdir, chmod, ...) pass through.
+ *
+ * The current size is computed by walking the fs on each mutation. That's O(n)
+ * per write; acceptable for the workspace's tens-to-hundreds of small files at
+ * worst. Revisit (cache + invalidate) if walk cost becomes measurable.
+ */
+export class CapacityLimitedFs implements IFileSystem {
+  private readonly inner: IFileSystem;
+  private readonly maxBytes: number;
+
+  constructor(inner: IFileSystem, maxBytes: number) {
+    this.inner = inner;
+    this.maxBytes = maxBytes;
+  }
+
+  // ---- mutation paths with capacity enforcement ----
+
+  async writeFile(
+    path: string,
+    content: string | Uint8Array,
+    options?: Parameters<IFileSystem['writeFile']>[2]
+  ): Promise<void> {
+    const newBytes = byteLengthOf(content);
+    const existingBytes = await this.tryGetFileSize(path);
+    const projected = (await this.currentSize()) - existingBytes + newBytes;
+    if (projected > this.maxBytes) throw enospc('writeFile', path, this.maxBytes);
+    return this.inner.writeFile(path, content, options);
+  }
+
+  async appendFile(
+    path: string,
+    content: string | Uint8Array,
+    options?: Parameters<IFileSystem['appendFile']>[2]
+  ): Promise<void> {
+    const addedBytes = byteLengthOf(content);
+    const projected = (await this.currentSize()) + addedBytes;
+    if (projected > this.maxBytes) throw enospc('appendFile', path, this.maxBytes);
+    return this.inner.appendFile(path, content, options);
+  }
+
+  async cp(src: string, dest: string, options?: Parameters<IFileSystem['cp']>[2]): Promise<void> {
+    const srcBytes = await this.subtreeSize(src);
+    const existingDestBytes = await this.subtreeSize(dest);
+    const projected = (await this.currentSize()) - existingDestBytes + srcBytes;
+    if (projected > this.maxBytes) throw enospc('cp', dest, this.maxBytes);
+    return this.inner.cp(src, dest, options);
+  }
+
+  // ---- pass-through (no growth or freeing) ----
+
+  async readFile(path: string, options?: Parameters<IFileSystem['readFile']>[1]): Promise<string> {
+    return this.inner.readFile(path, options);
+  }
+  async readFileBuffer(path: string): Promise<Uint8Array> {
+    return this.inner.readFileBuffer(path);
+  }
+  async readFileBytes(path: string): Promise<ByteString> {
+    return this.inner.readFileBytes!(path);
+  }
+  async exists(path: string): Promise<boolean> {
+    return this.inner.exists(path);
+  }
+  async stat(path: string): Promise<FsStat> {
+    return this.inner.stat(path);
+  }
+  async lstat(path: string): Promise<FsStat> {
+    return this.inner.lstat(path);
+  }
+  async readdir(path: string): Promise<string[]> {
+    return this.inner.readdir(path);
+  }
+  async readdirWithFileTypes(
+    path: string
+  ): Promise<
+    Array<{ name: string; isFile: boolean; isDirectory: boolean; isSymbolicLink: boolean }>
+  > {
+    return this.inner.readdirWithFileTypes!(path);
+  }
+  getAllPaths(): string[] {
+    return this.inner.getAllPaths();
+  }
+  resolvePath(base: string, path: string): string {
+    return this.inner.resolvePath(base, path);
+  }
+  async realpath(path: string): Promise<string> {
+    return this.inner.realpath(path);
+  }
+  async readlink(path: string): Promise<string> {
+    return this.inner.readlink(path);
+  }
+  async mkdir(path: string, options?: Parameters<IFileSystem['mkdir']>[1]): Promise<void> {
+    return this.inner.mkdir(path, options);
+  }
+  async rm(path: string, options?: Parameters<IFileSystem['rm']>[1]): Promise<void> {
+    return this.inner.rm(path, options);
+  }
+  async mv(src: string, dest: string): Promise<void> {
+    // Net-neutral within the same fs: bytes move from src to dest, total unchanged.
+    return this.inner.mv(src, dest);
+  }
+  async chmod(path: string, mode: number): Promise<void> {
+    return this.inner.chmod(path, mode);
+  }
+  async symlink(target: string, linkPath: string): Promise<void> {
+    return this.inner.symlink(target, linkPath);
+  }
+  async link(existingPath: string, newPath: string): Promise<void> {
+    return this.inner.link(existingPath, newPath);
+  }
+  async utimes(path: string, atime: Date, mtime: Date): Promise<void> {
+    return this.inner.utimes(path, atime, mtime);
+  }
+
+  // ---- size helpers ----
+
+  /** Total bytes used by all files in the fs. Directories don't count. */
+  private async currentSize(): Promise<number> {
+    return this.subtreeSize('/');
+  }
+
+  /** Recursively sums file sizes under `path`. Returns 0 if `path` doesn't exist. */
+  private async subtreeSize(path: string): Promise<number> {
+    if (!(await this.inner.exists(path))) return 0;
+    try {
+      const stat = await this.inner.stat(path);
+      if (stat.isFile) return stat.size;
+    } catch {
+      return 0;
+    }
+    let total = 0;
+    const walk = async (dir: string): Promise<void> => {
+      const entries = await this.inner.readdirWithFileTypes!(dir);
+      for (const entry of entries) {
+        const child = dir === '/' ? `/${entry.name}` : `${dir}/${entry.name}`;
+        if (entry.isFile) {
+          const s = await this.inner.stat(child);
+          total += s.size;
+        } else if (entry.isDirectory) {
+          await walk(child);
+        }
+      }
+    };
+    await walk(path);
+    return total;
+  }
+
+  private async tryGetFileSize(path: string): Promise<number> {
+    if (!(await this.inner.exists(path))) return 0;
+    try {
+      const stat = await this.inner.stat(path);
+      return stat.isFile ? stat.size : 0;
+    } catch {
+      return 0;
+    }
+  }
+}
diff --git a/x-pack/platform/plugins/shared/agent_builder/server/services/execution/filesystem/index.ts b/x-pack/platform/plugins/shared/agent_builder/server/services/execution/filesystem/index.ts
@@ -7,4 +7,9 @@
 
 export { FilesystemService, type FilesystemServiceDeps } from './filesystem_service';
 export { VolumeBackedReadOnlyFs } from './volume_backed_read_only_fs';
-export { WorkspaceVolume, type WorkspaceVolumeDeps } from './workspace_volume';
+export {
+  WorkspaceVolume,
+  type WorkspaceVolumeDeps,
+  DEFAULT_WORKSPACE_CAPACITY_BYTES,
+} from './workspace_volume';
+export { CapacityLimitedFs } from './capacity_limited_fs';
diff --git a/x-pack/platform/plugins/shared/agent_builder/server/services/execution/filesystem/workspace_volume.test.ts b/x-pack/platform/plugins/shared/agent_builder/server/services/execution/filesystem/workspace_volume.test.ts
@@ -101,6 +101,42 @@ describe('WorkspaceVolume', () => {
     });
   });
 
+  describe('capacity cap', () => {
+    it('writes through getFilesystem() throw ENOSPC when the cap is exceeded', async () => {
+      const v = new WorkspaceVolume({
+        workspaceClient: mockWorkspaceClient(),
+        capacityBytes: 50,
+      });
+      const fs = v.getFilesystem();
+      await fs.writeFile('/note.txt', 'x'.repeat(40));
+      await expect(fs.writeFile('/big.txt', 'y'.repeat(20))).rejects.toThrow(/ENOSPC/);
+    });
+
+    it('load() bypasses the cap (restoring persisted state)', async () => {
+      const client = mockWorkspaceClient();
+      // Persisted content is 60 bytes — would violate a 50-byte cap if it
+      // went through the wrapper. load() should restore it regardless.
+      client.load.mockResolvedValueOnce(
+        persistedDoc({
+          '/workspace/big.txt': {
+            content: Buffer.from('x'.repeat(60)).toString('base64'),
+            mode: 0o644,
+            mtime: '2025-01-01T00:00:00.000Z',
+          },
+        })
+      );
+      const v = new WorkspaceVolume({
+        workspaceClient: client,
+        initialWorkspaceId: 'ws-test',
+        capacityBytes: 50,
+      });
+      await v.load();
+      // Persisted file is there; new agent writes are still capped.
+      expect((await v.getFilesystem().readFile('/big.txt')).length).toBe(60);
+      await expect(v.getFilesystem().writeFile('/more.txt', 'y')).rejects.toThrow(/ENOSPC/);
+    });
+  });
+
   describe('flush', () => {
     it('is a no-op when no workspaceId is set', async () => {
       const client = mockWorkspaceClient();
diff --git a/x-pack/platform/plugins/shared/agent_builder/server/services/execution/filesystem/workspace_volume.ts b/x-pack/platform/plugins/shared/agent_builder/server/services/execution/filesystem/workspace_volume.ts
