tenant.openagents/omega

1#!/usr/bin/env node
2
3// ./script/cargo is a transparent wrapper around cargo that:
4// - When running in a clone of `./zed-industries/zed`
5// - outputs build timings to the ZED_DATA_DIR/build_timings
6// When Zed starts for staff-members it uploads the build timings to Snowflake
7// To use it:
8//  ./script/cargo --init
9// This will add a wrapper to your shell configuration files.
10// (Otherwise set up an alias `cargo=PATH_TO_THIS_FILE`)
11
12// We need to ignore SIGINT in this process so that we can continue
13// processing timing files after the child cargo process exits.
14// The signal will still be delivered to the child process.
15process.on("SIGINT", () => {});
16
17const { spawn, spawnSync } = require("child_process");
18const fs = require("fs");
19const os = require("os");
20const path = require("path");
21const readline = require("readline");
22
23const SUBCOMMANDS_WITH_TIMINGS = new Set(["build", "check", "run", "test"]);
24
25// Built-in cargo aliases
26const CARGO_ALIASES = {
27  b: "build",
28  c: "check",
29  t: "test",
30  r: "run",
31  d: "doc",
32};
33
34function expandAlias(subcommand) {
35  return CARGO_ALIASES[subcommand] || subcommand;
36}
37
38function detectShell() {
39  // Check for PowerShell first (works when running from pwsh)
40  if (process.env.PSModulePath && !process.env.BASH_VERSION) {
41    return "powershell";
42  }
43
44  const shell = process.env.SHELL || "";
45  if (shell.endsWith("/zsh")) return "zsh";
46  if (shell.endsWith("/bash")) return "bash";
47  if (shell.endsWith("/fish")) return "fish";
48  if (shell.endsWith("/pwsh") || shell.endsWith("/powershell")) return "powershell";
49  return path.basename(shell) || "unknown";
50}
51
52function getShellConfigPath(shell) {
53  const home = os.homedir();
54  switch (shell) {
55    case "zsh":
56      return path.join(process.env.ZDOTDIR || home, ".zshrc");
57    case "bash":
58      // Prefer .bashrc, fall back to .bash_profile
59      const bashrc = path.join(home, ".bashrc");
60      if (fs.existsSync(bashrc)) return bashrc;
61      return path.join(home, ".bash_profile");
62    case "fish":
63      return path.join(home, ".config", "fish", "config.fish");
64    case "powershell":
65      if (process.platform === "win32") {
66        // Spawn PowerShell to get the real $PROFILE path, since os.homedir() doesn't account
67        // for OneDrive folder redirection, and the subdirectory differs between Windows PowerShell
68        // 5.x ("WindowsPowerShell") and PowerShell Core ("PowerShell").
69        const psModulePath = process.env.PSModulePath || "";
70        const psExe = psModulePath.toLowerCase().includes("\\windowspowershell\\") ? "powershell" : "pwsh";
71        const result = spawnSync(psExe, ["-NoProfile", "-Command", "$PROFILE"], {
72          encoding: "utf-8",
73          stdio: ["pipe", "pipe", "pipe"],
74          timeout: 5000,
75        });
76        if (result.status === 0 && result.stdout.trim()) {
77          return result.stdout.trim();
78        }
79        // Fallback if spawning fails
80        return path.join(home, "Documents", "PowerShell", "Microsoft.PowerShell_profile.ps1");
81      } else {
82        return path.join(home, ".config", "powershell", "Microsoft.PowerShell_profile.ps1");
83      }
84    default:
85      return null;
86  }
87}
88
89function generateAlias(shell, scriptDir) {
90  const cargoWrapper = path.join(scriptDir, "cargo");
91
92  switch (shell) {
93    case "zsh":
94    case "bash":
95      return `\n# Zed cargo timing wrapper\ncargo() { local w="${cargoWrapper}"; if [[ -x "$w" ]]; then "$w" "$@"; else command cargo "$@"; fi; }\n`;
96    case "fish":
97      return `\n# Zed cargo timing wrapper\nfunction cargo\n    set -l w "${cargoWrapper}"\n    if test -x "$w"\n        "$w" $argv\n        return $status\n    else\n        command cargo $argv\n    end\nend\n`;
98    case "powershell":
99      return `\n# Zed cargo timing wrapper\nfunction cargo {\n    \$wrapper = "${cargoWrapper}"\n    if (Test-Path \$wrapper) {\n        node \$wrapper @args\n    } else {\n        & (Get-Command -Name cargo -CommandType Application | Select-Object -First 1).Source @args\n    }\n}\n`;
100    default:
101      return `cargo() { local w="${cargoWrapper}"; if [[ -x "$w" ]]; then "$w" "$@"; else command cargo "$@"; fi; }`;
102  }
103}
104
105function aliasBlockRegex(shell) {
106  switch (shell) {
107    case "zsh":
108    case "bash":
109      // Comment line + single-line cargo() { ... } function
110      return /\n?# Zed cargo timing wrapper\ncargo\(\) \{[^\n]*\}\n/;
111    case "fish":
112      // Comment line + multi-line function cargo...end block
113      return /\n?# Zed cargo timing wrapper\nfunction cargo\n[\s\S]*?\nend\n/;
114    case "powershell":
115      // Comment line + multi-line function cargo {...} block
116      return /\n?# Zed cargo timing wrapper\nfunction cargo \{[\s\S]*?\n\}\n/;
117    default:
118      return null;
119  }
120}
121
122function initShellAlias() {
123  const scriptDir = __dirname;
124  const shell = detectShell();
125  const configPath = getShellConfigPath(shell);
126  const alias = generateAlias(shell, scriptDir);
127
128  if (!configPath) {
129    console.log(`Unsupported shell: ${shell}`);
130    console.log("\nAdd the following to your shell configuration:\n");
131    console.log(alias);
132    return;
133  }
134
135  // Check if alias already exists; if so, replace it in-place
136  if (fs.existsSync(configPath)) {
137    const content = fs.readFileSync(configPath, "utf-8");
138    if (content.includes("Zed cargo timing wrapper")) {
139      const blockRegex = aliasBlockRegex(shell);
140      const updated = blockRegex ? content.replace(blockRegex, "") : content;
141      fs.writeFileSync(configPath, updated + alias);
142      console.log(`Updated cargo timing alias in ${configPath}`);
143      if (shell === "powershell") {
144        console.log(`\nRestart PowerShell or run: . "${configPath}"`);
145      } else {
146        console.log(`\nRestart your shell or run: source ${configPath}`);
147      }
148      return;
149    }
150  }
151
152  // Create parent directory if needed (for PowerShell on Linux/macOS)
153  const configDir = path.dirname(configPath);
154  if (!fs.existsSync(configDir)) {
155    fs.mkdirSync(configDir, { recursive: true });
156  }
157
158  // Append alias to config file
159  fs.appendFileSync(configPath, alias);
160  console.log(`Added cargo timing alias to ${configPath}`);
161
162  if (shell === "powershell") {
163    console.log(`\nRestart PowerShell or run: . "${configPath}"`);
164  } else {
165    console.log(`\nRestart your shell or run: source ${configPath}`);
166  }
167}
168
169function isZedRepo() {
170  try {
171    const result = spawnSync("git", ["remote", "-v"], {
172      encoding: "utf-8",
173      stdio: ["pipe", "pipe", "pipe"],
174      timeout: 5000,
175    });
176    if (result.status !== 0 || !result.stdout) {
177      return false;
178    }
179    return result.stdout.includes("zed-industries/zed");
180  } catch {
181    return false;
182  }
183}
184
185function findSubcommand(args) {
186  for (let i = 0; i < args.length; i++) {
187    const arg = args[i];
188    // Skip flags and their values
189    if (arg.startsWith("-")) {
190      // If this flag takes a value and it's not using = syntax, skip the next arg too
191      if (!arg.includes("=") && i + 1 < args.length && !args[i + 1].startsWith("-")) {
192        i++;
193      }
194      continue;
195    }
196    // First non-flag argument is the subcommand
197    return { subcommand: arg, index: i };
198  }
199  return null;
200}
201
202function findLatestTimingFile(targetDir) {
203  const timingsDir = path.join(targetDir, "cargo-timings");
204  if (!fs.existsSync(timingsDir)) {
205    return null;
206  }
207
208  const files = fs
209    .readdirSync(timingsDir)
210    .filter((f) => f.startsWith("cargo-timing-") && f.endsWith(".html") && f !== "cargo-timing.html")
211    .map((f) => ({
212      name: f,
213      path: path.join(timingsDir, f),
214      mtime: fs.statSync(path.join(timingsDir, f)).mtime.getTime(),
215    }))
216    .sort((a, b) => b.mtime - a.mtime);
217
218  return files.length > 0 ? files[0].path : null;
219}
220
221function getTargetDir(args) {
222  // Check for --target-dir flag
223  for (let i = 0; i < args.length; i++) {
224    if (args[i] === "--target-dir" && i + 1 < args.length) {
225      return args[i + 1];
226    }
227    if (args[i].startsWith("--target-dir=")) {
228      return args[i].substring("--target-dir=".length);
229    }
230  }
231  // Default target directory
232  return "target";
233}
234
235function runCargoPassthrough(args) {
236  const cargoCmd = process.env.CARGO || "cargo";
237  const result = spawnSync(cargoCmd, args, {
238    stdio: "inherit",
239    shell: false,
240  });
241  process.exit(result.status ?? 1);
242}
243
244async function main() {
245  const args = process.argv.slice(2);
246
247  // Handle --init flag
248  if (args[0] === "--init") {
249    if (process.env.NIX_WRAPPER === "1") {
250      console.error("`--init` not supported when going through the nix wrapper");
251      process.exit(1);
252    }
253    initShellAlias();
254    return;
255  }
256
257  // If not in zed repo, just pass through to cargo
258  if (!isZedRepo()) {
259    runCargoPassthrough(args);
260    return;
261  }
262
263  const cargoCmd = process.env.CARGO || "cargo";
264  const subcommandInfo = findSubcommand(args);
265  const expandedSubcommand = subcommandInfo ? expandAlias(subcommandInfo.subcommand) : null;
266  const shouldAddTimings = expandedSubcommand && SUBCOMMANDS_WITH_TIMINGS.has(expandedSubcommand);
267
268  // Build the final args array
269  let finalArgs = [...args];
270  if (shouldAddTimings) {
271    // Check if --timings is already present
272    const hasTimings = args.some((arg) => arg === "--timings" || arg.startsWith("--timings="));
273    if (!hasTimings) {
274      // Insert --timings after the subcommand
275      finalArgs.splice(subcommandInfo.index + 1, 0, "--timings");
276    }
277  }
278
279  // Run cargo asynchronously so we can handle signals properly
280  const child = spawn(cargoCmd, finalArgs, {
281    stdio: "inherit",
282    shell: false,
283  });
284
285  // Wait for the child to exit
286  const result = await new Promise((resolve) => {
287    child.on("exit", (code, signal) => {
288      resolve({ status: code, signal });
289    });
290  });
291
292  // If we added timings, try to process the timing file (regardless of cargo's exit status)
293  if (shouldAddTimings) {
294    const targetDir = getTargetDir(args);
295    const timingFile = findLatestTimingFile(targetDir);
296
297    if (timingFile) {
298      // Run cargo-timing-info.js in the background
299      const scriptDir = __dirname;
300      const timingScript = path.join(scriptDir, "cargo-timing-info.js");
301
302      if (fs.existsSync(timingScript)) {
303        const timingChild = spawn(process.execPath, [timingScript, timingFile, `cargo ${expandedSubcommand}`], {
304          detached: true,
305          stdio: "ignore",
306        });
307        timingChild.unref();
308      }
309    }
310  }
311
312  // Exit with cargo's exit code, or re-raise the signal if it was killed
313  if (result.signal) {
314    // Reset signal handler and re-raise so parent sees the signal
315    process.removeAllListeners(result.signal);
316    process.kill(process.pid, result.signal);
317  } else {
318    process.exit(result.status ?? 1);
319  }
320}
321
322main();
323