From ff6e39dd10006dee4cbd2b119caf052c42cc1376 Mon Sep 17 00:00:00 2001 From: Harivansh Rathi Date: Sun, 8 Mar 2026 23:38:13 -0700 Subject: [PATCH 1/3] grind mode baby --- package-lock.json | 33 ++ packages/coding-agent/src/cli.ts | 2 +- .../coding-agent/src/core/extensions/types.ts | 7 +- .../coding-agent/src/core/gateway/runtime.ts | 3 +- .../interactive/components/tool-execution.ts | 4 +- .../test/vercel-ai-stream.test.ts | 2 +- packages/pi-grind/README.md | 45 ++ packages/pi-grind/package.json | 35 ++ packages/pi-grind/src/config.ts | 64 +++ packages/pi-grind/src/index.ts | 480 ++++++++++++++++++ packages/pi-grind/src/parser.ts | 101 ++++ packages/pi-grind/src/prompts.ts | 60 +++ packages/pi-grind/src/state.ts | 132 +++++ packages/pi-grind/src/time.ts | 103 ++++ packages/pi-grind/src/types.ts | 53 ++ packages/pi-grind/test/parser.test.ts | 44 ++ packages/pi-grind/test/state.test.ts | 72 +++ 17 files changed, 1232 insertions(+), 8 deletions(-) create mode 100644 packages/pi-grind/README.md create mode 100644 packages/pi-grind/package.json create mode 100644 packages/pi-grind/src/config.ts create mode 100644 packages/pi-grind/src/index.ts create mode 100644 packages/pi-grind/src/parser.ts create mode 100644 packages/pi-grind/src/prompts.ts create mode 100644 packages/pi-grind/src/state.ts create mode 100644 packages/pi-grind/src/time.ts create mode 100644 packages/pi-grind/src/types.ts create mode 100644 packages/pi-grind/test/parser.test.ts create mode 100644 packages/pi-grind/test/state.test.ts diff --git a/package-lock.json b/package-lock.json index fff923f..396fed7 100644 --- a/package-lock.json +++ b/package-lock.json @@ -5819,6 +5819,10 @@ "integrity": "sha512-F3asv42UuXchdzt+xXqfW1OGlVBe+mxa2mqI0pg5yAHZPvFmY3Y6drSf/GQ1A86WgWEN9Kzh/WrgKa6iGcHXLg==", "license": "MIT" }, + "node_modules/pi-grind": { + "resolved": "packages/pi-grind", + "link": true + }, "node_modules/pi-teams": { "resolved": "packages/pi-teams", "link": true @@ -7595,6 +7599,35 @@ "@sinclair/typebox": "*" } }, + "packages/pi-grind": { + "version": "0.1.0", + "license": "MIT", + "devDependencies": { + "@types/node": "^24.3.0", + "typescript": "^5.9.3", + "vitest": "^3.2.4" + }, + "peerDependencies": { + "@mariozechner/pi-coding-agent": "*" + } + }, + "packages/pi-grind/node_modules/@types/node": { + "version": "24.12.0", + "resolved": "https://registry.npmjs.org/@types/node/-/node-24.12.0.tgz", + "integrity": "sha512-GYDxsZi3ChgmckRT9HPU0WEhKLP08ev/Yfcq2AstjrDASOYCSXeyjDsHg4v5t4jOj7cyDX3vmprafKlWIG9MXQ==", + "dev": true, + "license": "MIT", + "dependencies": { + "undici-types": "~7.16.0" + } + }, + "packages/pi-grind/node_modules/undici-types": { + "version": "7.16.0", + "resolved": "https://registry.npmjs.org/undici-types/-/undici-types-7.16.0.tgz", + "integrity": "sha512-Zz+aZWSj8LE6zoxD+xrjh4VfkIG8Ya6LvYkZqtUQGJPZjYl53ypCaUwWqo7eI0x66KBGeRo+mlBEkMSeSZ38Nw==", + "dev": true, + "license": "MIT" + }, "packages/pi-runtime-daemon": { "name": "@local/pi-runtime-daemon", "version": "0.0.1", diff --git a/packages/coding-agent/src/cli.ts b/packages/coding-agent/src/cli.ts index d05106c..f434e65 100644 --- a/packages/coding-agent/src/cli.ts +++ b/packages/coding-agent/src/cli.ts @@ -8,8 +8,8 @@ process.title = "pi"; import { setBedrockProviderModule } from "@mariozechner/pi-ai"; -import { bedrockProviderModule } from "@mariozechner/pi-ai/bedrock-provider"; import { EnvHttpProxyAgent, setGlobalDispatcher } from "undici"; +import { bedrockProviderModule } from "../../ai/src/bedrock-provider.js"; import { main } from "./main.js"; setGlobalDispatcher(new EnvHttpProxyAgent()); diff --git a/packages/coding-agent/src/core/extensions/types.ts b/packages/coding-agent/src/core/extensions/types.ts index 5ee24cf..d45015c 100644 --- a/packages/coding-agent/src/core/extensions/types.ts +++ b/packages/coding-agent/src/core/extensions/types.ts @@ -995,12 +995,15 @@ export interface RegisteredCommand { // Extension API // ============================================================================ +type ExtensionHandlerResult = [R] extends [undefined] + ? Promise | void + : Promise | R | undefined; + /** Handler function type for events */ -// biome-ignore lint/suspicious/noConfusingVoidType: void allows bare return statements export type ExtensionHandler = ( event: E, ctx: ExtensionContext, -) => Promise | R | void; +) => ExtensionHandlerResult; /** * ExtensionAPI passed to extension factory functions. diff --git a/packages/coding-agent/src/core/gateway/runtime.ts b/packages/coding-agent/src/core/gateway/runtime.ts index 8623047..6bebca6 100644 --- a/packages/coding-agent/src/core/gateway/runtime.ts +++ b/packages/coding-agent/src/core/gateway/runtime.ts @@ -9,10 +9,10 @@ import { join } from "node:path"; import { URL } from "node:url"; import type { AgentMessage } from "@mariozechner/pi-agent-core"; import type { AgentSession, AgentSessionEvent } from "../agent-session.js"; +import type { Settings } from "../settings-manager.js"; import { extractMessageText, getLastAssistantText } from "./helpers.js"; import { type GatewayEvent, - type GatewayQueuedMessage, HttpError, type ManagedGatewaySession, } from "./internal-types.js"; @@ -29,7 +29,6 @@ import type { HistoryPart, ModelInfo, } from "./types.js"; -import type { Settings } from "../settings-manager.js"; import { createVercelStreamListener, errorVercelStream, diff --git a/packages/coding-agent/src/modes/interactive/components/tool-execution.ts b/packages/coding-agent/src/modes/interactive/components/tool-execution.ts index 1f2e4f0..fc72044 100644 --- a/packages/coding-agent/src/modes/interactive/components/tool-execution.ts +++ b/packages/coding-agent/src/modes/interactive/components/tool-execution.ts @@ -836,10 +836,10 @@ export class ToolExecutionComponent extends Container { .join("\n"); if (remaining > 0) { text += - theme.fg( + `${theme.fg( "muted", `\n... (${remaining} more lines, ${totalLines} total,`, - ) + ` ${keyHint("expandTools", "to expand")})`; + )} ${keyHint("expandTools", "to expand")})`; } } diff --git a/packages/coding-agent/test/vercel-ai-stream.test.ts b/packages/coding-agent/test/vercel-ai-stream.test.ts index 182540b..ab98c69 100644 --- a/packages/coding-agent/test/vercel-ai-stream.test.ts +++ b/packages/coding-agent/test/vercel-ai-stream.test.ts @@ -3,7 +3,7 @@ import type { AgentSessionEvent } from "../src/core/agent-session.js"; import { createVercelStreamListener, extractUserText, -} from "../src/core/vercel-ai-stream.js"; +} from "../src/core/gateway/vercel-ai-stream.js"; describe("extractUserText", () => { it("extracts text from useChat v5+ format with parts", () => { diff --git a/packages/pi-grind/README.md b/packages/pi-grind/README.md new file mode 100644 index 0000000..2d1e3f3 --- /dev/null +++ b/packages/pi-grind/README.md @@ -0,0 +1,45 @@ +# pi-grind + +Explicit grind mode for Pi. + +Features: + +- Auto-activates only when the user uses explicit grind cues in a prompt +- Persists run state in session custom entries +- Continues work on a heartbeat while running in `pi daemon` +- Pauses automatically when the user sends a normal prompt + +Example prompts: + +- `Keep going on this until 5pm` +- `Don't stop until the refactor is done` + +Commands: + +- `/grind start --until "5pm" Ship the refactor` +- `/grind status` +- `/grind pause` +- `/grind resume` +- `/grind stop` + +Settings: + +```json +{ + "pi-grind": { + "enabled": true, + "pollIntervalMs": 30000, + "cueMode": "explicit-only", + "requireDaemon": true, + "userIntervention": "pause", + "cuePatterns": [ + "don't stop", + "keep going", + "keep running", + "run until", + "until done", + "stay on this until" + ] + } +} +``` diff --git a/packages/pi-grind/package.json b/packages/pi-grind/package.json new file mode 100644 index 0000000..70ded13 --- /dev/null +++ b/packages/pi-grind/package.json @@ -0,0 +1,35 @@ +{ + "name": "pi-grind", + "version": "0.1.0", + "description": "Explicit grind mode for pi with durable follow-up continuation in daemon mode", + "type": "module", + "keywords": [ + "pi-package" + ], + "license": "MIT", + "author": "OpenAI", + "main": "./src/index.ts", + "files": [ + "src", + "test", + "package.json", + "README.md" + ], + "pi": { + "extensions": [ + "./src/index.ts" + ] + }, + "peerDependencies": { + "@mariozechner/pi-agent-core": "*", + "@mariozechner/pi-coding-agent": "*" + }, + "devDependencies": { + "@types/node": "^24.3.0", + "typescript": "^5.9.3", + "vitest": "^3.2.4" + }, + "scripts": { + "test": "vitest --run" + } +} diff --git a/packages/pi-grind/src/config.ts b/packages/pi-grind/src/config.ts new file mode 100644 index 0000000..9fd7117 --- /dev/null +++ b/packages/pi-grind/src/config.ts @@ -0,0 +1,64 @@ +import { getAgentDir, SettingsManager } from "@mariozechner/pi-coding-agent"; +import { DEFAULT_POLL_INTERVAL_MS, GRIND_SETTINGS_KEY, type GrindConfig } from "./types.js"; + +const DEFAULT_CUE_PATTERNS = [ + "don't stop", + "keep going", + "keep running", + "run until", + "until done", + "stay on this until", +]; + +function isRecord(value: unknown): value is Record { + return value !== null && typeof value === "object"; +} + +function asStringArray(value: unknown): string[] | undefined { + if (!Array.isArray(value)) { + return undefined; + } + const strings = value.filter((item): item is string => typeof item === "string"); + return strings.length > 0 ? strings : undefined; +} + +export function loadConfig(cwd: string): GrindConfig { + const settingsManager = SettingsManager.create(cwd, getAgentDir()); + const globalSettings = settingsManager.getGlobalSettings() as Record; + const projectSettings = settingsManager.getProjectSettings() as Record; + + const globalConfig = isRecord(globalSettings[GRIND_SETTINGS_KEY]) ? globalSettings[GRIND_SETTINGS_KEY] : {}; + const projectConfig = isRecord(projectSettings[GRIND_SETTINGS_KEY]) ? projectSettings[GRIND_SETTINGS_KEY] : {}; + + const cuePatterns = + asStringArray(projectConfig.cuePatterns) ?? asStringArray(globalConfig.cuePatterns) ?? DEFAULT_CUE_PATTERNS; + + const pollIntervalMsRaw = + typeof projectConfig.pollIntervalMs === "number" + ? projectConfig.pollIntervalMs + : typeof globalConfig.pollIntervalMs === "number" + ? globalConfig.pollIntervalMs + : DEFAULT_POLL_INTERVAL_MS; + + return { + enabled: + typeof projectConfig.enabled === "boolean" + ? projectConfig.enabled + : typeof globalConfig.enabled === "boolean" + ? globalConfig.enabled + : true, + pollIntervalMs: + Number.isFinite(pollIntervalMsRaw) && pollIntervalMsRaw > 0 + ? Math.max(1_000, Math.floor(pollIntervalMsRaw)) + : DEFAULT_POLL_INTERVAL_MS, + cueMode: "explicit-only", + requireDaemon: + typeof projectConfig.requireDaemon === "boolean" + ? projectConfig.requireDaemon + : typeof globalConfig.requireDaemon === "boolean" + ? globalConfig.requireDaemon + : true, + userIntervention: "pause", + cuePatterns, + }; +} diff --git a/packages/pi-grind/src/index.ts b/packages/pi-grind/src/index.ts new file mode 100644 index 0000000..411db00 --- /dev/null +++ b/packages/pi-grind/src/index.ts @@ -0,0 +1,480 @@ +import type { AgentMessage } from "@mariozechner/pi-agent-core"; +import type { ExtensionAPI, ExtensionContext, RegisteredCommand } from "@mariozechner/pi-coding-agent"; +import { loadConfig } from "./config.js"; +import { parseAutoActivation, parseGrindStatus, parseStopCondition } from "./parser.js"; +import { buildContinuationPrompt, buildRepairPrompt, buildSystemPromptAddon } from "./prompts.js"; +import { createRunState, getLatestRunState, withLoopStatus, withStatus } from "./state.js"; +import { + DEFAULT_COMPLETION_CRITERION, + GRIND_STATE_ENTRY_TYPE, + type GrindConfig, + type GrindRunState, + MAX_PARSE_FAILURES, +} from "./types.js"; + +function isDaemonRuntime(): boolean { + if (process.env.PI_GRIND_FORCE_DAEMON === "1") { + return true; + } + if (process.env.PI_GRIND_FORCE_DAEMON === "0") { + return false; + } + + return ( + process.argv.includes("daemon") || + process.argv.includes("gateway") || + Boolean(process.env.PI_GATEWAY_BIND) || + Boolean(process.env.PI_GATEWAY_PORT) || + Boolean(process.env.PI_GATEWAY_TOKEN) + ); +} + +function getAssistantText(message: AgentMessage): string { + if (message.role !== "assistant") { + return ""; + } + + const parts: string[] = []; + for (const part of message.content) { + if (part.type === "text") { + parts.push(part.text); + } + } + return parts.join("\n"); +} + +function isGrindCommand(text: string): boolean { + return text.trim().toLowerCase().startsWith("/grind"); +} + +function note(ctx: ExtensionContext, message: string): void { + if (ctx.hasUI) { + ctx.ui.notify(message, "info"); + } +} + +function readState(ctx: ExtensionContext): GrindRunState | null { + return getLatestRunState(ctx.sessionManager.getEntries()); +} + +function persistState(pi: ExtensionAPI, ctx: ExtensionContext, state: GrindRunState): GrindRunState { + pi.appendEntry(GRIND_STATE_ENTRY_TYPE, state); + if (ctx.hasUI) { + ctx.ui.setStatus("pi-grind", state.status === "active" ? "GRIND" : state.status.toUpperCase()); + } + return state; +} + +function clearUiStatus(ctx: ExtensionContext): void { + if (ctx.hasUI) { + ctx.ui.setStatus("pi-grind", ""); + } +} + +function maybeExpireRun(pi: ExtensionAPI, ctx: ExtensionContext, state: GrindRunState): GrindRunState | null { + if (!state.deadlineAt) { + return state; + } + + if (Date.parse(state.deadlineAt) > Date.now()) { + return state; + } + + const expired = withStatus(state, "expired", { + lastCheckpoint: "Grind run expired at the configured deadline.", + lastNextAction: null, + pendingRepair: false, + }); + persistState(pi, ctx, expired); + note(ctx, "Grind mode stopped: deadline reached."); + return expired; +} + +function tokenizeArgs(input: string): string[] { + const tokens: string[] = []; + let current = ""; + let quote: '"' | "'" | null = null; + + for (let index = 0; index < input.length; index += 1) { + const char = input[index]; + if ((char === '"' || char === "'") && !quote) { + quote = char; + continue; + } + if (quote && char === quote) { + quote = null; + continue; + } + if (!quote && /\s/.test(char)) { + if (current) { + tokens.push(current); + current = ""; + } + continue; + } + current += char; + } + + if (current) { + tokens.push(current); + } + + return tokens; +} + +function parseStartCommandArgs(args: string): { + goal: string | null; + until: string | null; + criterion: string | null; +} { + const tokens = tokenizeArgs(args); + const goalParts: string[] = []; + let until: string | null = null; + let criterion: string | null = null; + + for (let index = 0; index < tokens.length; index += 1) { + const token = tokens[index]; + if (token === "--until") { + until = tokens[index + 1] ?? null; + index += 1; + continue; + } + if (token.startsWith("--until=")) { + until = token.slice("--until=".length) || null; + continue; + } + if (token === "--criterion") { + criterion = tokens[index + 1] ?? null; + index += 1; + continue; + } + if (token.startsWith("--criterion=")) { + criterion = token.slice("--criterion=".length) || null; + continue; + } + + goalParts.push(token); + } + + return { + goal: goalParts.length > 0 ? goalParts.join(" ") : null, + until, + criterion, + }; +} + +function startRun( + pi: ExtensionAPI, + ctx: ExtensionContext, + config: GrindConfig, + input: { + activation: "explicit" | "command"; + goal: string; + sourcePrompt: string; + deadlineAt: string | null; + completionCriterion: string | null; + }, +): GrindRunState | null { + if (!config.enabled) { + note(ctx, "Grind mode is disabled in settings."); + return null; + } + + if (config.requireDaemon && !isDaemonRuntime()) { + note(ctx, "Durable grind mode requires `pi daemon`."); + return null; + } + + const nextState = createRunState(input); + persistState(pi, ctx, nextState); + note(ctx, "Grind mode activated."); + return nextState; +} + +export default function grind(pi: ExtensionAPI) { + let config: GrindConfig | null = null; + let state: GrindRunState | null = null; + let heartbeat: NodeJS.Timeout | null = null; + + const getConfig = (cwd: string): GrindConfig => { + config = loadConfig(cwd); + return config; + }; + + const stopHeartbeat = () => { + if (heartbeat) { + clearInterval(heartbeat); + heartbeat = null; + } + }; + + const ensureHeartbeat = (ctx: ExtensionContext) => { + stopHeartbeat(); + heartbeat = setInterval(() => { + if (!config || !state || state.status !== "active") { + return; + } + if (config.requireDaemon && !isDaemonRuntime()) { + return; + } + if (!ctx.isIdle() || ctx.hasPendingMessages()) { + return; + } + + const expired = maybeExpireRun(pi, ctx, state); + state = expired; + if (!state || state.status !== "active") { + return; + } + + if (state.pendingRepair) { + pi.sendUserMessage(buildRepairPrompt(state), { + deliverAs: "followUp", + }); + } else { + pi.sendUserMessage(buildContinuationPrompt(state), { + deliverAs: "followUp", + }); + } + }, config?.pollIntervalMs ?? 30_000); + heartbeat.unref?.(); + }; + + const registerCommand = (name: string, command: Omit) => { + pi.registerCommand(name, command); + }; + + registerCommand("grind", { + description: "Manage grind mode: /grind [start|status|pause|resume|stop]", + getArgumentCompletions: (prefix: string) => + ["start", "status", "pause", "resume", "stop"] + .filter((value) => value.startsWith(prefix.trim().toLowerCase())) + .map((value) => ({ value, label: value })), + handler: async (args, ctx) => { + const currentConfig = getConfig(ctx.cwd); + const trimmed = args?.trim() ?? ""; + const [command] = tokenizeArgs(trimmed); + const subcommand = command?.toLowerCase() ?? "status"; + const restArgs = trimmed.slice(command?.length ?? 0).trim(); + + if (subcommand === "start") { + const parsed = parseStartCommandArgs(restArgs); + if (!parsed.goal) { + note(ctx, "Usage: /grind start [--until