src/worker/workspace.ts

// Workspace ảo: hệ thống file trong bộ nhớ + shell `just-bash` + bốn tool cho agent.

import { Bash } from 'just-bash/browser';
import type { SerializedEntry, ShellResult, ToolDef } from '../shared/types';

export const WORKSPACE_DIR = '/workspace';
export const MAX_FILE_BYTES = 1024 ** 2;
const MAX_TOOL_OUTPUT_CHARS = 6000;

export const DEFAULT_FILES: Record<string, string> = {
  '/workspace/README.md': '# Browser workspace\n\nYour files are saved in this browser.\n',
  '/workspace/sales.csv': 'product,units\nNotebook,12\nPencil,30\nFolder,8\n',
  '/workspace/products.json':
    JSON.stringify(
      [
        { product: 'Notebook', price: 5 },
        { product: 'Pencil', price: 2 },
        { product: 'Folder', price: 3 },
        { product: 'Eraser', price: 1.5 },
        { product: 'Stapler', price: 4.99 },
        { product: 'Ruler', price: 2.5 },
      ],
      null,
      2,
    ) + '\n',
};

export const TOOLS: ToolDef[] = [
  {
    name: 'bash',
    label: 'Bash',
    description:
      'Run a bash command in /workspace. Supports cat, ls, grep, sed, awk, jq, sort, pipes, redirects, and shell scripts. No network, Node.js, Python, npm, or native binaries. Files persist across commands; cd and variables reset each call.',
    parameters: {
      type: 'object',
      properties: { command: { type: 'string', description: 'The shell command to execute.' } },
      required: ['command'],
    },
  },
  {
    name: 'read',
    label: 'Read file',
    description: 'Read a UTF-8 text file. Use bash head/grep for large files.',
    parameters: { type: 'object', properties: { path: { type: 'string' } }, required: ['path'] },
  },
  {
    name: 'write',
    label: 'Write file',
    description: 'Create or replace a UTF-8 text file. Parent directories are created.',
    parameters: {
      type: 'object',
      properties: { path: { type: 'string' }, content: { type: 'string' } },
      required: ['path', 'content'],
    },
  },
  {
    name: 'edit',
    label: 'Edit file',
    description: 'Replace one exact occurrence of oldText with newText in a UTF-8 file. oldText must match exactly once.',
    parameters: {
      type: 'object',
      properties: { path: { type: 'string' }, oldText: { type: 'string' }, newText: { type: 'string' } },
      required: ['path', 'oldText', 'newText'],
    },
  },
];

export interface ToolOutput {
  content: Array<{ type: 'text'; text: string }>;
  details: { exitCode?: number; path?: string };
  /** true nếu tool chạy xong nhưng kết quả nên coi là lỗi (ví dụ exit code khác 0). */
  isError: boolean;
}

const truncate = (text: string) =>
  text.length > MAX_TOOL_OUTPUT_CHARS
    ? text.slice(0, MAX_TOOL_OUTPUT_CHARS) +
      '\n[Output truncated at 6,000 characters; use head, grep, or a narrower query.]'
    : text;

const output = (text: string, details: ToolOutput['details'] = {}, isError = false): ToolOutput => ({
  content: [{ type: 'text', text: truncate(text) }],
  details,
  isError,
});

const SYSTEM_DIRS = ['/bin', '/usr/bin', '/dev', '/proc'];
const isSystemPath = (p: string) => SYSTEM_DIRS.some((d) => p === d || p.startsWith(d + '/'));

function base64ToBytes(b64: string): Uint8Array {
  return Uint8Array.from(atob(b64), (c) => c.charCodeAt(0));
}

function bytesToBase64(bytes: Uint8Array): string {
  let bin = '';
  for (let i = 0; i < bytes.length; i += 32768) bin += String.fromCharCode(...bytes.subarray(i, i + 32768));
  return btoa(bin);
}

export class Workspace {
  readonly bash: Bash;
  readonly ready: Promise<void>;

