Fix oversized Mainline hook context injection

internal/hooks/agents/pi/pi_test.go

@@ -4,7 +4,9 @@ import (
 	"context"
 	"encoding/json"
 	"os"
+	"os/exec"
 	"path/filepath"
+	"runtime"
 	"strconv"
 	"strings"
 	"testing"
@@ -42,6 +44,11 @@ func TestInstallWritesManagedExtensionAndIsIdempotent(t *testing.T) {
 		`debug("running " + hookName + " in " + commandCwd);`,
 		`debug("exited " + hookName + " with code " + code`,
 		`systemPromptAppend`,
+		`const HOOK_STDOUT_MAX_CHARS = 256 * 1024;`,
+		`const MAINLINE_CONTEXT_MAX_CHARS = 48 * 1024;`,
+		`hook context exceeded Pi safety budget`,
+		`combined hook context exceeded Pi safety budget`,
+		`stdout from " + hookName + " exceeded safety budget`,
 	} {
 		if !strings.Contains(text, want) {
 			t.Fatalf("extension missing %q:\n%s", want, text)
@@ -68,6 +75,76 @@ func TestInstallWritesManagedExtensionAndIsIdempotent(t *testing.T) {
 	}
 }
 
+func TestGeneratedExtensionEnforcesContextBudgets(t *testing.T) {
+	if runtime.GOOS == "windows" {
+		t.Skip("test uses an executable script as the fake mainline command")
+	}
+	node, err := exec.LookPath("node")
+	if err != nil {
+		t.Skip("node is not available")
+	}
+	if err := exec.Command(node, "--experimental-strip-types", "--input-type=module", "--eval", "").Run(); err != nil {
+		t.Skip("node does not support TypeScript type stripping")
+	}
+
+	dir := t.TempDir()
+	fakeMainline := filepath.Join(dir, "fake-mainline")
+	fakeSource := `#!/usr/bin/env node
+const hook = process.argv.at(-1);
+const scenario = process.env.FAKE_MAINLINE_SCENARIO;
+let size = 2;
+if (scenario === "combined") size = hook === "session-start" ? 40000 : 20000;
+if (scenario === "single") size = hook === "user-prompt-submit" ? 50000 : 2;
+if (scenario === "stdout") size = hook === "user-prompt-submit" ? 300000 : 2;
+process.stdout.write(JSON.stringify({systemPromptAppend: "X".repeat(size)}));
+`
+	if err := os.WriteFile(fakeMainline, []byte(fakeSource), 0o755); err != nil {
+		t.Fatal(err)
+	}
+	if _, err := (Agent{}).Install(dir, hooks.InstallOptions{BinPath: fakeMainline}); err != nil {
+		t.Fatal(err)
+	}
+
+	harnessPath := filepath.Join(dir, "probe.mjs")
+	harness := `import mainlinePiExtension from "./.pi/extensions/mainline.ts";
+const handlers = new Map();
+mainlinePiExtension({on: (name, handler) => handlers.set(name, handler)});
+const ctx = {cwd: process.cwd(), sessionManager: {getSessionFile: () => "session.jsonl"}};
+await handlers.get("session_start")({reason: "new"}, ctx);
+const result = await handlers.get("before_agent_start")({prompt: "test", systemPrompt: "BASE"}, ctx);
+const prompt = result?.systemPrompt ?? "";
+process.stdout.write(JSON.stringify({
+  warning: prompt.includes("Mainline hook context omitted"),
+  leaked: prompt.includes("XXXXXXXXXX"),
+}));
+`
+	if err := os.WriteFile(harnessPath, []byte(harness), 0o644); err != nil {
+		t.Fatal(err)
+	}
+
+	for _, scenario := range []string{"single", "combined", "stdout"} {
+		t.Run(scenario, func(t *testing.T) {
+			cmd := exec.Command(node, "--experimental-strip-types", harnessPath)
+			cmd.Dir = dir
+			cmd.Env = append(os.Environ(), "FAKE_MAINLINE_SCENARIO="+scenario)
+			raw, err := cmd.Output()
+			if err != nil {
+				t.Fatalf("run generated extension: %v", err)
+			}
+			var got struct {
+				Warning bool `json:"warning"`
+				Leaked  bool `json:"leaked"`
+			}
+			if err := json.Unmarshal(raw, &got); err != nil {
+				t.Fatalf("decode probe output: %v\n%s", err, raw)
+			}
+			if !got.Warning || got.Leaked {
+				t.Fatalf("unexpected bounded context result: %#v", got)
+			}
+		})
+	}
+}
+
 func TestDefaultInstallUsesLocalDevInMainlineSourceRepo(t *testing.T) {
 	dir := t.TempDir()
 	if err := os.WriteFile(filepath.Join(dir, "go.mod"), []byte("module github.com/mainline-org/mainline\n"), 0o644); err != nil {

internal/hooks/agents/pi/templates/mainline.ts.tmpl

@@ -9,6 +9,10 @@ const MAINLINE_COMMAND = {{ .Command }};
 const MAINLINE_ARGS_PREFIX = {{ .ArgsPrefix }};
 const MAINLINE_COMMAND_CWD: string | undefined = {{ .CommandCwd }};
 const HOOK_TIMEOUT_MS = 30000;
+const HOOK_STDOUT_MAX_CHARS = 256 * 1024;
+const MAINLINE_CONTEXT_MAX_CHARS = 48 * 1024;
+const CONTEXT_LIMIT_WARNING =
+	"# Mainline hook context omitted\n\nThe hook returned more context than Pi's safety budget permits. Run `mainline preflight --json`, `mainline status --json`, and task-specific `mainline context` / `mainline show` commands to load details on demand.";
 
 type HookOutput = {
 	systemPromptAppend?: string;
@@ -29,7 +33,21 @@ function sessionId(ctx: { sessionManager?: { getSessionFile?: () => string | und
 }
 
 function extractContext(out: HookOutput | undefined): string | undefined {
-	return out?.systemPromptAppend ?? out?.additionalContext ?? out?.context;
+	const context = out?.systemPromptAppend ?? out?.additionalContext ?? out?.context;
+	if (!context) return undefined;
+	if (context.length <= MAINLINE_CONTEXT_MAX_CHARS) return context;
+	debug("hook context exceeded Pi safety budget; injecting compact warning");
+	return CONTEXT_LIMIT_WARNING;
+}
+
+function combineContexts(parts: Array<string | undefined>): string | undefined {
+	const combined = parts.filter(
+		(part): part is string => typeof part === "string" && part.length > 0,
+	).join("\n\n");
+	if (!combined) return undefined;
+	if (combined.length <= MAINLINE_CONTEXT_MAX_CHARS) return combined;
+	debug("combined hook context exceeded Pi safety budget; injecting compact warning");
+	return CONTEXT_LIMIT_WARNING;
 }
 
 async function runMainlineHook(
@@ -47,25 +65,37 @@ async function runMainlineHook(
 
 		let stdout = "";
 		let stderr = "";
+		let settled = false;
+		const finish = (value: HookOutput | undefined): void => {
+			if (settled) return;
+			settled = true;
+			clearTimeout(timer);
+			resolve(value);
+		};
 		const timer = setTimeout(() => {
 			debug("timed out running " + hookName);
 			child.kill("SIGTERM");
-			resolve(undefined);
+			finish(undefined);
 		}, HOOK_TIMEOUT_MS);
 
 		child.stdout.on("data", (chunk) => {
+			if (settled) return;
 			stdout += String(chunk);
+			if (stdout.length > HOOK_STDOUT_MAX_CHARS) {
+				debug("stdout from " + hookName + " exceeded safety budget");
+				child.kill("SIGTERM");
+				finish({ systemPromptAppend: CONTEXT_LIMIT_WARNING });
+			}
 		});
 		child.stderr.on("data", (chunk) => {
 			stderr += String(chunk);
 		});
 		child.on("error", (err) => {
-			clearTimeout(timer);
 			debug("failed to run " + hookName + ": " + err.message);
-			resolve(undefined);
+			finish(undefined);
 		});
 		child.on("close", (code, signal) => {
-			clearTimeout(timer);
+			if (settled) return;
 			if (signal) {
 				debug("exited " + hookName + " with signal " + signal);
 			} else if (code !== 0) {
@@ -74,14 +104,14 @@ async function runMainlineHook(
 			if (stderr.trim()) debug("stderr from " + hookName + ": " + stderr.trim());
 			const text = stdout.trim();
 			if (!text) {
-				resolve(undefined);
+				finish(undefined);
 				return;
 			}
 			try {
-				resolve(JSON.parse(text) as HookOutput);
+				finish(JSON.parse(text) as HookOutput);
 			} catch {
 				debug("non-JSON hook output from " + hookName);
-				resolve(undefined);
+				finish(undefined);
 			}
 		});
 		child.stdin.on("error", (err) => {
@@ -92,7 +122,7 @@ async function runMainlineHook(
 			child.stdin.end(JSON.stringify(payload));
 		} catch (err) {
 			debug("failed to write payload for " + hookName + ": " + String(err));
-			resolve(undefined);
+			finish(undefined);
 		}
 	});
 }
@@ -126,12 +156,10 @@ export default function mainlinePiExtension(pi: ExtensionAPI) {
 			},
 			ctx.cwd,
 		);
-		const parts = [pendingSessionContext, extractContext(out)].filter(
-			(part): part is string => typeof part === "string" && part.length > 0,
-		);
+		const context = combineContexts([pendingSessionContext, extractContext(out)]);
 		pendingSessionContext = undefined;
-		if (parts.length === 0) return undefined;
-		return { systemPrompt: event.systemPrompt + "\n\n" + parts.join("\n\n") };
+		if (!context) return undefined;
+		return { systemPrompt: event.systemPrompt + "\n\n" + context };
 	});
 
 	pi.on("agent_end", async (event, ctx) => {