resident_budget_recovery_test.go Raw History Folder

1package subagents
2
3import (
4	"context"
5	"encoding/json"
6	"errors"
7	"fmt"
8	"os"
9	"path/filepath"
10	"strings"
11	"testing"
12	"time"
13
14	"github.com/bnema/zut/packages/core"
15	"github.com/bnema/zut/packages/provider"
16)
17
18func TestResidentBudgetExhaustionHandoffAndRecovery(t *testing.T) {
19	for _, mode := range []string{"read-only", "shared", "worktree"} {
20		for _, restart := range []bool{false, true} {
21			t.Run(fmt.Sprintf("%s/restart=%t", mode, restart), func(t *testing.T) {
22				root, workspace := t.TempDir(), t.TempDir()
23				if mode == "worktree" {
24					workspace = initTestRepo(t)
25				}
26				spec := ResidentChildSpec{ID: "budget-recovery", SessionID: "session", InitialTurnID: "initial", Provider: "openai", Model: "test", Required: true, BudgetLimit: 100, BudgetSource: "model_context", RepositoryRoot: workspace, WorkspaceMode: WorkspaceShared}
27				if mode == "worktree" {
28					spec.WorkspaceMode, spec.WorkspaceCapture = WorkspaceWorktree, CapturePatch
29				}
30				factory := func(spec ResidentChildSpec, journal *ResidentJournal) (ResidentTurnRunner, error) {
31					return func(_ context.Context, prompt string) error {
32						if prompt == "continue remaining work" {
33							if journal.BudgetBaseline() != 100 {
34								return fmt.Errorf("recovery baseline = %d", journal.BudgetBaseline())
35							}
36							messages, err := ReadResidentTranscriptMessages(journal.Dir())
37							if err != nil || len(messages) < 3 {
38								return fmt.Errorf("missing restored progress: %d messages, %v", len(messages), err)
39							}
40							if mode != "read-only" {
41								data, err := os.ReadFile(filepath.Join(spec.Workspace, "README.md"))
42								if err != nil || string(data) != "partial artifact\n" {
43									return fmt.Errorf("lost artifact: %q, %v", data, err)
44								}
45							}
46							if err := journal.RecordAgentEvent(core.EvUsage{Cumulative: provider.Usage{InputTokens: 105}}); err != nil {
47								return err
48							}
49							return journal.RecordAgentEvent(core.EvAssistantMessage{Message: provider.Message{Role: provider.RoleAssistant, Content: []provider.Content{provider.TextBlock{Text: "verified remaining work"}}}})
50						}
51						tool, output := "read", "inspected README; verification unfinished"
52						if mode != "read-only" {
53							tool, output = "write", "wrote README.md; verification unfinished"
54							if err := os.WriteFile(filepath.Join(spec.Workspace, "README.md"), []byte("partial artifact\n"), 0o600); err != nil {
55								return err
56							}
57						}
58						args := json.RawMessage(`{"path":"README.md"}`)
59						for _, event := range []core.AgentEvent{
60							core.EvUserMessage{Message: provider.Message{Role: provider.RoleUser, Content: []provider.Content{provider.TextBlock{Text: prompt}}}},
61							core.EvAssistantMessage{Message: provider.Message{Role: provider.RoleAssistant, Content: []provider.Content{provider.TextBlock{Text: "working on README.md; verification unfinished"}, provider.ToolCallBlock{ID: "call", Name: tool, Arguments: args}}}},
62							core.EvToolCall{ID: "call", Name: tool, Args: args},
63							core.EvToolResult{ID: "call", Result: core.ToolResult{Content: []provider.Content{provider.TextBlock{Text: output}}}},
64							core.EvUsage{Cumulative: provider.Usage{InputTokens: 100}},
65						} {
66							if err := journal.RecordAgentEvent(event); err != nil {
67								return err
68							}
69						}
70						return ErrBudgetExceeded
71					}, nil
72				}
73				manager := NewResidentManager(root, factory)
74				t.Cleanup(func() { _ = manager.Close(context.Background()) })
75				completions := make(chan ResidentCompletion, 2)
76				manager.SetCompletionObserver(func(c ResidentCompletion) { completions <- c })
77				if _, err := manager.Spawn(t.Context(), spec, "inspect and verify README"); err != nil {
78					t.Fatal(err)
79				}
80				completion := awaitBudgetCompletion(t, completions)
81				if !errors.Is(completion.Err, ErrBudgetExceeded) || !strings.Contains(completion.Summary, "verification unfinished") {
82					t.Fatalf("terminal notification = %#v", completion)
83				}
84				assertBudgetPartial := func() ResidentResult {
85					t.Helper()
86					snapshot, ok := manager.SnapshotFor(spec.ID)
87					if !ok || snapshot.State != ResidentBudgetExhausted || snapshot.Budget.State != BudgetExceeded || len(manager.UnmetRequired()) != 1 {
88						t.Fatalf("partial status = %#v; unmet = %#v", snapshot, manager.UnmetRequired())
89					}
90					result, err := manager.Result(spec.ID)
91					if err != nil || result.ErrorCode != "budget_exhausted" || !strings.Contains(result.Handoff, "README.md") || len(result.Handoff) > residentHandoffBytes {
92						t.Fatalf("saved partial result = %#v, %v", result, err)
93					}
94					return result
95				}
96				result := assertBudgetPartial()
97				if mode == "worktree" {
98					if result.PatchRef != PatchRef(spec.ID) || len(result.ChangedFiles) != 1 {
99						t.Fatalf("missing partial patch: %#v", result)
100					}
101					patch, err := os.ReadFile(filepath.Join(root, spec.ID, residentPatchName))
102					if err != nil || !strings.Contains(string(patch), "partial artifact") {
103						t.Fatalf("partial patch = %q, %v", patch, err)
104					}
105				}
106				if restart {
107					if err := manager.Close(t.Context()); err != nil {
108						t.Fatal(err)
109					}
110					manager = NewResidentManager(root, factory)
111					manager.SetCompletionObserver(func(c ResidentCompletion) { completions <- c })
112					if errs := manager.Reconcile(); len(errs) != 0 {
113						t.Fatal(errs)
114					}
115					assertBudgetPartial()
116				}
117				if err := manager.Resume(t.Context(), spec.ID, "continue remaining work"); err != nil {
118					t.Fatal(err)
119				}
120				if completion := awaitBudgetCompletion(t, completions); completion.Err != nil {
121					t.Fatal(completion.Err)
122				}
123				if unmet := manager.UnmetRequired(); len(unmet) != 0 {
124					t.Fatalf("successful recovery remains unmet: %#v", unmet)
125				}
126				snapshot, _ := manager.SnapshotFor(spec.ID)
127				if snapshot.Budget.Used != 5 || snapshot.Budget.Limit != 100 || snapshot.Usage.InputTokens != 105 {
128					t.Fatalf("recovery lost cumulative accounting: %#v", snapshot)
129				}
130				if err := manager.Close(t.Context()); err != nil {
131					t.Fatal(err)
132				}
133				metadata, err := ReconcileResidentJournal(filepath.Join(root, spec.ID))
134				if err != nil || metadata.BudgetBaseline != 100 || metadata.Usage.InputTokens != 105 {
135					t.Fatalf("recovery baseline did not survive reconciliation: %#v, %v", metadata, err)
136				}
137			})
138		}
139	}
140}
141
142func awaitBudgetCompletion(t *testing.T, completions <-chan ResidentCompletion) ResidentCompletion {
143	t.Helper()
144	select {
145	case completion := <-completions:
146		return completion
147	case <-time.After(5 * time.Second):
148		t.Fatal("no terminal budget notification")
149		return ResidentCompletion{}
150	}
151}
resident_budget_recovery_test.go RawHistoryFolder

resident_budget_recovery_test.go Raw History Folder
