fix(aiops): stop model calls when token budget is exceeded

internal/manager/biz/aiops/graph/budget_stop_model.go

@@ -9,6 +9,8 @@ import (
 
 	einomodel "github.com/cloudwego/eino/components/model"
 	"github.com/cloudwego/eino/schema"
+
+	"github.com/ongridio/ongrid/internal/pkg/llm"
 )
 
 type budgetStopModel struct {
@@ -23,6 +25,10 @@ func wrapBudgetStopModel(inner einomodel.ToolCallingChatModel) einomodel.ToolCal
 }
 
 func (m *budgetStopModel) Generate(ctx context.Context, input []*schema.Message, opts ...einomodel.Option) (*schema.Message, error) {
+	// 回调只能标记预算拒绝；在普通推理和工具熔断总结之前终止模型调用。
+	if err := llm.BudgetRejectionFromContext(ctx); err != nil {
+		return nil, err
+	}
 	if env, ok := latestTerminalToolBudget(input); ok {
 		return m.synthesizeAfterToolBudget(ctx, input, env.Tool, opts...)
 	}
@@ -38,6 +44,9 @@ func (m *budgetStopModel) Generate(ctx context.Context, input []*schema.Message,
 }
 
 func (m *budgetStopModel) Stream(ctx context.Context, input []*schema.Message, opts ...einomodel.Option) (*schema.StreamReader[*schema.Message], error) {
+	if err := llm.BudgetRejectionFromContext(ctx); err != nil {
+		return nil, err
+	}
 	if env, ok := latestTerminalToolBudget(input); ok {
 		msg, err := m.synthesizeAfterToolBudget(ctx, input, env.Tool, opts...)
 		if err != nil {

internal/manager/biz/aiops/graph/token_budget_test.go

@@ -0,0 +1,86 @@
+package graph
+
+import (
+	"context"
+	"errors"
+	"io"
+	"testing"
+
+	"github.com/cloudwego/eino/compose"
+	"github.com/cloudwego/eino/schema"
+
+	"github.com/ongridio/ongrid/internal/pkg/llm"
+)
+
+func TestBuildReActGraph_TokenBudgetStopsModelCalls(t *testing.T) {
+	t.Parallel()
+	for _, mode := range []string{"invoke", "stream"} {
+		for _, tc := range []struct {
+			name       string
+			limit      int
+			toolBudget bool
+			wantReject bool
+		}{
+			{name: "exhausted", limit: 1, wantReject: true},
+			{name: "exhausted_before_synthesis", limit: 1, toolBudget: true, wantReject: true},
+			{name: "available", limit: 100000},
+			{name: "available_for_synthesis", limit: 100000, toolBudget: true},
+			{name: "unlimited", limit: 0},
+		} {
+			t.Run(mode+"/"+tc.name, func(t *testing.T) {
+				t.Parallel()
+				inner := newScriptedChatModel(makeAssistantNoTools("diagnosis complete"))
+				g, err := BuildReActGraph(inner, nil, Config{})
+				if err != nil {
+					t.Fatal(err)
+				}
+				handler := llm.NewBudgetCallbackHandler(llm.NewInMemoryBudget(tc.limit), 0)
+				input := &Input{UserText: "Check the service health and summarize the evidence."}
+				if tc.toolBudget {
+					input = &Input{History: []*schema.Message{
+						schema.UserMessage(input.UserText),
+						makeAssistantToolCall("", "call_1", "query_logql", `{}`),
+						schema.ToolMessage(`{"status":"call_budget_exceeded","tool":"query_logql","final_answer_required":true}`, "call_1", schema.WithToolName("query_logql")),
+					}}
+				}
+				err = runBudgetGraph(context.Background(), g, input, mode, compose.WithCallbacks(handler))
+				wantCalls := int32(1)
+				if tc.wantReject {
+					wantCalls = 0
+					if !errors.Is(err, llm.ErrBudgetExceeded) {
+						t.Errorf("error = %v, want ErrBudgetExceeded", err)
+					}
+					if handler.Stats().Rejects != 1 {
+						t.Errorf("budget stats = %+v, want one rejection", handler.Stats())
+					}
+				} else if err != nil {
+					t.Errorf("unexpected error: %v", err)
+				}
+				if calls := inner.generateCalls(); calls != wantCalls {
+					t.Errorf("underlying model calls = %d, want %d", calls, wantCalls)
+				}
+			})
+		}
+	}
+}
+
+func runBudgetGraph(ctx context.Context, g compose.Runnable[*Input, *Output], input *Input, mode string, opts ...compose.Option) error {
+	if mode == "invoke" {
+		_, err := g.Invoke(ctx, input, opts...)
+		return err
+	}
+	stream, err := g.Stream(ctx, input, opts...)
+	if err != nil {
+		return err
+	}
+	defer stream.Close()
+	for {
+		_, err := stream.Recv()
+		if errors.Is(err, io.EOF) {
+			return nil
+		}
+		if err != nil {
+			return err
+		}
+	}
+}

internal/pkg/llm/budget_callback.go

@@ -7,8 +7,8 @@
 // Design intent:
 //   - OnStart estimates prompt tokens (cheap rule-of-thumb: joined content
 //     length / 4) and asks BudgetChecker.Check. On rejection we mark the
-//     context so OnEnd / OnError can short-circuit reporting and PR-2 graph
-//     code can surface ErrBudgetExceeded.
+//     context so the graph's ChatModel wrapper can return ErrBudgetExceeded
+//     before invoking the underlying model.
 //   - OnEnd reads schema.ResponseMeta.Usage from the model's reply and
 //     records it via BudgetChecker.Record.
 //   - We do NOT touch the legacy budget gate inside openaiClient.Chat — it
@@ -19,9 +19,8 @@
 //
 // Why not return an error from OnStart? eino's callback contract gives
 // handlers no way to short-circuit a component. We attach the rejection
-// to ctx so the graph node (PR-2) can read it back out and fail-fast
-// before the network call. PR-1 ships the plumbing; PR-2 wires the check
-// into the node.
+// to ctx so the graph's ChatModel wrapper can read it back and fail before
+// either a normal model call or a tool-budget synthesis call.
 package llm
 
 import (
@@ -39,9 +38,9 @@ import (
 type budgetRejectKey struct{}
 
 // BudgetRejectionFromContext returns the budget-rejection error attached
-// to ctx by BudgetCallbackHandler.OnStart, or nil if none. Graph nodes
-// (PR-2) call this immediately after invoking the model to convert the
-// soft signal into a hard error.
+// to ctx by BudgetCallbackHandler.OnStart, or nil if none. The graph's
+// ChatModel wrapper checks it before invoking the model, including Stream,
+// to convert the callback signal into a hard error.
 func BudgetRejectionFromContext(ctx context.Context) error {
 	v := ctx.Value(budgetRejectKey{})
 	if v == nil {