diff --git a/internal/runtime/material_final_reserve.go b/internal/runtime/material_final_reserve.go new file mode 100644 index 0000000..6bf4e26 --- /dev/null +++ b/internal/runtime/material_final_reserve.go @@ -0,0 +1,31 @@ +package runtime + +// materialBudgetState describes only usage returned by the provider for a +// material subagent. A response with no reported usage cannot authorize a +// further request. +type materialBudgetState struct { + Limit, Used, LastPrompt, OutputCap, SavedCheckpoints, Responses int + UsageKnown bool +} + +// materialFinalDecision reserves room for an answer before allowing another +// checkpoint request. It is deliberately an estimate, not a billing limit for +// an already in-flight request. +func materialFinalDecision(s materialBudgetState) (finalOnly, stop bool) { + if s.Responses == 0 { + return false, false + } + if !s.UsageKnown { + return false, true + } + margin := s.LastPrompt / 4 + if margin < 1024 { + margin = 1024 + } + remaining := s.Limit - s.Used + if remaining <= s.LastPrompt+margin { + return false, true + } + reserve := s.LastPrompt + s.OutputCap + margin + return s.SavedCheckpoints >= 2 || remaining <= 2*reserve, false +} diff --git a/internal/runtime/material_final_reserve_test.go b/internal/runtime/material_final_reserve_test.go new file mode 100644 index 0000000..9e1315c --- /dev/null +++ b/internal/runtime/material_final_reserve_test.go @@ -0,0 +1,29 @@ +package runtime + +import "testing" + +func TestMaterialFinalReserveDecision(t *testing.T) { + tests := []struct { + name string + state materialBudgetState + wantFinal, wantStop bool + }{ + {"first request", materialBudgetState{Limit: 30000}, false, false}, + {"two saved checkpoints", materialBudgetState{Limit: 30000, Used: 16000, LastPrompt: 5000, OutputCap: 4000, SavedCheckpoints: 2, Responses: 2, UsageKnown: true}, true, false}, + {"one checkpoint and limited remainder", materialBudgetState{Limit: 30000, Used: 15000, LastPrompt: 5000, OutputCap: 4000, SavedCheckpoints: 1, Responses: 1, UsageKnown: true}, true, false}, + {"enough for another checkpoint", materialBudgetState{Limit: 30000, Used: 6000, LastPrompt: 3000, OutputCap: 4000, SavedCheckpoints: 1, Responses: 1, UsageKnown: true}, false, false}, + {"final prompt cannot fit", materialBudgetState{Limit: 30000, Used: 27000, LastPrompt: 8000, OutputCap: 4000, SavedCheckpoints: 1, Responses: 2, UsageKnown: true}, false, true}, + // Old A01's three material responses totalled 35,632 tokens; the third + // consumed 14,231. This offline check would stop before that third call. + {"old A01 before third request", materialBudgetState{Limit: 30000, Used: 35632 - 14231, LastPrompt: 10134, OutputCap: 4000, SavedCheckpoints: 1, Responses: 2, UsageKnown: true}, false, true}, + {"provider usage missing", materialBudgetState{Limit: 30000, Responses: 1}, false, true}, + } + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + final, stop := materialFinalDecision(tc.state) + if final != tc.wantFinal || stop != tc.wantStop { + t.Fatalf("final=%v stop=%v; want final=%v stop=%v", final, stop, tc.wantFinal, tc.wantStop) + } + }) + } +} diff --git a/internal/runtime/sub_loop.go b/internal/runtime/sub_loop.go index a82b097..9a86145 100644 --- a/internal/runtime/sub_loop.go +++ b/internal/runtime/sub_loop.go @@ -102,7 +102,40 @@ func (l *SubLLMRunner) Execute(rc RunCtx) error { } wrapped := func(ctx context.Context, model string, messages []agentllm.Message, _ []agentllm.ToolDecl, opts agentllm.StreamOptions) <-chan agentllm.Event { opts.MaxTokens = outputTokenCap(rc.Profile().Budget, rc.Settings(), "models.chat.maxOutputTokens") - in := stream(ctx, model, messages, decls, opts) + requestDecls := decls + if id == materialSubagentID && rc.Profile().Budget.MaxTokens != nil && *rc.Profile().Budget.MaxTokens > 0 { + if st.finalOnlyRejected { + return subagentBudgetError("checkpoint requested after final-only transition") + } + outputCap := opts.MaxTokens + if configured := int(rc.Settings().GetNumber("models.chat.maxOutputTokens", 0)); configured > 0 && configured < outputCap { + outputCap = configured + } + finalOnly, stop := materialFinalDecision(materialBudgetState{ + Limit: *rc.Profile().Budget.MaxTokens, Used: st.usageTotal.TotalTokens, + LastPrompt: st.usageTotal.LastPromptTokens, OutputCap: outputCap, + SavedCheckpoints: st.savedCheckpoints, Responses: st.responses, + UsageKnown: st.usageKnown, + }) + if stop { + return subagentBudgetError("final-answer reserve unavailable") + } + if finalOnly { + st.finalOnly = true + requestDecls = make([]agentllm.ToolDecl, 0, len(decls)) + for _, d := range decls { + if d.Name != workCheckpointTool { + requestDecls = append(requestDecls, d) + } + } + reminder := "现在用现有材料写出真实最终结论与未决项，不再发布检查点。" + if st.savedCheckpoints > 0 { + reminder = "阶段成果已保存；" + reminder + } + messages = append([]agentllm.Message{{Role: agentllm.RoleSystem, Content: reminder}}, messages...) + } + } + in := stream(ctx, model, messages, requestDecls, opts) out := make(chan agentllm.Event) go func() { defer close(out) @@ -216,13 +249,18 @@ type subSlot struct { } type subRun struct { - localCheckpoint bool - rc RunCtx - parts []core.Part - queued []agentllm.ToolCall - inflight *subSlot - turnText strings.Builder - usageTotal core.Usage + localCheckpoint bool + finalOnly bool + finalOnlyRejected bool + savedCheckpoints int + responses int + usageKnown bool + rc RunCtx + parts []core.Part + queued []agentllm.ToolCall + inflight *subSlot + turnText strings.Builder + usageTotal core.Usage } func (st *subRun) onEvent(ev nruntime.Event) { @@ -233,9 +271,19 @@ func (st *subRun) onEvent(ev nruntime.Event) { st.turnText.WriteString(ev.Text) _ = st.rc.Emit(core.EvAgentProgress, core.AgentProgressPayload{AgentRun: st.rc.Run().ID, Note: truncateRun(st.turnText.String(), 80)}) case "message_end": + st.responses++ + reportedTotal := ev.TotalTokens + if reportedTotal == 0 { + reportedTotal = ev.PromptTokens + ev.CompletionTokens + } + if st.responses == 1 { + st.usageKnown = reportedTotal > 0 + } else { + st.usageKnown = st.usageKnown && reportedTotal > 0 + } st.usageTotal.PromptTokens += ev.PromptTokens st.usageTotal.CompletionTokens += ev.CompletionTokens - st.usageTotal.TotalTokens += ev.TotalTokens + st.usageTotal.TotalTokens += reportedTotal st.usageTotal.LastPromptTokens = ev.PromptTokens st.usageTotal.LastCompletionTokens = ev.CompletionTokens st.usageTotal.CacheReadTokens += ev.CacheReadTokens @@ -274,7 +322,19 @@ func (st *subRun) onEvent(ev nruntime.Event) { func (st *subRun) runTool(name, args string) agenttools.Result { if name == workCheckpointTool && st.localCheckpoint { + if st.finalOnly { + st.finalOnlyRejected = true + res := agenttools.Result{Content: "最终答复阶段不再接受检查点", IsError: true} + if st.inflight != nil { + st.inflight.summary = res.Content + st.inflight.status = "error" + } + return res + } res := publishCheckpoint(st.rc, args) + if !res.IsError { + st.savedCheckpoints++ + } if st.inflight != nil { st.inflight.summary = res.Content st.inflight.status = "ok" @@ -309,6 +369,13 @@ func (st *subRun) runTool(name, args string) agenttools.Result { return agenttools.Result{Content: summary} } +func subagentBudgetError(reason string) <-chan agentllm.Event { + out := make(chan agentllm.Event, 1) + out <- agentllm.Event{Kind: agentllm.KindError, Err: fmt.Errorf("budget_exceeded: %s", reason), ErrorMessage: reason} + close(out) + return out +} + // delegateMaterialsNote 模型自派后台任务的材料指引（R2）：子代理看不到父回合 // 内联的附件全文，指明 knowledge_search 检索通道（上传时解析文本已进索引）。 func delegateMaterialsNote(docs []core.ParsedDocument) string { diff --git a/internal/runtime/usage_n3_test.go b/internal/runtime/usage_n3_test.go index 686abed..650250d 100644 --- a/internal/runtime/usage_n3_test.go +++ b/internal/runtime/usage_n3_test.go @@ -69,3 +69,19 @@ func TestSubRunAccountsTokens(t *testing.T) { t.Fatalf("子代理用量未入账：%+v", u) } } + +func TestSubRunAccountsPromptAndCompletionWhenTotalIsMissing(t *testing.T) { + rc := &fakeRunCtx{ctx: context.Background(), prof: core.AgentProfile{Model: core.ModelRoute{Chat: "test-model"}}} + stream := func(ctx context.Context, model string, msgs []agentllm.Message, _ []agentllm.ToolDecl, _ agentllm.StreamOptions) <-chan agentllm.Event { + ch := make(chan agentllm.Event, 1) + ch <- agentllm.Event{Kind: agentllm.KindMessageEnd, Message: agentllm.Message{Role: agentllm.RoleAssistant, Content: "子结果", StopReason: agentllm.StopEnd}, Usage: agentllm.Usage{PromptTokens: 66, CompletionTokens: 6}} + close(ch) + return ch + } + if err := (&SubLLMRunner{Stream: stream}).Execute(rc); err != nil { + t.Fatal(err) + } + if rc.usageTotal.TotalTokens != 72 { + t.Fatalf("sub-run lost reported prompt/completion usage: %+v", rc.usageTotal) + } +} diff --git a/internal/runtime/work_checkpoint_test.go b/internal/runtime/work_checkpoint_test.go index 8488346..06bf014 100644 --- a/internal/runtime/work_checkpoint_test.go +++ b/internal/runtime/work_checkpoint_test.go @@ -4,14 +4,173 @@ import ( agentllm "agent-native/runtime/llm" core "agent-system-core" "context" + "strings" "testing" ) +func TestMaterialFinalReserveAfterTwoCheckpoints(t *testing.T) { + limit := 100000 + prof := subProfile() + prof.Budget.MaxTokens = &limit + rc := &cappedCheckpointRC{&checkpointRC{subRC: subRC{fakeRunCtx: fakeRunCtx{ctx: context.Background(), prof: prof}, id: materialSubagentID}}} + turn := func(body string, call ...agentllm.ToolCall) []agentllm.Event { + events := textDeltaTurn(body, call...) + events[len(events)-1].Usage = agentllm.Usage{PromptTokens: 4000, CompletionTokens: 1000, TotalTokens: 5000} + return events + } + base := scriptStream([][]agentllm.Event{ + turn("", toolCall(workCheckpointTool, checkpointJSON())), + turn("", toolCall(workCheckpointTool, strings.Replace(checkpointJSON(), `"cp1"`, `"cp2"`, 1))), + turn("整理完成"), + }, nil) + calls, finalOnly := 0, false + r := &SubLLMRunner{Stream: func(ctx context.Context, model string, msgs []agentllm.Message, decls []agentllm.ToolDecl, opts agentllm.StreamOptions) <-chan agentllm.Event { + calls++ + if calls == 3 { + finalOnly = true + for _, d := range decls { + if d.Name == workCheckpointTool { + finalOnly = false + } + } + reminder := false + for _, msg := range msgs { + reminder = reminder || (msg.Role == agentllm.RoleSystem && strings.Contains(msg.Content, "不再发布检查点")) + } + finalOnly = finalOnly && reminder + } + return base(ctx, model, msgs, decls, opts) + }} + if err := r.Execute(rc); err != nil { + t.Fatal(err) + } + if calls != 3 || len(rc.checkpoints) != 2 || !finalOnly { + t.Fatalf("final-only transition missing: calls=%d checkpoints=%d finalOnly=%v", calls, len(rc.checkpoints), finalOnly) + } +} + +func TestMaterialFinalReserveRejectsHiddenCheckpointWithoutAnotherProviderCall(t *testing.T) { + limit := 100000 + prof := subProfile() + prof.Budget.MaxTokens = &limit + rc := &cappedCheckpointRC{&checkpointRC{subRC: subRC{fakeRunCtx: fakeRunCtx{ctx: context.Background(), prof: prof}, id: materialSubagentID}}} + turn := func(call agentllm.ToolCall) []agentllm.Event { + events := textDeltaTurn("", call) + events[len(events)-1].Usage = agentllm.Usage{PromptTokens: 4000, CompletionTokens: 1000, TotalTokens: 5000} + return events + } + base := scriptStream([][]agentllm.Event{ + turn(toolCall(workCheckpointTool, checkpointJSON())), + turn(toolCall(workCheckpointTool, strings.Replace(checkpointJSON(), `"cp1"`, `"cp2"`, 1))), + turn(toolCall(workCheckpointTool, strings.Replace(checkpointJSON(), `"cp1"`, `"cp3"`, 1))), + textDeltaTurn("unexpected extra response"), + }, nil) + providerCalls := 0 + r := &SubLLMRunner{Stream: func(ctx context.Context, model string, msgs []agentllm.Message, decls []agentllm.ToolDecl, opts agentllm.StreamOptions) <-chan agentllm.Event { + providerCalls++ + return base(ctx, model, msgs, decls, opts) + }} + _ = r.Execute(rc) + if providerCalls != 3 || len(rc.checkpoints) != 2 { + t.Fatalf("hidden checkpoint caused a publication or another provider call: provider=%d saved=%d", providerCalls, len(rc.checkpoints)) + } +} + +func TestMaterialFinalReserveAnswersAfterOneCheckpointWhenBudgetNarrows(t *testing.T) { + limit := 30000 + prof := subProfile() + prof.Budget.MaxTokens = &limit + rc := &cappedCheckpointRC{&checkpointRC{subRC: subRC{fakeRunCtx: fakeRunCtx{ctx: context.Background(), prof: prof}, id: materialSubagentID}}} + first := textDeltaTurn("", toolCall(workCheckpointTool, checkpointJSON())) + first[len(first)-1].Usage = agentllm.Usage{PromptTokens: 8000, CompletionTokens: 2000, TotalTokens: 10000} + base := scriptStream([][]agentllm.Event{first, textDeltaTurn("目前只核对出一项，其他内容尚待核对")}, nil) + providerCalls, finalOnly := 0, false + r := &SubLLMRunner{Stream: func(ctx context.Context, model string, msgs []agentllm.Message, decls []agentllm.ToolDecl, opts agentllm.StreamOptions) <-chan agentllm.Event { + providerCalls++ + if providerCalls == 2 { + finalOnly = true + for _, d := range decls { + finalOnly = finalOnly && d.Name != workCheckpointTool + } + } + return base(ctx, model, msgs, decls, opts) + }} + if err := r.Execute(rc); err != nil { + t.Fatal(err) + } + if providerCalls != 2 || len(rc.checkpoints) != 1 || !finalOnly { + t.Fatalf("partial final answer was not requested: calls=%d saved=%d finalOnly=%v", providerCalls, len(rc.checkpoints), finalOnly) + } +} + +func TestMaterialFinalReserveStopsBeforeUnaffordableRequest(t *testing.T) { + for _, tc := range []struct { + name string + usage agentllm.Usage + }{ + {"reported usage exceeds final prompt allowance", agentllm.Usage{PromptTokens: 25000, CompletionTokens: 2000, TotalTokens: 27000}}, + {"missing usage", agentllm.Usage{}}, + } { + t.Run(tc.name, func(t *testing.T) { + limit := 30000 + prof := subProfile() + prof.Budget.MaxTokens = &limit + rc := &cappedCheckpointRC{&checkpointRC{subRC: subRC{fakeRunCtx: fakeRunCtx{ctx: context.Background(), prof: prof}, id: materialSubagentID}}} + first := textDeltaTurn("", toolCall(workCheckpointTool, checkpointJSON())) + first[len(first)-1].Usage = tc.usage + base := scriptStream([][]agentllm.Event{first, textDeltaTurn("unexpected extra response")}, nil) + providerCalls := 0 + r := &SubLLMRunner{Stream: func(ctx context.Context, model string, msgs []agentllm.Message, decls []agentllm.ToolDecl, opts agentllm.StreamOptions) <-chan agentllm.Event { + providerCalls++ + return base(ctx, model, msgs, decls, opts) + }} + if err := r.Execute(rc); err == nil { + t.Fatal("expected honest budget/usage failure") + } + if providerCalls != 1 || len(rc.checkpoints) != 1 { + t.Fatalf("unaffordable request made or saved checkpoint lost: calls=%d saved=%d", providerCalls, len(rc.checkpoints)) + } + }) + } +} + +func TestMaterialFinalReserveDoesNotChangeOrdinarySubagents(t *testing.T) { + limit := 10000 + prof := subProfile() + prof.Budget.MaxTokens = &limit + rc := &subRC{fakeRunCtx: fakeRunCtx{ctx: context.Background(), prof: prof}, id: "agent.sub.explore", tools: subTools()} + first := textDeltaTurn("", toolCall("knowledge_search", `{"query":"x"}`)) + first[len(first)-1].Usage = agentllm.Usage{PromptTokens: 5000, CompletionTokens: 1000, TotalTokens: 6000} + base := scriptStream([][]agentllm.Event{first, textDeltaTurn("普通子代理已完成")}, nil) + providerCalls, reminder := 0, false + r := &SubLLMRunner{Stream: func(ctx context.Context, model string, msgs []agentllm.Message, decls []agentllm.ToolDecl, opts agentllm.StreamOptions) <-chan agentllm.Event { + providerCalls++ + for _, msg := range msgs { + reminder = reminder || strings.Contains(msg.Content, "不再发布检查点") + } + return base(ctx, model, msgs, decls, opts) + }} + if err := r.Execute(rc); err != nil { + t.Fatal(err) + } + if providerCalls != 2 || reminder { + t.Fatalf("ordinary subagent changed: provider calls=%d material reminder=%v", providerCalls, reminder) + } +} + type checkpointRC struct { subRC checkpoints []core.WorkCheckpointCandidate } +type cappedCheckpointRC struct{ *checkpointRC } + +func (c *cappedCheckpointRC) Settings() core.EffectiveSettings { + eff := c.checkpointRC.Settings() + eff["models"].(map[string]any)["chat"].(map[string]any)["maxOutputTokens"] = float64(4000) + return eff +} + func (c *checkpointRC) PublishWorkCheckpoint(v core.WorkCheckpointCandidate) error { c.checkpoints = append(c.checkpoints, v) return nil