Guard task completion after failed tools

Author: timothyerwin

internal/agent/runtime/apply_patch_test.go

@@ -1,6 +1,7 @@
 package runtime
 
 import (
+	"bytes"
 	"context"
 	"encoding/json"
 	"io"
@@ -15,13 +16,18 @@ import (
 )
 
 func applyPatchSess(t *testing.T, root string) *Session {
+	t.Helper()
+	return applyPatchSessWithOut(t, root, io.Discard)
+}
+
+func applyPatchSessWithOut(t *testing.T, root string, out io.Writer) *Session {
 	t.Helper()
 	st, err := store.Open(context.Background(), filepath.Join(t.TempDir(), "state.db"))
 	if err != nil {
 		t.Fatal(err)
 	}
 	t.Cleanup(func() { st.Close() })
-	return newSess(st, captureProviderNil{}, root, "allow-all", permissions.ModeAllowAll, io.Discard)
+	return newSess(st, captureProviderNil{}, root, "allow-all", permissions.ModeAllowAll, out)
 }
 
 // A multi-file patch where every edit is valid applies them ALL.
@@ -89,3 +95,45 @@ func TestApplyPatchRollbackDeletesCreated(t *testing.T) {
 		t.Fatal("a file created by a rolled-back patch must be deleted")
 	}
 }
+
+func TestEditFileWriteRendersCreatedFilePreview(t *testing.T) {
+	root := t.TempDir()
+	var out bytes.Buffer
+	s := applyPatchSessWithOut(t, root, &out)
+
+	in, _ := json.Marshal(tools.EditFileInput{
+		Path:      "new.go",
+		OldString: "",
+		NewString: "package p\n\nfunc Answer() int {\n\treturn 42\n}\n",
+	})
+	if r := s.editFile(context.Background(), in); r.isError {
+		t.Fatalf("write should apply: %q", r.text())
+	}
+	got := ansiSeq.ReplaceAllString(out.String(), "")
+	for _, want := range []string{"Write(", "new.go", "Created file (5 lines)", "func Answer() int", "return 42"} {
+		if !strings.Contains(got, want) {
+			t.Errorf("write preview missing %q in:\n%s", want, got)
+		}
+	}
+}
+
+func TestApplyPatchWriteRendersCreatedFilePreview(t *testing.T) {
+	root := t.TempDir()
+	var out bytes.Buffer
+	s := applyPatchSessWithOut(t, root, &out)
+
+	in, _ := json.Marshal(tools.ApplyPatchInput{Edits: []tools.EditFileInput{{
+		Path:      "new.go",
+		OldString: "",
+		NewString: "package p\n\nfunc Answer() int {\n\treturn 42\n}\n",
+	}}})
+	if r := s.applyPatch(context.Background(), in); r.isError {
+		t.Fatalf("patch should apply: %q", r.text())
+	}
+	got := ansiSeq.ReplaceAllString(out.String(), "")
+	for _, want := range []string{"Write(", "new.go", "Created file (5 lines)", "func Answer() int", "return 42"} {
+		if !strings.Contains(got, want) {
+			t.Errorf("apply_patch write preview missing %q in:\n%s", want, got)
+		}
+	}
+}

internal/agent/runtime/diagnostics.go

@@ -158,7 +158,7 @@ func (s *Session) lspDiagResult(lang, path string, diags []lsp.Diagnostic) toolR
 	out := strings.TrimRight(b.String(), "\n")
 	s.toolLine(true, "Diagnostics", lang+" (lsp)", strconv.Itoa(len(diags))+" issue(s)", true)
 	s.toolResult(linesPreview(out, maxDiagPreviewLines)) // show WHAT the issues are, not just a red count
-	return textResult(s.redactor.Redact(truncate(out, maxToolOutput)))
+	return errResult(s.redactor.Redact(truncate(out, maxToolOutput)))
 }
 
 // diagnosticsLang picks the language from the path extension, else the repo markers.
@@ -238,5 +238,5 @@ func (s *Session) diagResult(lang, out string) toolResult {
 	n := strings.Count(out, "\n") + 1
 	s.toolLine(true, "Diagnostics", lang, strconv.Itoa(n)+" line(s)", true)
 	s.toolResult(linesPreview(out, maxDiagPreviewLines))
-	return textResult(s.redactor.Redact(truncate(out, maxToolOutput)))
+	return errResult(s.redactor.Redact(truncate(out, maxToolOutput)))
 }

internal/agent/runtime/diagnostics_test.go

