Refactor subagent tool, fix custom tool discovery, fix JSON mode stdout flush

Breaking changes:
- Custom tools now require index.ts entry point in subdirectory
  (e.g., tools/mytool/index.ts instead of tools/mytool.ts)

Subagent tool improvements:
- Refactored to use Message[] from ai package instead of custom types
- Extracted agent discovery to separate agents.ts module
- Added parallel mode streaming (shows progress from all tasks)
- Added turn count to usage stats footer
- Removed redundant Query section from scout output

Fixes:
- JSON mode stdout flush: Fixed race condition where pi --mode json
  could exit before all output was written, causing consumers to
  miss final events

Also:
- Added signal/timeout support to pi.exec() for custom tools and hooks
- Renamed pi-pods bin to avoid conflict with pi
This commit is contained in:
Mario Zechner 2025-12-19 04:10:09 +01:00
parent 1151975afe
commit 4fb3af93fb
15 changed files with 894 additions and 698 deletions

View file

@ -6,8 +6,16 @@
- **Subagent orchestration example**: Added comprehensive custom tool example for spawning and orchestrating sub-agents with isolated context windows. Includes scout/planner/reviewer/worker agents and workflow commands for multi-agent pipelines. ([#215](https://github.com/badlogic/pi-mono/pull/215) by [@nicobailon](https://github.com/nicobailon))
- **`pi.exec()` signal and timeout support**: Custom tools and hooks can now pass `{ signal, timeout }` options to `pi.exec()` for cancellation and timeout handling. The result includes a `killed` flag when the process was terminated.
### Fixed
- **JSON mode stdout flush**: Fixed race condition where `pi --mode json` could exit before all output was written to stdout, causing consumers to miss final events.
### Breaking Changes
- **Custom tools now require `index.ts` entry point**: Auto-discovered custom tools must be in a subdirectory with an `index.ts` file. The old pattern `~/.pi/agent/tools/mytool.ts` must become `~/.pi/agent/tools/mytool/index.ts`. This allows multi-file tools to import helper modules. Explicit paths via `--tool` or `settings.json` still work with any `.ts` file.
- **Hook `tool_result` event restructured**: The `ToolResultEvent` now exposes full tool result data instead of just text. ([#233](https://github.com/badlogic/pi-mono/pull/233))
- Removed: `result: string` field
- Added: `content: (TextContent | ImageContent)[]` - full content array

View file

@ -595,10 +595,12 @@ export default function (pi: HookAPI) {
Custom tools let you extend the built-in toolset (read, write, edit, bash, ...) and are called by the LLM directly. They are TypeScript modules that define tools with optional custom TUI integration for getting user input and custom tool call and result rendering.
**Tool locations:**
- Global: `~/.pi/agent/tools/*.ts`
- Project: `.pi/tools/*.ts`
- CLI: `--tool <path>`
**Tool locations (auto-discovered):**
- Global: `~/.pi/agent/tools/*/index.ts`
- Project: `.pi/tools/*/index.ts`
**Explicit paths:**
- CLI: `--tool <path>` (any .ts file)
- Settings: `customTools` array in `settings.json`
**Quick example:**

View file

@ -22,7 +22,7 @@ See [examples/custom-tools/](../examples/custom-tools/) for working examples.
## Quick Start
Create a file `~/.pi/agent/tools/hello.ts`:
Create a file `~/.pi/agent/tools/hello/index.ts`:
```typescript
import { Type } from "@sinclair/typebox";
@ -51,13 +51,26 @@ The tool is automatically discovered and available in your next pi session.
## Tool Locations
Tools must be in a subdirectory with an `index.ts` entry point:
| Location | Scope | Auto-discovered |
|----------|-------|-----------------|
| `~/.pi/agent/tools/*.ts` | Global (all projects) | Yes |
| `.pi/tools/*.ts` | Project-local | Yes |
| `~/.pi/agent/tools/*/index.ts` | Global (all projects) | Yes |
| `.pi/tools/*/index.ts` | Project-local | Yes |
| `settings.json` `customTools` array | Configured paths | Yes |
| `--tool <path>` CLI flag | One-off/debugging | No |
**Example structure:**
```
~/.pi/agent/tools/
├── hello/
│ └── index.ts # Entry point (auto-discovered)
└── complex-tool/
├── index.ts # Entry point (auto-discovered)
├── helpers.ts # Helper module (not loaded directly)
└── types.ts # Type definitions (not loaded directly)
```
**Priority:** Later sources win on name conflicts. CLI `--tool` takes highest priority.
**Reserved names:** Custom tools cannot use built-in tool names (`read`, `write`, `edit`, `bash`, `grep`, `find`, `ls`).
@ -125,7 +138,7 @@ The factory receives a `ToolAPI` object (named `pi` by convention):
```typescript
interface ToolAPI {
cwd: string; // Current working directory
exec(command: string, args: string[]): Promise<ExecResult>;
exec(command: string, args: string[], options?: ExecOptions): Promise<ExecResult>;
ui: {
select(title: string, options: string[]): Promise<string | null>;
confirm(title: string, message: string): Promise<boolean>;
@ -134,10 +147,36 @@ interface ToolAPI {
};
hasUI: boolean; // false in --print or --mode rpc
}
interface ExecOptions {
signal?: AbortSignal; // Cancel the process
timeout?: number; // Timeout in milliseconds
}
interface ExecResult {
stdout: string;
stderr: string;
code: number;
killed?: boolean; // True if process was killed by signal/timeout
}
```
Always check `pi.hasUI` before using UI methods.
### Cancellation Example
Pass the `signal` from `execute` to `pi.exec` to support cancellation:
```typescript
async execute(toolCallId, params, signal) {
const result = await pi.exec("long-running-command", ["arg"], { signal });
if (result.killed) {
return { content: [{ type: "text", text: "Cancelled" }] };
}
return { content: [{ type: "text", text: result.stdout }] };
}
```
## Session Lifecycle
Tools can implement `onSession` to react to session changes:

View file

@ -363,15 +363,23 @@ ctx.ui.notify("Operation complete", "info");
ctx.ui.notify("Something went wrong", "error");
```
### ctx.exec(command, args)
### ctx.exec(command, args, options?)
Execute a command and get the result.
Execute a command and get the result. Supports cancellation via `AbortSignal` and timeout.
```typescript
const result = await ctx.exec("git", ["status"]);
// result.stdout: string
// result.stderr: string
// result.code: number
// result.killed?: boolean // True if killed by signal/timeout
// With timeout (5 seconds)
const result = await ctx.exec("slow-command", [], { timeout: 5000 });
// With abort signal
const controller = new AbortController();
const result = await ctx.exec("long-command", [], { signal: controller.signal });
```
### ctx.cwd

View file

@ -2,12 +2,22 @@
Delegate tasks to specialized subagents with isolated context windows.
## Features
- **Isolated context**: Each subagent runs in a separate `pi` process
- **Streaming output**: See tool calls and progress as they happen
- **Parallel streaming**: All parallel tasks stream updates simultaneously
- **Markdown rendering**: Final output rendered with proper formatting (expanded view)
- **Usage tracking**: Shows turns, tokens, cost, and context usage per agent
- **Abort support**: Ctrl+C propagates to kill subagent processes
## Structure
```
subagent/
├── README.md # This file
├── subagent.ts # The custom tool
├── subagent.ts # The custom tool (entry point)
├── agents.ts # Agent discovery logic
├── agents/ # Sample agent definitions
│ ├── scout.md # Fast recon, returns compressed context
│ ├── planner.md # Creates implementation plans
@ -21,57 +31,54 @@ subagent/
## Installation
From the `examples/custom-tools/subagent/` directory:
From the repository root, symlink the files:
```bash
# Copy the tool
mkdir -p ~/.pi/agent/tools
cp subagent.ts ~/.pi/agent/tools/
# Symlink the tool (must be in a subdirectory with index.ts)
mkdir -p ~/.pi/agent/tools/subagent
ln -sf "$(pwd)/packages/coding-agent/examples/custom-tools/subagent/subagent.ts" ~/.pi/agent/tools/subagent/index.ts
ln -sf "$(pwd)/packages/coding-agent/examples/custom-tools/subagent/agents.ts" ~/.pi/agent/tools/subagent/agents.ts
# Copy agents
# Symlink agents
mkdir -p ~/.pi/agent/agents
cp agents/*.md ~/.pi/agent/agents/
for f in packages/coding-agent/examples/custom-tools/subagent/agents/*.md; do
ln -sf "$(pwd)/$f" ~/.pi/agent/agents/$(basename "$f")
done
# Copy workflow commands
# Symlink workflow commands
mkdir -p ~/.pi/agent/commands
cp commands/*.md ~/.pi/agent/commands/
for f in packages/coding-agent/examples/custom-tools/subagent/commands/*.md; do
ln -sf "$(pwd)/$f" ~/.pi/agent/commands/$(basename "$f")
done
```
## Security Model
This example intentionally executes a separate `pi` subprocess with a delegated system prompt and tool/model configuration.
This tool executes a separate `pi` subprocess with a delegated system prompt and tool/model configuration.
Treat **project-local agent definitions as repo-controlled prompts**:
- A project can define agents in `.pi/agents/*.md`.
- Those prompts can instruct the model to read files, run bash commands, etc. (depending on the allowed tools).
**Project-local agents** (`.pi/agents/*.md`) are repo-controlled prompts that can instruct the model to read files, run bash commands, etc.
**Default behavior:** the tool only loads **user-level agents** from `~/.pi/agent/agents`.
**Default behavior:** Only loads **user-level agents** from `~/.pi/agent/agents`.
To enable project-local agents, pass `agentScope: "both"` (or `"project"`) explicitly. Only do this for repositories you trust.
To enable project-local agents, pass `agentScope: "both"` (or `"project"`). Only do this for repositories you trust.
When running interactively, the tool will prompt for confirmation before running project-local agents. Set `confirmProjectAgents: false` to disable the prompt.
When running interactively, the tool prompts for confirmation before running project-local agents. Set `confirmProjectAgents: false` to disable.
## Usage
### Single agent
```
> Use the subagent tool with agent "scout" and task "find all authentication code"
Use scout to find all authentication code
```
### Parallel execution
```
> Use subagent with tasks:
> - scout: "analyze the auth module"
> - scout: "analyze the api module"
> - scout: "analyze the database module"
Run 2 scouts in parallel: one to find models, one to find providers
```
### Chained workflow
```
> Use subagent chain:
> 1. scout: "find code related to caching"
> 2. planner: "plan Redis integration using: {previous}"
> 3. worker: "implement: {previous}"
Use a chain: first have scout find the read tool, then have planner suggest improvements
```
### Workflow commands
@ -86,119 +93,32 @@ When running interactively, the tool will prompt for confirmation before running
| Mode | Parameter | Description |
|------|-----------|-------------|
| Single | `{ agent, task }` | One agent, one task |
| Parallel | `{ tasks: [...] }` | Multiple agents run concurrently |
| Parallel | `{ tasks: [...] }` | Multiple agents run concurrently (max 8, 4 concurrent) |
| Chain | `{ chain: [...] }` | Sequential with `{previous}` placeholder |
<details>
<summary>Flow Diagrams</summary>
## Output Display
### Single Mode
**Collapsed view** (default):
- Status icon (✓/✗/⏳) and agent name
- Last 5-10 items (tool calls and text)
- Usage stats: `3 turns ↑input ↓output RcacheRead WcacheWrite $cost ctx:contextTokens model`
```
┌─────────────────┐
│ Main Agent │
└────────┬────────┘
│ "use scout to find auth code"
┌─────────────────┐
│ subagent tool │
└────────┬────────┘
│ pi -p --model haiku ...
┌─────────────────┐
│ Scout │
│ (subprocess) │
└────────┬────────┘
│ stdout
┌─────────────────┐
│ Tool Result │
└─────────────────┘
```
**Expanded view** (Ctrl+O):
- Full task text
- All tool calls with formatted arguments
- Final output rendered as Markdown
- Per-task usage (for chain/parallel)
### Parallel Mode
**Parallel mode streaming**:
- Shows all tasks with live status (⏳ running, ✓ done, ✗ failed)
- Updates as each task makes progress
- Shows "2/3 done, 1 running" status
```
┌──────────────────────┐
│ Main Agent │
└──────────┬───────────┘
┌──────────────────────┐
│ subagent tool │
│ Promise.all() │
└──────────┬───────────┘
┌─────┼─────┐
▼ ▼ ▼
┌──────┐┌──────┐┌──────┐
│Scout ││Scout ││Scout │
│ auth ││ api ││ db │
└──┬───┘└──┬───┘└──┬───┘
│ │ │
└───────┼───────┘
┌──────────────────────┐
│ Combined Result │
└──────────────────────┘
```
### Chain Mode
```
┌─────────────────┐
│ Main Agent │
└────────┬────────┘
┌─────────────────┐
│ subagent tool │
└────────┬────────┘
┌─────────────────┐
│ Step 1: Scout │
└────────┬────────┘
│ {previous} = scout output
┌─────────────────┐
│ Step 2: Planner │
└────────┬────────┘
│ {previous} = planner output
┌─────────────────┐
│ Step 3: Worker │
└────────┬────────┘
┌─────────────────┐
│ Chain Result │
└─────────────────┘
```
### Workflow Command Expansion
```
/implement add Redis
┌─────────────────────────────────────────┐
│ Expands to chain: │
│ 1. scout: "find code for add Redis" │
│ 2. planner: "plan using {previous}" │
│ 3. worker: "implement {previous}" │
└─────────────────────────────────────────┘
┌─────────────────────────────────────────┐
│ Chain Execution │
│ │
│ scout ──► planner ──► worker │
│ (haiku) (sonnet) (sonnet) │
└─────────────────────────────────────────┘
```
</details>
**Tool call formatting** (mimics built-in tools):
- `$ command` for bash
- `read ~/path:1-10` for read
- `grep /pattern/ in ~/path` for grep
- etc.
## Agent Definitions
@ -216,30 +136,37 @@ System prompt for the agent goes here.
```
**Locations:**
- `~/.pi/agent/agents/*.md` - User-level (global)
- `.pi/agents/*.md` - Project-level (only loaded if `agentScope` includes `"project"`)
- `~/.pi/agent/agents/*.md` - User-level (always loaded)
- `.pi/agents/*.md` - Project-level (only with `agentScope: "project"` or `"both"`)
Project agents override user agents with the same name when `agentScope: "both"`.
## Sample Agents
| Agent | Purpose | Model |
|-------|---------|-------|
| `scout` | Fast codebase recon, returns compressed context | Haiku |
| `planner` | Creates implementation plans from context | Sonnet |
| `reviewer` | Code review for quality/security | Sonnet |
| `worker` | General-purpose with full capabilities | Sonnet |
| Agent | Purpose | Model | Tools |
|-------|---------|-------|-------|
| `scout` | Fast codebase recon | Haiku | read, grep, find, ls, bash |
| `planner` | Implementation plans | Sonnet | read, grep, find, ls |
| `reviewer` | Code review | Sonnet | read, grep, find, ls, bash |
| `worker` | General-purpose | Sonnet | (all default) |
## Workflow Commands
Commands are prompt templates that invoke the subagent tool:
| Command | Flow |
|---------|------|
| `/implement <query>` | scout -> planner -> worker |
| `/scout-and-plan <query>` | scout -> planner |
| `/implement-and-review <query>` | worker -> reviewer -> worker |
| `/implement <query>` | scout → planner → worker |
| `/scout-and-plan <query>` | scout → planner |
| `/implement-and-review <query>` | worker → reviewer → worker |
## Error Handling
- **Exit code != 0**: Tool returns error with stderr/output
- **stopReason "error"**: LLM error propagated with error message
- **stopReason "aborted"**: User abort (Ctrl+C) kills subprocess, throws error
- **Chain mode**: Stops at first failing step, reports which step failed
## Limitations
- No timeout/cancellation (subprocess limitation)
- Output truncated to 500 lines / 50KB per agent
- Agents discovered fresh on each invocation
- Output truncated to last 10 items in collapsed view (expand to see all)
- Agents discovered fresh on each invocation (allows editing mid-session)
- Parallel mode limited to 8 tasks, 4 concurrent

View file

@ -0,0 +1,157 @@
/**
* Agent discovery and configuration
*/
import * as fs from "node:fs";
import * as os from "node:os";
import * as path from "node:path";
export type AgentScope = "user" | "project" | "both";
export interface AgentConfig {
name: string;
description: string;
tools?: string[];
model?: string;
systemPrompt: string;
source: "user" | "project";
filePath: string;
}
export interface AgentDiscoveryResult {
agents: AgentConfig[];
projectAgentsDir: string | null;
}
function parseFrontmatter(content: string): { frontmatter: Record<string, string>; body: string } {
const frontmatter: Record<string, string> = {};
const normalized = content.replace(/\r\n/g, "\n");
if (!normalized.startsWith("---")) {
return { frontmatter, body: normalized };
}
const endIndex = normalized.indexOf("\n---", 3);
if (endIndex === -1) {
return { frontmatter, body: normalized };
}
const frontmatterBlock = normalized.slice(4, endIndex);
const body = normalized.slice(endIndex + 4).trim();
for (const line of frontmatterBlock.split("\n")) {
const match = line.match(/^([\w-]+):\s*(.*)$/);
if (match) {
let value = match[2].trim();
if ((value.startsWith('"') && value.endsWith('"')) || (value.startsWith("'") && value.endsWith("'"))) {
value = value.slice(1, -1);
}
frontmatter[match[1]] = value;
}
}
return { frontmatter, body };
}
function loadAgentsFromDir(dir: string, source: "user" | "project"): AgentConfig[] {
const agents: AgentConfig[] = [];
if (!fs.existsSync(dir)) {
return agents;
}
let entries: fs.Dirent[];
try {
entries = fs.readdirSync(dir, { withFileTypes: true });
} catch {
return agents;
}
for (const entry of entries) {
if (!entry.name.endsWith(".md")) continue;
if (!entry.isFile() && !entry.isSymbolicLink()) continue;
const filePath = path.join(dir, entry.name);
let content: string;
try {
content = fs.readFileSync(filePath, "utf-8");
} catch {
continue;
}
const { frontmatter, body } = parseFrontmatter(content);
if (!frontmatter.name || !frontmatter.description) {
continue;
}
const tools = frontmatter.tools
?.split(",")
.map((t) => t.trim())
.filter(Boolean);
agents.push({
name: frontmatter.name,
description: frontmatter.description,
tools: tools && tools.length > 0 ? tools : undefined,
model: frontmatter.model,
systemPrompt: body,
source,
filePath,
});
}
return agents;
}
function isDirectory(p: string): boolean {
try {
return fs.statSync(p).isDirectory();
} catch {
return false;
}
}
function findNearestProjectAgentsDir(cwd: string): string | null {
let currentDir = cwd;
while (true) {
const candidate = path.join(currentDir, ".pi", "agents");
if (isDirectory(candidate)) return candidate;
const parentDir = path.dirname(currentDir);
if (parentDir === currentDir) return null;
currentDir = parentDir;
}
}
export function discoverAgents(cwd: string, scope: AgentScope): AgentDiscoveryResult {
const userDir = path.join(os.homedir(), ".pi", "agent", "agents");
const projectAgentsDir = findNearestProjectAgentsDir(cwd);
const userAgents = scope === "project" ? [] : loadAgentsFromDir(userDir, "user");
const projectAgents =
scope === "user" || !projectAgentsDir ? [] : loadAgentsFromDir(projectAgentsDir, "project");
const agentMap = new Map<string, AgentConfig>();
if (scope === "both") {
for (const agent of userAgents) agentMap.set(agent.name, agent);
for (const agent of projectAgents) agentMap.set(agent.name, agent);
} else if (scope === "user") {
for (const agent of userAgents) agentMap.set(agent.name, agent);
} else {
for (const agent of projectAgents) agentMap.set(agent.name, agent);
}
return { agents: Array.from(agentMap.values()), projectAgentsDir };
}
export function formatAgentList(agents: AgentConfig[], maxItems: number): { text: string; remaining: number } {
if (agents.length === 0) return { text: "none", remaining: 0 };
const listed = agents.slice(0, maxItems);
const remaining = agents.length - listed.length;
return {
text: listed.map((a) => `${a.name} (${a.source}): ${a.description}`).join("; "),
remaining,
};
}

View file

@ -22,9 +22,6 @@ Strategy:
Output format:
## Query
One line summary of what was searched.
## Files Retrieved
List with exact line ranges:
1. `path/to/file.ts` (lines 10-50) - Description of what's here

View file

@ -11,7 +11,14 @@ import { fileURLToPath } from "node:url";
import { createJiti } from "jiti";
import { getAgentDir } from "../../config.js";
import type { HookUIContext } from "../hooks/types.js";
import type { CustomToolFactory, CustomToolsLoadResult, ExecResult, LoadedCustomTool, ToolAPI } from "./types.js";
import type {
CustomToolFactory,
CustomToolsLoadResult,
ExecOptions,
ExecResult,
LoadedCustomTool,
ToolAPI,
} from "./types.js";
// Create require function to resolve module paths at runtime
const require = createRequire(import.meta.url);
@ -69,8 +76,9 @@ function resolveToolPath(toolPath: string, cwd: string): string {
/**
* Execute a command and return stdout/stderr/code.
* Supports cancellation via AbortSignal and timeout.
*/
async function execCommand(command: string, args: string[], cwd: string): Promise<ExecResult> {
async function execCommand(command: string, args: string[], cwd: string, options?: ExecOptions): Promise<ExecResult> {
return new Promise((resolve) => {
const proc = spawn(command, args, {
cwd,
@ -80,6 +88,37 @@ async function execCommand(command: string, args: string[], cwd: string): Promis
let stdout = "";
let stderr = "";
let killed = false;
let timeoutId: NodeJS.Timeout | undefined;
const killProcess = () => {
if (!killed) {
killed = true;
proc.kill("SIGTERM");
// Force kill after 5 seconds if SIGTERM doesn't work
setTimeout(() => {
if (!proc.killed) {
proc.kill("SIGKILL");
}
}, 5000);
}
};
// Handle abort signal
if (options?.signal) {
if (options.signal.aborted) {
killProcess();
} else {
options.signal.addEventListener("abort", killProcess, { once: true });
}
}
// Handle timeout
if (options?.timeout && options.timeout > 0) {
timeoutId = setTimeout(() => {
killProcess();
}, options.timeout);
}
proc.stdout.on("data", (data) => {
stdout += data.toString();
@ -90,18 +129,28 @@ async function execCommand(command: string, args: string[], cwd: string): Promis
});
proc.on("close", (code) => {
if (timeoutId) clearTimeout(timeoutId);
if (options?.signal) {
options.signal.removeEventListener("abort", killProcess);
}
resolve({
stdout,
stderr,
code: code ?? 0,
killed,
});
});
proc.on("error", (err) => {
if (timeoutId) clearTimeout(timeoutId);
if (options?.signal) {
options.signal.removeEventListener("abort", killProcess);
}
resolve({
stdout,
stderr: stderr || err.message,
code: 1,
killed,
});
});
});
@ -182,7 +231,7 @@ export async function loadCustomTools(
// Shared API object - all tools get the same instance
const sharedApi: ToolAPI = {
cwd,
exec: (command: string, args: string[]) => execCommand(command, args, cwd),
exec: (command: string, args: string[], options?: ExecOptions) => execCommand(command, args, cwd, options),
ui: createNoOpUIContext(),
hasUI: false,
};
@ -224,21 +273,32 @@ export async function loadCustomTools(
/**
* Discover tool files from a directory.
* Returns all .ts files (and symlinks to .ts files) in the directory (non-recursive).
* Only loads index.ts files from subdirectories (e.g., tools/mytool/index.ts).
*/
function discoverToolsInDir(dir: string): string[] {
if (!fs.existsSync(dir)) {
return [];
}
const tools: string[] = [];
try {
const entries = fs.readdirSync(dir, { withFileTypes: true });
return entries
.filter((e) => (e.isFile() || e.isSymbolicLink()) && e.name.endsWith(".ts"))
.map((e) => path.join(dir, e.name));
for (const entry of entries) {
if (entry.isDirectory() || entry.isSymbolicLink()) {
// Check for index.ts in subdirectory
const indexPath = path.join(dir, entry.name, "index.ts");
if (fs.existsSync(indexPath)) {
tools.push(indexPath);
}
}
}
} catch {
return [];
}
return tools;
}
/**

View file

@ -19,6 +19,15 @@ export interface ExecResult {
stdout: string;
stderr: string;
code: number;
/** True if the process was killed due to signal or timeout */
killed?: boolean;
}
export interface ExecOptions {
/** AbortSignal to cancel the process */
signal?: AbortSignal;
/** Timeout in milliseconds */
timeout?: number;
}
/** API passed to custom tool factory (stable across session changes) */
@ -26,7 +35,7 @@ export interface ToolAPI {
/** Current working directory */
cwd: string;
/** Execute a command */
exec(command: string, args: string[]): Promise<ExecResult>;
exec(command: string, args: string[], options?: ExecOptions): Promise<ExecResult>;
/** UI methods for user interaction (select, confirm, input, notify) */
ui: ToolUIContext;
/** Whether UI is available (false in print/RPC mode) */

View file

@ -6,6 +6,7 @@ import { spawn } from "node:child_process";
import type { LoadedHook, SendHandler } from "./loader.js";
import type {
BranchEventResult,
ExecOptions,
ExecResult,
HookError,
HookEvent,
@ -28,13 +29,45 @@ export type HookErrorListener = (error: HookError) => void;
/**
* Execute a command and return stdout/stderr/code.
* Supports cancellation via AbortSignal and timeout.
*/
async function exec(command: string, args: string[], cwd: string): Promise<ExecResult> {
async function exec(command: string, args: string[], cwd: string, options?: ExecOptions): Promise<ExecResult> {
return new Promise((resolve) => {
const proc = spawn(command, args, { cwd, shell: false });
let stdout = "";
let stderr = "";
let killed = false;
let timeoutId: NodeJS.Timeout | undefined;
const killProcess = () => {
if (!killed) {
killed = true;
proc.kill("SIGTERM");
// Force kill after 5 seconds if SIGTERM doesn't work
setTimeout(() => {
if (!proc.killed) {
proc.kill("SIGKILL");
}
}, 5000);
}
};
// Handle abort signal
if (options?.signal) {
if (options.signal.aborted) {
killProcess();
} else {
options.signal.addEventListener("abort", killProcess, { once: true });
}
}
// Handle timeout
if (options?.timeout && options.timeout > 0) {
timeoutId = setTimeout(() => {
killProcess();
}, options.timeout);
}
proc.stdout?.on("data", (data) => {
stdout += data.toString();
@ -45,11 +78,19 @@ async function exec(command: string, args: string[], cwd: string): Promise<ExecR
});
proc.on("close", (code) => {
resolve({ stdout, stderr, code: code ?? 0 });
if (timeoutId) clearTimeout(timeoutId);
if (options?.signal) {
options.signal.removeEventListener("abort", killProcess);
}
resolve({ stdout, stderr, code: code ?? 0, killed });
});
proc.on("error", (_err) => {
resolve({ stdout, stderr, code: 1 });
if (timeoutId) clearTimeout(timeoutId);
if (options?.signal) {
options.signal.removeEventListener("abort", killProcess);
}
resolve({ stdout, stderr, code: 1, killed });
});
});
}
@ -166,7 +207,7 @@ export class HookRunner {
*/
private createContext(): HookEventContext {
return {
exec: (command: string, args: string[]) => exec(command, args, this.cwd),
exec: (command: string, args: string[], options?: ExecOptions) => exec(command, args, this.cwd, options),
ui: this.uiContext,
hasUI: this.hasUI,
cwd: this.cwd,

View file

@ -27,6 +27,15 @@ export interface ExecResult {
stdout: string;
stderr: string;
code: number;
/** True if the process was killed due to signal or timeout */
killed?: boolean;
}
export interface ExecOptions {
/** AbortSignal to cancel the process */
signal?: AbortSignal;
/** Timeout in milliseconds */
timeout?: number;
}
/**
@ -65,7 +74,7 @@ export interface HookUIContext {
*/
export interface HookEventContext {
/** Execute a command and return stdout/stderr/code */
exec(command: string, args: string[]): Promise<ExecResult>;
exec(command: string, args: string[], options?: ExecOptions): Promise<ExecResult>;
/** UI methods for user interaction */
ui: HookUIContext;
/** Whether UI is available (false in print mode) */

View file

@ -445,6 +445,10 @@ export async function main(args: string[]) {
await runPrintMode(session, mode, parsed.messages, initialMessage, initialAttachments);
// Clean up and exit (file watchers keep process alive)
stopThemeWatcher();
// Wait for stdout to fully flush before exiting
if (process.stdout.writableLength > 0) {
await new Promise<void>((resolve) => process.stdout.once("drain", resolve));
}
process.exit(0);
}
}

View file

@ -109,4 +109,13 @@ export async function runPrintMode(
}
}
}
// Ensure stdout is fully flushed before returning
// This prevents race conditions where the process exits before all output is written
await new Promise<void>((resolve, reject) => {
process.stdout.write("", (err) => {
if (err) reject(err);
else resolve();
});
});
}

View file

@ -4,7 +4,7 @@
"description": "CLI tool for managing vLLM deployments on GPU pods",
"type": "module",
"bin": {
"pi": "dist/cli.js"
"pi-pods": "dist/cli.js"
},
"scripts": {
"clean": "rm -rf dist",