  /**
   * @param entries  trạng thái đã lưu (file + thư mục + symlink); ưu tiên hơn `files`
   * @param files    file văn bản ban đầu (đường dẫn tuyệt đối → nội dung)
   */
  constructor(entries?: Array<[string, SerializedEntry]>, files: Record<string, string> = DEFAULT_FILES) {
    const initial = entries
      ? Object.fromEntries(
          entries
            .filter((e): e is [string, Extract<SerializedEntry, { type: 'file' }>] => e[1].type === 'file')
            .map(([path, e]) => [path, { mode: e.mode, content: base64ToBytes(e.base64) }]),
        )
      : files;

    const bash = new Bash({
      files: initial,
      cwd: WORKSPACE_DIR,
      executionLimitProfile: 'hardened',
      executionLimits: {
        maxExecutionTimeMs: 10_000,
        maxCommandCount: 3000,
        maxFileSystemBytes: 16 * 1024 ** 2,
        maxOutputSize: 262_144,
        maxSourceBytes: 131_072,
      },
    });
    this.bash = bash;

    this.ready = (async () => {
      for (const [path, e] of entries ?? []) {
        if (e.type === 'directory') {
          await bash.fs.mkdir(path, { recursive: true });
          await bash.fs.chmod(path, e.mode);
        } else if (e.type === 'symlink') {
          await bash.fs.symlink(e.target, path);
        }
      }
    })();
  }

  resolve(path: string): string {
    return this.bash.fs.resolvePath(WORKSPACE_DIR, path);
  }

  async read(path: string): Promise<string> {
    return this.bash.fs.readFile(this.resolve(path));
  }

  async write(path: string, content: string): Promise<void> {
    const full = this.resolve(path);
    if (new TextEncoder().encode(content).length > MAX_FILE_BYTES) {
      throw new Error('The editor accepts files up to 1 MiB.');
    }
    await this.bash.fs.mkdir(full.slice(0, full.lastIndexOf('/')) || '/', { recursive: true });
    await this.bash.fs.writeFile(full, content);
  }

  async exec(command: string, signal?: AbortSignal): Promise<ShellResult> {
    const r = await this.bash.exec(command, { signal });
    return { stdout: r.stdout, stderr: r.stderr, exitCode: r.exitCode };
  }

  /** Tất cả file thường (bỏ thư mục hệ thống) → nội dung văn bản. */
  async snapshot(): Promise<Record<string, string>> {
    const out: Record<string, string> = {};
    for (const p of this.bash.fs.getAllPaths().sort()) {
      if (isSystemPath(p)) continue;
      if ((await this.bash.fs.lstat(p)).isFile) out[p] = await this.bash.fs.readFile(p);
    }
    return out;
  }

  /** Bản đầy đủ (kể cả thư mục, symlink, quyền) để lưu vào IndexedDB. */
  async serialize(): Promise<Array<[string, SerializedEntry]>> {
    const out: Array<[string, SerializedEntry]> = [];
    for (const p of this.bash.fs.getAllPaths().sort()) {
      if (isSystemPath(p)) continue;
      const st = await this.bash.fs.lstat(p);
      if (st.isSymbolicLink) out.push([p, { type: 'symlink', target: await this.bash.fs.readlink(p) }]);
      else if (st.isDirectory) out.push([p, { type: 'directory', mode: st.mode }]);
      else if (st.isFile) {
        out.push([p, { type: 'file', base64: bytesToBase64(await this.bash.fs.readFileBuffer(p)), mode: st.mode }]);
      }
    }
    return out;
  }

  /** Chạy một tool của agent. Ném lỗi nếu tham số sai hoặc thao tác thất bại. */
  async runTool(name: string, args: Record<string, unknown>, signal?: AbortSignal): Promise<ToolOutput> {
    const str = (key: string) => {
      const v = args[key];
      if (typeof v !== 'string') throw new Error(`Missing or invalid "${key}".`);
      return v;
    };
    switch (name) {
      case 'bash': {
        const r = await this.bash.exec(str('command'), { signal });
        const text = `${r.stdout}${r.stderr ? '\nSTDERR:\n' + r.stderr : ''}\nExit code: ${r.exitCode}`;
        return output(text, { exitCode: r.exitCode }, r.exitCode !== 0);
      }
      case 'read': {
        const real = await this.bash.fs.realpath(this.resolve(str('path')));
        return output(await this.bash.fs.readFile(real), { path: real });
      }
      case 'write': {
        await this.write(str('path'), str('content'));
        return output('Wrote ' + this.resolve(str('path')));
      }
      case 'edit': {
        const path = str('path');
        const oldText = str('oldText');
        const newText = str('newText');
        const current = await this.bash.fs.readFile(this.resolve(path));
        if (!oldText || current.split(oldText).length !== 2) {
          throw new Error('oldText must match exactly once. Read the file first.');
        }
        await this.write(path, current.replace(oldText, () => newText));
        return output('Edited ' + this.resolve(path));
      }
      default:
        throw new Error('Unknown tool: ' + name);
    }
  }
}