@@ -45,6 +45,9 @@ func TestGoDiagnostics(t *testing.T) {
 	// Introduce a type error → diagnostics surfaces it with the file name.
 	write("bad.go", "package tmpdiag\n\nvar X int = \"not an int\"\n")
 	r := s.diagnosticsTool(context.Background(), in)
+	if !r.isError {
+		t.Fatal("diagnostics with compiler output should be an error tool result")
+	}
 	if !strings.Contains(r.text(), "bad.go") {
 		t.Fatalf("compile error not surfaced with file: %q", r.text())
 	}

internal/agent/runtime/exec.go

@@ -154,16 +154,18 @@ func isParallelSafe(name string) bool {
 func (s *Session) execute(ctx context.Context, u wire.Block) toolResult {
 	s.mu.Lock()
 	s.metrics.toolCalls++
+	seq := s.metrics.toolCalls
 	s.mu.Unlock()
 	// Redact the tool input before it is persisted — the model may pass a secret
 	// (e.g. a bearer token in a command).
 	s.emit(ctx, events.KindToolCalled, map[string]any{"tool": u.Name, "input": s.redactor.Redact(string(u.Input))})
 	tr := s.dispatch(ctx, u)
+	s.mu.Lock()
 	if tr.isError {
-		s.mu.Lock()
 		s.metrics.toolErrors++
-		s.mu.Unlock()
 	}
+	s.noteCompletionBlockerLocked(seq, u, tr)
+	s.mu.Unlock()
 	// Central guarantee: no known secret value ever reaches the model — redact
 	// the text of every block (images carry no text).
 	for i := range tr.blocks {
@@ -174,6 +176,62 @@ func (s *Session) execute(ctx context.Context, u wire.Block) toolResult {
 	return tr
 }
 
+// noteCompletionBlockerLocked records failed tools that are themselves a
+// deliverable the user can inspect (PR creation, artifact publish, etc.). The
+// todo tracker refuses to mark work done while one is unresolved, so a model
+// cannot paper over "GitHub(create PR) · failed" with "Task(done)".
+func (s *Session) noteCompletionBlockerLocked(seq int, u wire.Block, tr toolResult) {
+	label, ok := completionBlockerLabel(u)
+	if !ok {
+		return
+	}
+	if tr.isError {
+		s.metrics.blockerSeq = seq
+		s.metrics.blockerLabel = label
+		s.metrics.blockerDetail = s.redactor.Redact(firstLine(tr.text()))
+		return
+	}
+	if s.metrics.blockerLabel == label {
+		s.metrics.blockerSeq = 0
+		s.metrics.blockerLabel = ""
+		s.metrics.blockerDetail = ""
+	}
+}
+
+func completionBlockerLabel(u wire.Block) (string, bool) {
+	switch u.Name {
+	case tools.RunTests:
+		return "test run", true
+	case tools.Diagnostics:
+		return "diagnostics check", true
+	case tools.GitHub:
+		var in tools.GitHubInput
+		if json.Unmarshal(u.Input, &in) != nil {
+			return "", false
+		}
+		switch in.Action {
+		case "pr_create":
+			return "GitHub create PR", true
+		case "comment":
+			return "GitHub comment", true
+		}
+	case tools.Artifact:
+		var in tools.ArtifactInput
+		if json.Unmarshal(u.Input, &in) != nil {
+			return "", false
+		}
+		switch in.Action {
+		case "publish":
+			return "artifact publish", true
+		case "update":
+			return "artifact update", true
+		case "delete":
+			return "artifact delete", true
+		}
+	}
+	return "", false
+}
+
 func (s *Session) dispatch(ctx context.Context, u wire.Block) toolResult {
 	// A reader sub-agent has no mutating tools; reject them defensively even if one
 	// is somehow proposed (they aren't advertised — see toolDefs). Bash is the
@@ -1375,7 +1433,7 @@ func (s *Session) editFile(ctx context.Context, input json.RawMessage) toolResul
 		verb = "Write"
 	}
 	safeDiff := s.redactor.Redact(res.Diff)
-	newContent := s.redactor.Redact(in.NewString)
+	newContent := s.newFilePreviewContent(res.Path, in.NewString)
 	if secrets.IsSecretPath(in.Path) {
 		// Editing a credential file: mask values in both the diff and new content.
 		safeDiff = secrets.RedactSecretFile(safeDiff)
@@ -1536,13 +1594,19 @@ func (s *Session) applyPatch(ctx context.Context, input json.RawMessage) toolRes
 
 	var b strings.Builder
 	fmt.Fprintf(&b, "OK, applied %d edits across %d file(s) atomically:\n", len(applied), len(dedupePaths(paths)))
-	for _, res := range applied {
+	for i, res := range applied {
 		verb := "Update"
 		if res.Created {
 			verb = "Write"
 		}
 		s.toolLine(true, verb, res.Path, "", false)
-		if d := s.redactor.Redact(res.Diff); d != "" {
+		if res.Created {
+			newContent := s.newFilePreviewContent(res.Path, in.Edits[i].NewString)
+			if secrets.IsSecretPath(res.Path) {
+				newContent = secrets.RedactSecretFile(newContent)
+			}
+			renderNewFile(s.out, newContent, res.Path, s.diffWidth())
+		} else if d := s.redactor.Redact(res.Diff); d != "" {
 			renderDiff(s.out, d, res.Path, s.diffWidth())
 		}
 		fmt.Fprintf(&b, "  %s %s\n", verb, res.Path)
@@ -1555,6 +1619,15 @@ func (s *Session) applyPatch(ctx context.Context, input json.RawMessage) toolRes
 	return textResult(out)
 }
 
+func (s *Session) newFilePreviewContent(path, fallback string) string {
+	if abs, err := safeJoin(s.root, path); err == nil {
+		if b, err := os.ReadFile(abs); err == nil {
+			return s.redactor.Redact(string(b))
+		}
+	}
+	return s.redactor.Redact(fallback)
+}
+
 // dedupePaths returns the unique paths in first-seen order.
 func dedupePaths(paths []string) []string {
 	seen := map[string]bool{}

internal/agent/runtime/github.go

@@ -7,6 +7,7 @@ import (
 	"os/exec"
 	"strconv"
 	"strings"
+	"time"
 
 	"github.com/memcode-ai/memcode/internal/agent/permissions"
 	"github.com/memcode-ai/memcode/internal/agent/tools"
@@ -51,6 +52,10 @@ func (s *Session) githubTool(ctx context.Context, input json.RawMessage) toolRes
 		if strings.TrimSpace(in.Title) == "" || strings.TrimSpace(in.Body) == "" {
 			return errResult("pr_create needs a `title` and `body`.")
 		}
+		if msg := s.prCreatePreflight(ctx, in.Base); msg != "" {
+			s.toolLine(true, "GitHub", "create PR", "blocked: "+clip(msg, 160), true)
+			return errResult(msg)
+		}
 		if ok, reason := s.gate(ctx, permissions.Medium, false, ApprovalRequest{
 			Title: in.Title, Label: "Open a GitHub PR", Detail: "gh pr create", Risk: permissions.Medium.String(),
 		}); !ok {
@@ -78,6 +83,56 @@ func (s *Session) githubTool(ctx context.Context, input json.RawMessage) toolRes
 	}
 }
 
+// prCreatePreflight keeps PR creation at the end of the work: a PR without
+// committed, pushed branch work is at best a noisy gh prompt and at worst a
+// false "done" marker. Return "" when ready, else a user/model-facing reason.
+func (s *Session) prCreatePreflight(ctx context.Context, base string) string {
+	run := func(args ...string) (string, error) {
+		cctx, cancel := context.WithTimeout(ctx, 10*time.Second)
+		defer cancel()
+		cmd := exec.CommandContext(cctx, "git", args...)
+		cmd.Dir = s.root
+		out, err := cmd.CombinedOutput()
+		return strings.TrimSpace(string(out)), err
+	}
+	if out, err := run("status", "--porcelain"); err != nil {
+		return "pr_create preflight failed: git status failed: " + firstLine(out)
+	} else if strings.TrimSpace(out) != "" {
+		return "pr_create is a final step: the worktree still has uncommitted changes. Commit the branch work and push it before creating the PR."
+	}
+	branch, err := run("branch", "--show-current")
+	if err != nil || branch == "" {
+		return "pr_create needs a named feature branch; commit on a branch before creating the PR."
+	}
+	if branch == "main" || branch == "master" {
+		return "pr_create needs a feature branch, not " + branch + ". Create and push a branch before opening the PR."
+	}
+	if strings.TrimSpace(base) == "" {
+		base = "origin/main"
+	} else if !strings.Contains(base, "/") {
+		base = "origin/" + strings.TrimSpace(base)
+	}
+	ahead, err := run("rev-list", "--count", base+"..HEAD")
+	if err != nil {
+		return "pr_create preflight failed: cannot compare against " + base + " — fetch the base branch or pass a valid base."
+	}
+	if ahead == "0" {
+		return "pr_create has nothing to open: this branch has no commits ahead of " + base + ". Commit the completed work first."
+	}
+	upstream, err := run("rev-parse", "--abbrev-ref", "--symbolic-full-name", "@{u}")
+	if err != nil || upstream == "" {
+		return "pr_create needs the current branch pushed upstream. Run git push -u origin " + branch + " after committing, then create the PR."
+	}
+	unpushed, err := run("rev-list", "--count", upstream+"..HEAD")
+	if err != nil {
+		return "pr_create preflight failed: cannot compare against upstream " + upstream + ". Push the branch again, then create the PR."
+	}
+	if unpushed != "0" {
+		return "pr_create needs the latest commits pushed first (" + unpushed + " local commit(s) are not on " + upstream + ")."
+	}
+	return ""
+}
+
 // ghRun invokes gh with the given args in the repo root, returns its output (redacted,
 // truncated), and surfaces a tool line. A non-zero exit returns the error text.
 func (s *Session) ghRun(ctx context.Context, label string, args ...string) toolResult {
@@ -91,7 +146,11 @@ func (s *Session) ghRun(ctx context.Context, label string, args ...string) toolR
 	out, err := cmd.CombinedOutput()
 	text := strings.TrimSpace(string(out))
 	if err != nil {
-		s.toolLine(true, "GitHub", label, "failed", true)
+		reason := strings.TrimSpace(firstLine(text))
+		if reason == "" {
+			reason = err.Error()
+		}
+		s.toolLine(true, "GitHub", label, "failed: "+clip(s.redactor.Redact(reason), 160), true)
 		return errResult(fmt.Sprintf("gh %s failed: %v\n%s", args[0], err, truncate(text, 2000)))
 	}
 	s.toolLine(true, "GitHub", label, "", false)

internal/agent/runtime/github_test.go

@@ -4,6 +4,8 @@ import (
 	"context"
 	"encoding/json"
 	"io"
+	"os"
+	"os/exec"
 	"path/filepath"
 	"strings"
 	"testing"
@@ -54,3 +56,58 @@ func TestGitHubToolAdvertised(t *testing.T) {
 		t.Fatal("github tool should be advertised in normal chat")
 	}
 }
+
+func TestGitHubPRCreatePreflightBlocksDirtyWorktree(t *testing.T) {
+	root := t.TempDir()
+	git(t, root, "init", "-b", "feat/pr")
+	if err := os.WriteFile(filepath.Join(root, "draft.txt"), []byte("uncommitted\n"), 0o644); err != nil {
+		t.Fatal(err)
+	}
+
+	s := githubSess(t)
+	s.root = root
+	in, _ := json.Marshal(tools.GitHubInput{Action: "pr_create", Title: "Draft PR", Body: "Body"})
+	r := s.githubTool(context.Background(), in)
+	if !r.isError {
+		t.Fatal("pr_create should fail before gh when the worktree is dirty")
+	}
+	if !strings.Contains(r.text(), "worktree still has uncommitted changes") {
+		t.Fatalf("dirty-worktree reason missing: %q", r.text())
+	}
+}
+
+func TestGitHubPRCreatePreflightBlocksBranchWithNoCommitsAhead(t *testing.T) {
+	root := t.TempDir()
+	git(t, root, "init", "-b", "main")
+	git(t, root, "config", "user.email", "test@example.com")
+	git(t, root, "config", "user.name", "Test User")
+	if err := os.WriteFile(filepath.Join(root, "README.md"), []byte("# test\n"), 0o644); err != nil {
+		t.Fatal(err)
+	}
+	git(t, root, "add", "README.md")
+	git(t, root, "commit", "-m", "initial")
+	git(t, root, "update-ref", "refs/remotes/origin/main", "HEAD")
+	git(t, root, "switch", "-c", "feat/noop")
+
+	s := githubSess(t)
+	s.root = root
+	in, _ := json.Marshal(tools.GitHubInput{Action: "pr_create", Title: "Draft PR", Body: "Body"})
+	r := s.githubTool(context.Background(), in)
+	if !r.isError {
+		t.Fatal("pr_create should fail before gh when the branch has no commits")
+	}
+	if !strings.Contains(r.text(), "no commits ahead of origin/main") {
+		t.Fatalf("no-ahead reason missing: %q", r.text())
+	}
+}
+
+func git(t *testing.T, root string, args ...string) string {
+	t.Helper()
+	cmd := exec.Command("git", args...)
+	cmd.Dir = root
+	out, err := cmd.CombinedOutput()
+	if err != nil {
+		t.Fatalf("git %s failed: %v\n%s", strings.Join(args, " "), err, out)
+	}
+	return strings.TrimSpace(string(out))
+}

internal/agent/runtime/metricsstate.go

@@ -12,6 +12,9 @@ type metricsState struct {
 	didVerify       bool              // any build/test command run this session
 	lastEditSeq     int               // tool-call seq of the most recent edit
 	lastVerifyOKSeq int               // tool-call seq of the most recent passing verification
+	blockerSeq      int               // failed deliverable tool that must be resolved before a todo is done
+	blockerLabel    string            // user-facing name of the blocked deliverable (e.g. GitHub create PR)
+	blockerDetail   string            // first useful error line to feed back on refused todo completion
 	readHashes      map[string]string // path → content hash when last read/wrote (stale-edit guard; lazily inited)
 	reportsSpilled  int               // sub-agent reports written to .memcode/sessions/<id>/reports/ (names the files)
 }