pi-shepherd 0.1.0

package/shepherd/state-tracker.ts

@@ -0,0 +1,119 @@
+/**
+ * Guard StateTracker — 有状态规则引擎的状态收集器
+ *
+ * 自动追踪工具调用次数、字符量和连续错误次数，
+ * 供 shepherd 的 tool_result hook 查询匹配。
+ *
+ * 状态只存在内存中，不持久化（重启 pi 重置）。
+ */
+
+export interface StateCondition {
+	tools?: string[];
+	gte?: number;
+	lte?: number;
+	countKind?: "calls" | "errors" | "chars";
+}
+
+export interface ResettableRule {
+	comment: string;
+	state?: StateCondition;
+	resetOn?: string[];
+	/** 运行时：已触发标记 */
+	_triggered?: boolean;
+}
+
+export class StateTracker {
+	private toolCounts = new Map<string, number>();
+	private toolChars = new Map<string, number>();
+	private consecutiveErrors = 0;
+	/** 规则级触发状态：key = comment（规则唯一标识），value = { triggered, firedCount } */
+	private ruleStates = new Map<
+		string,
+		{ triggered: boolean; firedCount: number }
+	>();
+
+	/** 每次 tool_result 调用，自动更新状态 */
+	update(toolName: string, resultChars: number, isError: boolean): void {
+		this.toolCounts.set(toolName, (this.toolCounts.get(toolName) ?? 0) + 1);
+		this.toolChars.set(
+			toolName,
+			(this.toolChars.get(toolName) ?? 0) + resultChars,
+		);
+		this.consecutiveErrors = isError ? this.consecutiveErrors + 1 : 0;
+	}
+
+	/** 判断规则的状态条件是否满足 */
+	matches(condition: StateCondition): boolean {
+		const kind = condition.countKind ?? "calls";
+		let value: number;
+
+		if (kind === "errors") {
+			value = this.consecutiveErrors;
+		} else if (kind === "chars") {
+			value = (condition.tools ?? []).reduce(
+				(sum, t) => sum + (this.toolChars.get(t) ?? 0),
+				0,
+			);
+		} else {
+			value = (condition.tools ?? []).reduce(
+				(sum, t) => sum + (this.toolCounts.get(t) ?? 0),
+				0,
+			);
+		}
+
+		if (condition.gte != null && value < condition.gte) return false;
+		if (condition.lte != null && value > condition.lte) return false;
+		return true;
+	}
+
+	/** 检查当前 toolName 是否触发 resetOn → 重置匹配规则的计数 */
+	resetIf(toolName: string, rules: ResettableRule[]): void {
+		for (const rule of rules) {
+			if (rule.resetOn?.includes(toolName)) {
+				if (rule.state?.tools) {
+					for (const t of rule.state.tools) {
+						this.toolCounts.delete(t);
+						this.toolChars.delete(t);
+					}
+				}
+				// 重置规则触发状态（持久化的 ruleStates）
+				const key = rule.comment;
+				if (this.ruleStates.has(key)) {
+					this.ruleStates.get(key)!.triggered = false;
+				}
+				rule._triggered = false;
+			}
+		}
+	}
+
+	/** 获取格式化信息（用于 reason 模板替换） */
+	getStats(tools: string[]): { count: number; chars: number } {
+		return {
+			count: tools.reduce((s, t) => s + (this.toolCounts.get(t) ?? 0), 0),
+			chars: tools.reduce((s, t) => s + (this.toolChars.get(t) ?? 0), 0),
+		};
+	}
+
+	/** 规则是否已触发（持久化在 StateTracker 内，不受 loadRules 重新创建实例影响） */
+	isTriggered(ruleKey: string): boolean {
+		return this.ruleStates.get(ruleKey)?.triggered ?? false;
+	}
+
+	/** 标记规则已触发，返回当前触发次数 */
+	markTriggered(ruleKey: string): number {
+		const s = this.ruleStates.get(ruleKey) ?? {
+			triggered: true,
+			firedCount: 0,
+		};
+		s.triggered = true;
+		s.firedCount++;
+		this.ruleStates.set(ruleKey, s);
+		return s.firedCount;
+	}
+
+	/** 计算递增阈值：baseThreshold + firedCount * step */
+	nextThreshold(baseThreshold: number, ruleKey: string, step = 5): number {
+		const firedCount = this.ruleStates.get(ruleKey)?.firedCount ?? 0;
+		return baseThreshold + firedCount * step;
+	}
+}

