Skip to content

Commit b131ec8

Browse files
nathanwhitnathanwhit
committed
replanLoopWithOptions` rebuilt the full store snapshot repeatedly per
1 parent ee842c6 commit b131ec8

5 files changed

Lines changed: 263 additions & 19 deletions

File tree

‎internal/orchestrator/pull_request_service.go‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2044,6 +2044,10 @@ func latestPullRequestFollowUpIsQueued(snapshot core.Snapshot, taskID string) bo
20442044
}
20452045

20462046
func pendingPullRequestFeedback(snapshot core.Snapshot, taskID string) []PullRequestFeedbackItem {
2047+
return pendingPullRequestFeedbackFromSnapshot(snapshot, taskID)
2048+
}
2049+
2050+
func pendingPullRequestFeedbackFromSnapshot(snapshot core.Snapshot, taskID string) []PullRequestFeedbackItem {
20472051
pullRequests := map[string]core.PullRequest{}
20482052
for _, pr := range snapshot.PullRequests {
20492053
if pr.TaskID == taskID {

‎internal/orchestrator/replan_prompt_budgeter.go‎

Lines changed: 89 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -167,13 +167,15 @@ func (b ReplanPromptBudgeter) compactRecentResults(results []WorkerTurnResult, b
167167
}
168168
recent = append(recent, b.compactWorkerResult(result))
169169
}
170-
for approxJSONTokens(recent) > b.RecentResultsTokens {
170+
recentTokens := newJSONArrayTokenSizer(recent)
171+
for recentTokens.tokens() > b.RecentResultsTokens {
171172
dropped := false
172173
for i, result := range recent {
173174
if isHighPriorityPromptResult(result, blocked) {
174175
continue
175176
}
176177
recent = append(recent[:i], recent[i+1:]...)
178+
recentTokens.drop(i)
177179
dropped = true
178180
break
179181
}
@@ -216,7 +218,8 @@ func (b ReplanPromptBudgeter) compactWorkerResult(result WorkerTurnResult) Worke
216218

217219
func (b ReplanPromptBudgeter) compactContextLedger(entries []ContextLedgerEntry) []ContextLedgerEntry {
218220
compact := compactContextLedgerForPrompt(entries)
219-
for approxJSONTokens(compact) > b.ContextLedgerTokens && len(compact) > 0 {
221+
compactTokens := newJSONArrayTokenSizer(compact)
222+
for compactTokens.tokens() > b.ContextLedgerTokens && len(compact) > 0 {
220223
dropIndex := -1
221224
for i, entry := range compact {
222225
if entry.Error == "" && !strings.Contains(entry.Kind, "candidate") {
@@ -228,6 +231,7 @@ func (b ReplanPromptBudgeter) compactContextLedger(entries []ContextLedgerEntry)
228231
dropIndex = 0
229232
}
230233
compact = append(compact[:dropIndex], compact[dropIndex+1:]...)
234+
compactTokens.drop(dropIndex)
231235
}
232236
return compact
233237
}
@@ -237,10 +241,13 @@ func (b ReplanPromptBudgeter) compactPullRequestFeedback(items []PullRequestFeed
237241
for i := range compact {
238242
compact[i].Prompt = truncateStringForPrompt(compact[i].Prompt, tokensToApproxChars(1000))
239243
}
240-
for approxJSONTokens(compact) > b.PullRequestFeedbackTokens && len(compact) > 0 {
244+
compactTokens := newJSONArrayTokenSizer(compact)
245+
for compactTokens.tokens() > b.PullRequestFeedbackTokens && len(compact) > 0 {
241246
compact[0].Prompt = truncateStringForPrompt(compact[0].Prompt, tokensToApproxChars(250))
242-
if approxJSONTokens(compact) > b.PullRequestFeedbackTokens {
247+
compactTokens.update(0, compact[0])
248+
if compactTokens.tokens() > b.PullRequestFeedbackTokens {
243249
compact = compact[1:]
250+
compactTokens.drop(0)
244251
}
245252
}
246253
return compact
@@ -259,7 +266,8 @@ func (b ReplanPromptBudgeter) compactArtifacts(artifacts []core.TaskArtifact) []
259266
}
260267
compact = append(compact, item)
261268
}
262-
for approxJSONTokens(compact) > b.ArtifactsTokens && len(compact) > 0 {
269+
compactTokens := newJSONArrayTokenSizer(compact)
270+
for compactTokens.tokens() > b.ArtifactsTokens && len(compact) > 0 {
263271
dropIndex := 0
264272
for i, artifact := range compact {
265273
if artifact.Kind != "worker_result_digest" {
@@ -268,12 +276,14 @@ func (b ReplanPromptBudgeter) compactArtifacts(artifacts []core.TaskArtifact) []
268276
}
269277
}
270278
compact = append(compact[:dropIndex], compact[dropIndex+1:]...)
279+
compactTokens.drop(dropIndex)
271280
}
272281
return compact
273282
}
274283

275284
func (b ReplanPromptBudgeter) degradeToTotalBudget(payload map[string]any, state ReplanPromptState) ReplanPromptState {
276-
for approxJSONTokens(payload) > b.TotalTokens {
285+
currentTokens := approxJSONTokens(payload)
286+
for currentTokens > b.TotalTokens {
277287
switch {
278288
case len(state.Artifacts) > 0:
279289
state.Artifacts = state.Artifacts[1:]
@@ -288,16 +298,18 @@ func (b ReplanPromptBudgeter) degradeToTotalBudget(payload map[string]any, state
288298
state.RecentResults[i].Changes.ChangedFiles = nil
289299
}
290300
state.RecoveryHint = truncateStringForPrompt(state.RecoveryHint, tokensToApproxChars(250))
301+
payload["state"] = state
291302
state.PromptBudget.ApproxTokens = approxJSONTokens(payload)
292303
return state
293304
}
294-
state.PromptBudget.ApproxTokens = approxJSONTokens(state)
295305
state.PromptBudget.RecentResultCount = len(state.RecentResults)
296306
state.PromptBudget.ArtifactCount = len(state.Artifacts)
297307
state.PromptBudget.ContextLedgerCount = len(state.ContextLedger)
298308
payload["state"] = state
309+
currentTokens = approxJSONTokens(payload)
310+
state.PromptBudget.ApproxTokens = currentTokens
299311
}
300-
state.PromptBudget.ApproxTokens = approxJSONTokens(payload)
312+
state.PromptBudget.ApproxTokens = currentTokens
301313
return state
302314
}
303315

@@ -317,8 +329,10 @@ func compactTaskSteeringForPrompt(items []string) []string {
317329
for i := range compact {
318330
compact[i] = truncateStringForPrompt(compact[i], tokensToApproxChars(1000))
319331
}
320-
for approxJSONTokens(compact) > 4000 && len(compact) > 1 {
332+
compactTokens := newJSONArrayTokenSizer(compact)
333+
for compactTokens.tokens() > 4000 && len(compact) > 1 {
321334
compact = compact[1:]
335+
compactTokens.drop(0)
322336
}
323337
return compact
324338
}
@@ -404,11 +418,76 @@ func approxJSONTokens(value any) int {
404418
return approxTokensForString(string(data))
405419
}
406420

421+
type jsonArrayTokenSizer[T any] struct {
422+
nilSlice bool
423+
itemBytes []int
424+
}
425+
426+
func newJSONArrayTokenSizer[T any](items []T) *jsonArrayTokenSizer[T] {
427+
sizer := &jsonArrayTokenSizer[T]{
428+
nilSlice: items == nil,
429+
itemBytes: make([]int, 0, len(items)),
430+
}
431+
for _, item := range items {
432+
sizer.itemBytes = append(sizer.itemBytes, jsonEncodedBytes(item))
433+
}
434+
return sizer
435+
}
436+
437+
func (s *jsonArrayTokenSizer[T]) tokens() int {
438+
return approxTokensForBytes(s.bytes())
439+
}
440+
441+
func (s *jsonArrayTokenSizer[T]) bytes() int {
442+
if s == nil || s.nilSlice {
443+
return len("null")
444+
}
445+
if len(s.itemBytes) == 0 {
446+
return len("[]")
447+
}
448+
total := len("[]") + len(s.itemBytes) - 1
449+
for _, bytes := range s.itemBytes {
450+
total += bytes
451+
}
452+
return total
453+
}
454+
455+
func (s *jsonArrayTokenSizer[T]) drop(index int) {
456+
if s == nil || index < 0 || index >= len(s.itemBytes) {
457+
return
458+
}
459+
s.nilSlice = false
460+
s.itemBytes = append(s.itemBytes[:index], s.itemBytes[index+1:]...)
461+
}
462+
463+
func (s *jsonArrayTokenSizer[T]) update(index int, item T) {
464+
if s == nil || index < 0 || index >= len(s.itemBytes) {
465+
return
466+
}
467+
s.nilSlice = false
468+
s.itemBytes[index] = jsonEncodedBytes(item)
469+
}
470+
471+
func jsonEncodedBytes(value any) int {
472+
data, err := json.Marshal(value)
473+
if err != nil {
474+
return 0
475+
}
476+
return len(data)
477+
}
478+
479+
func approxTokensForBytes(bytes int) int {
480+
if bytes <= 0 {
481+
return 0
482+
}
483+
return (bytes + 3) / 4
484+
}
485+
407486
func approxTokensForString(value string) int {
408487
if value == "" {
409488
return 0
410489
}
411-
return (len(value) + 3) / 4
490+
return approxTokensForBytes(len(value))
412491
}
413492

414493
func tokensToApproxChars(tokens int) int {

‎internal/orchestrator/replan_prompt_budgeter_test.go‎

Lines changed: 82 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,82 @@
1+
package orchestrator
2+
3+
import (
4+
"fmt"
5+
"strings"
6+
"testing"
7+
)
8+
9+
func TestJSONArrayTokenSizerMatchesMarshalEstimate(t *testing.T) {
10+
var nilItems []string
11+
if got, want := newJSONArrayTokenSizer(nilItems).tokens(), approxJSONTokens(nilItems); got != want {
12+
t.Fatalf("nil slice tokens = %d, want %d", got, want)
13+
}
14+
15+
items := []string{"alpha", strings.Repeat("beta", 20), "gamma"}
16+
sizer := newJSONArrayTokenSizer(items)
17+
if got, want := sizer.tokens(), approxJSONTokens(items); got != want {
18+
t.Fatalf("initial tokens = %d, want %d", got, want)
19+
}
20+
21+
items = append(items[:1], items[2:]...)
22+
sizer.drop(1)
23+
if got, want := sizer.tokens(), approxJSONTokens(items); got != want {
24+
t.Fatalf("tokens after drop = %d, want %d", got, want)
25+
}
26+
27+
items[0] = strings.Repeat("delta", 30)
28+
sizer.update(0, items[0])
29+
if got, want := sizer.tokens(), approxJSONTokens(items); got != want {
30+
t.Fatalf("tokens after update = %d, want %d", got, want)
31+
}
32+
33+
items = items[:0]
34+
sizer.drop(0)
35+
sizer.drop(0)
36+
if got, want := sizer.tokens(), approxJSONTokens(items); got != want {
37+
t.Fatalf("empty slice tokens = %d, want %d", got, want)
38+
}
39+
}
40+
41+
func TestCompactContextLedgerUsesBudget(t *testing.T) {
42+
budgeter := DefaultReplanPromptBudgeter()
43+
budgeter.ContextLedgerTokens = 600
44+
entries := largeContextLedgerEntries(24)
45+
46+
compact := budgeter.compactContextLedger(entries)
47+
if got := approxJSONTokens(compact); got > budgeter.ContextLedgerTokens {
48+
t.Fatalf("compact context ledger tokens = %d, want <= %d", got, budgeter.ContextLedgerTokens)
49+
}
50+
if len(compact) >= len(entries) {
51+
t.Fatalf("compact context ledger count = %d, want fewer than %d", len(compact), len(entries))
52+
}
53+
}
54+
55+
func BenchmarkCompactContextLedgerLarge(b *testing.B) {
56+
budgeter := DefaultReplanPromptBudgeter()
57+
budgeter.ContextLedgerTokens = 6000
58+
entries := largeContextLedgerEntries(240)
59+
60+
b.ReportAllocs()
61+
for i := 0; i < b.N; i++ {
62+
compact := budgeter.compactContextLedger(entries)
63+
if len(compact) == 0 {
64+
b.Fatal("empty compact ledger")
65+
}
66+
}
67+
}
68+
69+
func largeContextLedgerEntries(count int) []ContextLedgerEntry {
70+
entries := make([]ContextLedgerEntry, 0, count)
71+
for i := 0; i < count; i++ {
72+
entries = append(entries, ContextLedgerEntry{
73+
Kind: "task_action",
74+
Status: "rejected",
75+
Summary: fmt.Sprintf("context fact %03d: %s", i, strings.Repeat("large prompt detail ", 80)),
76+
Metadata: map[string]any{
77+
"turn": i,
78+
},
79+
})
80+
}
81+
return entries
82+
}

‎internal/orchestrator/service.go‎

Lines changed: 23 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -2588,19 +2588,27 @@ func (s *Service) taskIsTerminal(ctx context.Context, taskID string) (bool, erro
25882588
if err != nil {
25892589
return false, err
25902590
}
2591+
return taskIsTerminalFromSnapshot(snapshot, taskID), nil
2592+
}
2593+
2594+
func taskIsTerminalFromSnapshot(snapshot core.Snapshot, taskID string) bool {
25912595
for _, task := range snapshot.Tasks {
25922596
if task.ID == taskID {
2593-
return isTerminalTaskStatus(task.Status), nil
2597+
return isTerminalTaskStatus(task.Status)
25942598
}
25952599
}
2596-
return false, nil
2600+
return false
25972601
}
25982602

25992603
func (s *Service) taskArtifacts(ctx context.Context, taskID string) []core.TaskArtifact {
26002604
snapshot, err := s.store.Snapshot(ctx)
26012605
if err != nil {
26022606
return nil
26032607
}
2608+
return taskArtifactsFromSnapshot(snapshot, taskID)
2609+
}
2610+
2611+
func taskArtifactsFromSnapshot(snapshot core.Snapshot, taskID string) []core.TaskArtifact {
26042612
for _, task := range snapshot.Tasks {
26052613
if task.ID == taskID {
26062614
return append([]core.TaskArtifact{}, task.Artifacts...)
@@ -2614,6 +2622,10 @@ func (s *Service) taskPullRequestStates(ctx context.Context, taskID string) []Re
26142622
if err != nil {
26152623
return nil
26162624
}
2625+
return taskPullRequestStatesFromSnapshot(snapshot, taskID)
2626+
}
2627+
2628+
func taskPullRequestStatesFromSnapshot(snapshot core.Snapshot, taskID string) []ReplanPullRequestState {
26172629
states := []ReplanPullRequestState{}
26182630
for _, pr := range snapshot.PullRequests {
26192631
if pr.TaskID != taskID {
@@ -7030,10 +7042,12 @@ func (s *Service) replanLoopWithOptions(ctx context.Context, task core.Task, ini
70307042
limitUnproductiveTurns := !taskIsBroadObjective(task)
70317043
currentWorkPlan := initial.WorkPlan
70327044
for turn := 1; ; turn++ {
7033-
if terminal, err := s.taskIsTerminal(ctx, task.ID); err != nil {
7045+
stateSnapshot, err := s.store.Snapshot(ctx)
7046+
if err != nil {
70347047
_ = s.failTask(ctx, task.ID, err)
70357048
return false, "", "", results
7036-
} else if terminal && !options.FinalizationRecovery {
7049+
}
7050+
if taskIsTerminalFromSnapshot(stateSnapshot, task.ID) && !options.FinalizationRecovery {
70377051
return false, "", "", results
70387052
}
70397053
if limitUnproductiveTurns && stalledTurns >= maxConsecutiveUnproductiveReplanTurns {
@@ -7048,11 +7062,11 @@ func (s *Service) replanLoopWithOptions(ctx context.Context, task core.Task, ini
70487062
WorkPlan: currentWorkPlan,
70497063
Results: results,
70507064
ContextLedger: s.taskContextLedger(ctx, task.ID),
7051-
Artifacts: s.taskArtifacts(ctx, task.ID),
7052-
PullRequests: s.taskPullRequestStates(ctx, task.ID),
7053-
TaskSteering: s.taskSteering(ctx, task.ID),
7054-
PendingPullRequestFeedback: s.pendingPullRequestFeedback(ctx, task.ID),
7055-
PendingWorkerSteering: s.pendingWorkerSteering(ctx, task.ID),
7065+
Artifacts: taskArtifactsFromSnapshot(stateSnapshot, task.ID),
7066+
PullRequests: taskPullRequestStatesFromSnapshot(stateSnapshot, task.ID),
7067+
TaskSteering: taskSteering(stateSnapshot, task.ID),
7068+
PendingPullRequestFeedback: pendingPullRequestFeedbackFromSnapshot(stateSnapshot, task.ID),
7069+
PendingWorkerSteering: pendingWorkerSteering(stateSnapshot, task.ID),
70567070
Turn: turn,
70577071
BlockedFinalCandidateIDs: blockedFinalCandidateIDs,
70587072
RecoveryHint: recoveryHint,

0 commit comments

Comments
 (0)