WeKnora deployment source with the reviewed wiki identity, summary and recovery fixes

internal/workbench/metering.go

@@ -0,0 +1,341 @@
+package workbench
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"strings"
+	"time"
+
+	"github.com/Tencent/WeKnora/internal/agent/token"
+	"github.com/Tencent/WeKnora/internal/models/chat"
+	"github.com/Tencent/WeKnora/internal/types"
+	"github.com/google/uuid"
+	"gorm.io/gorm"
+	"gorm.io/gorm/clause"
+)
+
+// This rate preserves the imported ledger. It is not exposed in member APIs.
+const pointsPerToken int64 = 10
+
+type meteringContextKey struct{}
+type meteringAttribution struct{ AgentID, JobID, Action string }
+
+func (s *Service) WrapChat(model chat.Chat) chat.Chat { return &meteredChat{model: model, service: s} }
+
+type meteredChat struct {
+	model   chat.Chat
+	service *Service
+}
+
+func (m *meteredChat) GetModelID() string   { return m.model.GetModelID() }
+func (m *meteredChat) GetModelName() string { return m.model.GetModelName() }
+
+func (s *Service) reserveBalance(ctx context.Context, amount int64) (*BalanceHold, error) {
+	t, ok := types.TenantIDFromContext(ctx)
+	if !ok || t == 0 {
+		return nil, ErrForbidden
+	}
+	u, _ := types.UserIDFromContext(ctx)
+	// Scheduled maintenance is attributed to the workspace, never a random member.
+	if u == "" || types.IsSyntheticUserID(u) {
+		return nil, nil
+	}
+	if amount <= 0 || amount > 1_000_000_000_000 {
+		return nil, ErrInvalid
+	}
+	hold := &BalanceHold{ID: uuid.NewString(), TenantID: t, UserID: u, Amount: amount, Status: "reserved"}
+	err := s.db.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
+		q := tx.Model(&Balance{}).Where("tenant_id=? AND user_id=? AND balance-reserved>=?", t, u, amount).Updates(map[string]any{"reserved": gorm.Expr("reserved+?", amount), "updated_at": time.Now()})
+		if q.Error != nil {
+			return q.Error
+		}
+		if q.RowsAffected != 1 {
+			return fmt.Errorf("%w: 余额不足以完成本次生成，请联系管理员补充余额或调低输出长度", ErrInvalid)
+		}
+		hold.CreatedAt = time.Now().UTC()
+		hold.ExpiresAt = hold.CreatedAt.Add(30 * time.Minute)
+		if e := checkQuotaReservation(tx, hold); e != nil {
+			return e
+		}
+		return tx.Create(hold).Error
+	})
+	return hold, err
+}
+
+func (s *Service) settleBalance(ctx context.Context, hold *BalanceHold, usage Usage) error {
+	if usage.CreatedAt.IsZero() {
+		usage.CreatedAt = time.Now().UTC()
+	} else {
+		usage.CreatedAt = usage.CreatedAt.UTC()
+	}
+	return s.db.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
+		if hold != nil {
+			var current BalanceHold
+			if e := tx.Clauses(clause.Locking{Strength: "UPDATE"}).Where("id=? AND tenant_id=? AND user_id=?", hold.ID, hold.TenantID, hold.UserID).Take(&current).Error; e != nil {
+				return e
+			}
+			if current.Status == "settled" {
+				return nil
+			}
+			if current.Status != "reserved" {
+				return ErrConflict
+			}
+			var account Balance
+			if e := tx.Clauses(clause.Locking{Strength: "UPDATE"}).Where("tenant_id=? AND user_id=?", hold.TenantID, hold.UserID).Take(&account).Error; e != nil {
+				return e
+			}
+			if account.Reserved < current.Amount {
+				return errors.New("余额预留记录不一致")
+			}
+			// Actual usage above the conservative reservation is retained as debt,
+			// never silently dropped or charged to another user's account.
+			if e := tx.Model(&account).Updates(map[string]any{"balance": account.Balance - usage.Points, "reserved": account.Reserved - current.Amount, "updated_at": time.Now()}).Error; e != nil {
+				return e
+			}
+			if usage.Points > 0 {
+				row := Transaction{ID: uuid.NewString(), TenantID: hold.TenantID, UserID: hold.UserID, Delta: -usage.Points, Kind: "debit", Reason: usage.Action, IdempotencyKey: "call:" + hold.ID, Metadata: JSON{"call_id": usage.CallID, "job_id": usage.JobID, "model_id": usage.ModelID}, CreatedAt: time.Now()}
+				if e := tx.Create(&row).Error; e != nil {
+					return e
+				}
+			}
+			if e := tx.Model(&current).Update("status", "settled").Error; e != nil {
+				return e
+			}
+		}
+		return tx.Clauses(clause.OnConflict{DoNothing: true}).Create(&usage).Error
+	})
+}
+
+func (s *Service) releaseExpiredBalances(ctx context.Context) error {
+	var holds []BalanceHold
+	if e := s.db.WithContext(ctx).Where("status=? AND expires_at<=?", "reserved", time.Now().UTC()).Limit(100).Find(&holds).Error; e != nil {
+		return e
+	}
+	for _, hold := range holds {
+		if e := s.db.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
+			var current BalanceHold
+			if e := tx.Clauses(clause.Locking{Strength: "UPDATE"}).Where("id=?", hold.ID).Take(&current).Error; e != nil {
+				return e
+			}
+			if current.Status != "reserved" || current.ExpiresAt.After(time.Now()) {
+				return nil
+			}
+			q := tx.Model(&Balance{}).Where("tenant_id=? AND user_id=? AND reserved>=?", current.TenantID, current.UserID, current.Amount).Updates(map[string]any{"reserved": gorm.Expr("reserved-?", current.Amount), "updated_at": time.Now()})
+			if q.Error != nil {
+				return q.Error
+			}
+			if q.RowsAffected != 1 {
+				return ErrConflict
+			}
+			return tx.Model(&current).Update("status", "expired").Error
+		}); e != nil {
+			return e
+		}
+	}
+	return nil
+}
+
+type meteredCall struct {
+	hold      *BalanceHold
+	estimator *token.Estimator
+	input     int
+	usage     Usage
+	started   time.Time
+	budgetMS  int64
+	purpose   string
+}
+
+func (m *meteredChat) start(ctx context.Context, messages []chat.Message, options *chat.ChatOptions) (*meteredCall, error) {
+	estimator, e := token.NewEstimator()
+	if e != nil {
+		return nil, e
+	}
+	input := estimator.EstimateMessages(messages)
+	output := 8192
+	if options != nil {
+		input += estimator.EstimateTools(options.Tools)
+		if options.MaxCompletionTokens > 0 {
+			output = options.MaxCompletionTokens
+		} else if options.MaxTokens > 0 {
+			output = options.MaxTokens
+		}
+	}
+	if input > 10_000_000 || output > 10_000_000 {
+		return nil, ErrInvalid
+	}
+	hold, e := m.service.reserveBalance(ctx, int64(input+input/5+output)*pointsPerToken)
+	if e != nil {
+		return nil, e
+	}
+	t, _ := types.TenantIDFromContext(ctx)
+	u, _ := types.UserIDFromContext(ctx)
+	callID := uuid.NewString()
+	started := time.Now().UTC()
+	if hold != nil {
+		callID = hold.ID
+		started = hold.CreatedAt
+	}
+	a, _ := ctx.Value(meteringContextKey{}).(meteringAttribution)
+	purpose, _ := types.LLMCallMetadataFromContext(ctx)
+	if a.Action == "" {
+		a.Action = actionName(purpose)
+	}
+	if a.Action == "" {
+		a.Action = "AI 生成"
+	}
+	call := &meteredCall{hold: hold, estimator: estimator, input: input, started: time.Now(), purpose: purpose, usage: Usage{ID: uuid.NewString(), CallID: callID, TenantID: t, OwnerID: u, JobID: a.JobID, AgentID: a.AgentID, ModelID: m.GetModelID(), Action: a.Action, Status: "succeeded", Metadata: JSON{}, CreatedAt: started}}
+	if deadline, ok := ctx.Deadline(); ok {
+		call.budgetMS = max(0, time.Until(deadline).Milliseconds())
+	}
+	return call, nil
+}
+
+func (m *meteredChat) finish(ctx context.Context, call *meteredCall, reported types.TokenUsage, output string, failure error) error {
+	call.usage.InputTokens = int64(max(0, reported.PromptTokens))
+	call.usage.OutputTokens = int64(max(0, reported.CompletionTokens))
+	measured := reported.TotalTokens > 0 || reported.PromptTokens > 0 || reported.CompletionTokens > 0
+	if !measured && (failure == nil || output != "") {
+		call.usage.InputTokens = int64(call.input)
+		call.usage.OutputTokens = int64(call.estimator.EstimateString(output))
+	}
+	if measured && int64(reported.TotalTokens) > call.usage.InputTokens+call.usage.OutputTokens {
+		call.usage.OutputTokens = int64(reported.TotalTokens) - call.usage.InputTokens
+	}
+	call.usage.Metadata = JSON{"estimated": !measured, "model_name": m.GetModelName(), "elapsed_ms": time.Since(call.started).Milliseconds(), "timeout_budget_ms": call.budgetMS, "purpose": call.purpose}
+	failure = phaseFailure(failure, "model", call.started)
+	if failure != nil {
+		call.usage.Status = "failed"
+		category, message := jobErrorDetails(failure)
+		call.usage.Metadata["error_class"], call.usage.Metadata["error_message"] = category, message
+	}
+	if call.hold != nil {
+		call.usage.Points = (call.usage.InputTokens + call.usage.OutputTokens) * pointsPerToken
+	}
+	finish, cancel := context.WithTimeout(context.WithoutCancel(ctx), 10*time.Second)
+	defer cancel()
+	if e := m.service.settleBalance(finish, call.hold, call.usage); e != nil {
+		return fmt.Errorf("模型用量结算失败，请联系管理员核对调用 %s: %w", call.usage.CallID, e)
+	}
+	if failure != nil {
+		_, message := jobErrorDetails(failure)
+		diagnostic := jobErrorMetadata(failure)
+		diagnostic["model_id"], diagnostic["agent_id"], diagnostic["purpose"] = m.GetModelID(), call.usage.AgentID, call.purpose
+		diagnostic["elapsed_ms"], diagnostic["timeout_budget_ms"] = call.usage.Metadata["elapsed_ms"], call.budgetMS
+		if e := m.service.recordOperation(finish, OperationLog{Kind: "system", Level: "error", Action: call.usage.Action, Message: message, TraceID: call.usage.CallID, ResourceID: call.usage.JobID, Metadata: diagnostic}); e != nil {
+			logJobFailure(finish, nil, "model_error_log", e)
+		}
+	}
+	return nil
+}
+
+func (m *meteredChat) Chat(ctx context.Context, messages []chat.Message, options *chat.ChatOptions) (*types.ChatResponse, error) {
+	ctx, cancel := context.WithTimeout(ctx, 25*time.Minute)
+	defer cancel()
+	call, e := m.start(ctx, messages, options)
+	if e != nil {
+		return nil, e
+	}
+	response, err := m.model.Chat(ctx, messages, options)
+	err = phaseFailure(err, "model", call.started)
+	usage := types.TokenUsage{}
+	output := ""
+	if response != nil {
+		usage = response.Usage
+		output = response.Content + response.ReasoningContent
+		if len(response.ToolCalls) > 0 {
+			output += compact(response.ToolCalls)
+		}
+	}
+	if settleErr := m.finish(ctx, call, usage, output, err); settleErr != nil {
+		return nil, settleErr
+	}
+	return response, err
+}
+
+func (m *meteredChat) ChatStream(ctx context.Context, messages []chat.Message, options *chat.ChatOptions) (<-chan types.StreamResponse, error) {
+	ctx, cancel := context.WithTimeout(ctx, 25*time.Minute)
+	call, e := m.start(ctx, messages, options)
+	if e != nil {
+		cancel()
+		return nil, e
+	}
+	upstream, err := m.model.ChatStream(ctx, messages, options)
+	err = phaseFailure(err, "model", call.started)
+	if upstream == nil && err == nil {
+		err = errors.New("AI 返回了空的流式响应")
+	}
+	if err != nil {
+		defer cancel()
+		if e = m.finish(ctx, call, types.TokenUsage{}, "", err); e != nil {
+			return nil, e
+		}
+		return nil, err
+	}
+	out := make(chan types.StreamResponse)
+	go func() {
+		defer close(out)
+		defer cancel()
+		var usage types.TokenUsage
+		var output strings.Builder
+		var failure error
+		var final *types.StreamResponse
+		active := true
+	stream:
+		for {
+			var chunk types.StreamResponse
+			select {
+			case <-ctx.Done():
+				failure = ctx.Err()
+				active = false
+				break stream
+			case next, ok := <-upstream:
+				if !ok {
+					break stream
+				}
+				chunk = next
+			}
+			if chunk.Usage != nil {
+				usage = *chunk.Usage
+			}
+			output.WriteString(chunk.Content)
+			if chunk.FinishReason == types.FinishReasonIncomplete || chunk.ResponseType == types.ResponseTypeError {
+				failure = errors.New("AI 响应中断")
+			}
+			if chunk.Done {
+				copy := chunk
+				final = &copy
+				continue
+			}
+			if active {
+				select {
+				case out <- chunk:
+				case <-ctx.Done():
+					active = false
+					failure = ctx.Err()
+				}
+			}
+		}
+		if final == nil && failure == nil {
+			failure = errors.New("AI 未返回完成标记")
+		}
+		failure = phaseFailure(failure, "model", call.started)
+		if e := m.finish(ctx, call, usage, output.String(), failure); e != nil {
+			failure = e
+		}
+		if final == nil {
+			final = &types.StreamResponse{Done: true}
+		}
+		if failure != nil {
+			final.FinishReason = types.FinishReasonIncomplete
+			final.Data = map[string]any{"error": publicJobError(failure)}
+		}
+		if active {
+			select {
+			case out <- *final:
+			case <-ctx.Done():
+			}
+		}
+	}()
+	return out, nil
+}