package/shepherd/tool-event-types.ts

@@ -0,0 +1,31 @@
+/**
+ * pi SDK 工具事件类型定义
+ * 统一 tool_call 和 tool_result 事件的联合类型
+ */
+
+import type {
+	BashToolCallEvent,
+	BashToolResultEvent,
+	CustomToolCallEvent,
+	CustomToolResultEvent,
+	EditToolCallEvent,
+	EditToolResultEvent,
+	GrepToolCallEvent,
+	GrepToolResultEvent,
+	LsToolCallEvent,
+	LsToolResultEvent,
+	ReadToolCallEvent,
+	ReadToolResultEvent,
+	WriteToolCallEvent,
+	WriteToolResultEvent,
+} from "@earendil-works/pi-coding-agent";
+
+/** getMatchTargets 接收的事件类型（tool_call 或 tool_result） */
+export type ToolEvent =
+	| BashToolCallEvent | BashToolResultEvent
+	| CustomToolCallEvent | CustomToolResultEvent
+	| EditToolCallEvent | EditToolResultEvent
+	| GrepToolCallEvent | GrepToolResultEvent
+	| LsToolCallEvent | LsToolResultEvent
+	| ReadToolCallEvent | ReadToolResultEvent
+	| WriteToolCallEvent | WriteToolResultEvent;

package/shepherd/tool-hooks.ts

@@ -0,0 +1,176 @@
+/**
+ * tool_call 和 tool_result hook 处理逻辑
+ */
+
+import type { ExtensionAPI } from "@earendil-works/pi-coding-agent";
+
+/** pi 工具结果中的 content block 类型 */
+interface ContentBlock {
+	type: string;
+	text?: string;
+	[key: string]: unknown;
+}
+import { pushWarning } from "./ephemeral.js";
+import { checkLineCount } from "./line-count.js";
+import {
+	getMatchTargets,
+	isRtkAvailable,
+	isSubagent,
+	type LoadRulesOptions,
+	loadRules,
+	type Rule,
+	ruleMatches,
+} from "./rules.js";
+import type { ResettableRule, StateTracker } from "./state-tracker.js";
+
+/** 工具 hook 共享的可变状态 */
+export interface ToolState {
+	hasEdits: boolean;
+	tracker: StateTracker;
+	cachedTools: Set<string> | null;
+}
+
+export function getAvailableTools(
+	pi: ExtensionAPI,
+	state: ToolState,
+): Set<string> {
+	if (!state.cachedTools) {
+		state.cachedTools = new Set(pi.getActiveTools());
+	}
+	return state.cachedTools;
+}
+
+export function toolsAvailable(
+	rule: Rule,
+	pi: ExtensionAPI,
+	state: ToolState,
+): boolean {
+	if (!rule.requiresTools || rule.requiresTools.length === 0) return true;
+	const tools = getAvailableTools(pi, state);
+	return rule.requiresTools.every((t) => tools.has(t));
+}
+
+/** 注册 tool_call hook */
+export function registerToolCall(
+	pi: ExtensionAPI,
+	state: ToolState,
+	rulesDir?: string,
+	rulesOptions?: LoadRulesOptions,
+): void {
+	pi.on("tool_call", async (event) => {
+		if (event.toolName === "edit" || event.toolName === "write") {
+			state.hasEdits = true;
+		}
+
+		const rules = loadRules(rulesDir, rulesOptions).filter(
+			(r) => r.hook === "tool_call" && r.tool === event.toolName,
+		);
+		if (rules.length === 0) return;
+
+		const targets = getMatchTargets(event.toolName!, event, "tool_call");
+		if (!targets || Object.keys(targets).length === 0) return;
+
+		for (const rule of rules) {
+			if (isSubagent() && rule.subagent === false) continue;
+			if (!toolsAvailable(rule, pi, state)) continue;
+			if (!ruleMatches(rule, event.toolName!, targets)) continue;
+
+			if (rule.action === "block") {
+				return { block: true, reason: `⛔ shepherd: ${rule.reason}` };
+			}
+
+			if (rule.action === "rewrite" && event.toolName === "bash") {
+				if (!isRtkAvailable) continue;
+				const cmd = targets.command;
+				if (cmd && !cmd.startsWith("rtk ")) {
+					(event.input as any).command = `rtk ${cmd}`;
+				}
+			}
+
+			if (rule.action === "notify") {
+				pushWarning(rule.reason, rule.comment);
+			}
+		}
+	});
+}
+
+/** 注册 tool_result hook */
+export function registerToolResult(
+	pi: ExtensionAPI,
+	state: ToolState,
+	rulesDir?: string,
+	rulesOptions?: LoadRulesOptions,
+): void {
+	pi.on("tool_result", async (event, ctx) => {
+		// 行数检查（edit/write/memory_update 后）
+		if (event.toolName === "edit" || event.toolName === "write") {
+			const filePath = (event.input as any)?.path as string;
+			if (filePath) {
+				checkLineCount(filePath);
+			}
+		} else if (event.toolName === "memory_update") {
+			// memory_update 工具内部写文件，从 tool_result 文本中提取路径
+			const resultText =
+				event.content
+					?.filter((c: ContentBlock) => c.type === "text")
+					?.map((c: ContentBlock) => c.text)
+					.join("") ?? "";
+			const pathMatch = resultText.match(/(?:写入|更新|创建).*?`([^`]+\.md)`/);
+			if (pathMatch?.[1]) checkLineCount(pathMatch[1]);
+		}
+
+		// 状态更新
+		const resultText =
+			event.content
+				?.filter((c: ContentBlock) => c.type === "text")
+				?.map((c: ContentBlock) => c.text)
+				.join("") ?? "";
+		state.tracker.update(event.toolName, resultText.length, !!event.isError);
+
+		const allRules = loadRules(rulesDir, rulesOptions);
+		const rules = allRules.filter((r) => r.hook === "tool_result");
+
+		// resetOn 检查
+		state.tracker.resetIf(event.toolName, rules as ResettableRule[]);
+
+		// 规则匹配
+		for (const rule of rules) {
+			if ((rule as any).enabled === false) continue;
+			if (isSubagent() && rule.subagent === false) continue;
+			if (!toolsAvailable(rule, pi, state)) continue;
+			if (rule.requireSuccess && event.isError) continue;
+			if (rule.tool && rule.tool !== event.toolName) continue;
+
+			// 正则条件匹配
+			if (rule.conditions || rule.pattern) {
+				const targets = getMatchTargets(event.toolName!, event, "tool_result");
+				if (!targets || Object.keys(targets).length === 0) continue;
+				if (!ruleMatches(rule, event.toolName!, targets)) continue;
+			}
+
+			// 状态条件匹配
+			if (rule.state) {
+				const ruleKey = rule.comment;
+				if (state.tracker.isTriggered(ruleKey)) {
+					const baseGte = rule.state.gte ?? 1;
+					const nextThresh = state.tracker.nextThreshold(baseGte, ruleKey);
+					const stats = state.tracker.getStats(rule.state.tools ?? []);
+					if (stats.count < nextThresh) continue;
+				}
+				if (!state.tracker.matches(rule.state)) continue;
+			}
+
+			// 执行 action：注入到临时缓冲区，由 before_provider_request 消费
+			if (rule.action === "steer") {
+				if (rule.state) state.tracker.markTriggered(rule.comment);
+				const stats = state.tracker.getStats(rule.state?.tools ?? []);
+				const reason = rule.reason
+					.replace("{count}", String(stats.count))
+					.replace("{chars}", String(Math.round(stats.chars / 1000)));
+				pushWarning(reason, rule.comment);
+			} else if (rule.action === "notify") {
+				pushWarning(rule.reason, rule.comment);
+			}
+		}
+	});
+}

package/shepherd/worktree-check.ts

@@ -0,0 +1,130 @@
+/**
+ * Worktree 检测与清理
+ * session_start 时自动检测未合并 worktree 并提醒
+ */
+
+import { execFileSync } from "node:child_process";
+import { existsSync, rmSync } from "node:fs";
+import { join } from "node:path";
+import { isInWorktree, isSubagent } from "./rules.js";
+
+interface UIContext {
+	notify?: (msg: string, level?: "error" | "info" | "warning") => void;
+}
+
+function git(args: string[], opts: { cwd: string; timeout?: number }): string {
+	try {
+		return execFileSync("git", args, {
+			timeout: opts.timeout ?? 5000,
+			stdio: ["pipe", "pipe", "pipe"],
+			cwd: opts.cwd,
+			encoding: "utf-8",
+		}).trim();
+	} catch {
+		return "";
+	}
+}
+
+export function checkWorktrees(ctx: UIContext): void {
+	if (isSubagent()) return;
+	if (isInWorktree()) return;
+
+	try {
+		const cwd = process.cwd();
+		const branches = git(["branch", "--list", "worktree/*", "--no-color"], {
+			cwd,
+		});
+		if (!branches) return;
+
+		const branchList = branches
+			.split("\n")
+			.map((b) => b.replace(/^\s*[+*]\s*/, "").trim())
+			.filter(Boolean);
+		if (branchList.length === 0) return;
+
+		const merged: string[] = [];
+		const unmerged: string[] = [];
+
+		for (const branch of branchList) {
+			const logResult = git(["log", `main..${branch}`, "--oneline"], {
+				cwd,
+				timeout: 3000,
+			});
+			if (logResult) {
+				unmerged.push(branch);
+			} else {
+				const verify = git(["rev-parse", "--verify", `${branch}^{commit}`], {
+					cwd,
+					timeout: 2000,
+				});
+				if (verify) {
+					merged.push(branch);
+				}
+				// 无 commit 的分支，忽略
+			}
+		}
+
+		if (merged.length > 0) {
+			for (const branch of merged) {
+				const name = branch.replace(/^worktree\//, "");
+				const wtPath = join(cwd, ".worktrees", name);
+				try {
+					// 先尝试 git worktree remove
+					git(["worktree", "remove", join(".worktrees", name), "--force"], {
+						cwd,
+						timeout: 10000,
+					});
+				} catch {
+					// fallback: 用 fs.rmSync 而非 shell rm -rf
+					if (existsSync(wtPath)) {
+						try {
+							rmSync(wtPath, { recursive: true, force: true });
+						} catch {
+							/* ignore */
+						}
+					}
+				}
+				git(["branch", "-d", branch], { cwd, timeout: 3000 });
+			}
+			ctx.notify?.(
+				`🧹 shepherd: 自动清理了 ${merged.length} 个已合并 worktree: ${merged.map((b) => b.replace("worktree/", "")).join(", ")}`,
+				"info",
+			);
+		}
+
+		if (unmerged.length > 0) {
+			const details = unmerged.map((b) => {
+				const name = b.replace(/^worktree\//, "");
+				const wtPath = join(cwd, ".worktrees", name);
+				let uncommitted = 0;
+				let aheadCount = 0;
+				if (existsSync(wtPath)) {
+					const statusLines = git(["-C", wtPath, "status", "--short"], {
+						cwd,
+						timeout: 3000,
+					});
+					if (statusLines) uncommitted = statusLines.split("\n").length;
+				}
+				const logLines = git(["log", `main..${b}`, "--oneline"], {
+					cwd,
+					timeout: 3000,
+				});
+				if (logLines) aheadCount = logLines.split("\n").length;
+				const parts: string[] = [];
+				if (aheadCount > 0) parts.push(`${aheadCount} 个未合并提交`);
+				if (uncommitted > 0) parts.push(`${uncommitted} 个未提交改动`);
+				const summary = parts.length > 0 ? ` (${parts.join(", ")})` : "";
+				return `  - ${name}${summary}`;
+			});
+
+			ctx.notify?.(
+				`⚠️ shepherd: ${unmerged.length} 个未合并 worktree:\n` +
+					details.join("\n") +
+					"\n合并: /worktree-merge ｜ 删除: /worktree destroy <名称>",
+				"warning",
+			);
+		}
+	} catch {
+		// 非 git 仓库，忽略
+	}
+}

package/tsconfig.json

@@ -0,0 +1,14 @@
+{
+	"extends": "../tsconfig.base.json",
+	"compilerOptions": {
+		"strict": false,
+		"outDir": "dist",
+		"rootDir": ".",
+		"paths": {
+			"@pi-atelier/shared-utils": [
+				"./node_modules/@pi-atelier/shared-utils/src/index.ts"
+			]
+		}
+	},
+	"include": ["index.ts", "shepherd/**/*.ts", "tests/**/*.ts"]
+}

package/vitest.config.ts

@@ -0,0 +1,13 @@
+import path from "node:path";
+import { createConfig } from "../vitest.config.base";
+
+export default createConfig({
+	alias: {
+		"@earendil-works/pi-coding-agent": true,
+		"@pi-atelier/shepherd": path.resolve(__dirname, "./shepherd"),
+	},
+	include: ["tests/**/*.test.ts"],
+	test: {
+		testTimeout: 10000,
+	},
+});