xinmi/internal/workbench/metering_test.go

package workbench

import (
	"context"
	"errors"
	"testing"
	"time"

	"github.com/Tencent/WeKnora/internal/models/chat"
	"github.com/Tencent/WeKnora/internal/types"
)

type billingChatStub struct {
	response *types.ChatResponse
	stream   <-chan types.StreamResponse
	err      error
}

func (m billingChatStub) GetModelID() string   { return "model" }
func (m billingChatStub) GetModelName() string { return "test" }
func (m billingChatStub) Chat(context.Context, []chat.Message, *chat.ChatOptions) (*types.ChatResponse, error) {
	return m.response, m.err
}
func (m billingChatStub) ChatStream(context.Context, []chat.Message, *chat.ChatOptions) (<-chan types.StreamResponse, error) {
	return m.stream, m.err
}

func seedBalance(t *testing.T, s *Service, tenant uint64, user string, amount int64) {
	t.Helper()
	if e := s.db.Create(&Balance{TenantID: tenant, UserID: user, Balance: amount}).Error; e != nil {
		t.Fatal(e)
	}
}
func accountBalance(t *testing.T, s *Service, tenant uint64, user string) Balance {
	t.Helper()
	var row Balance
	if e := s.db.Where("tenant_id=? AND user_id=?", tenant, user).Take(&row).Error; e != nil {
		t.Fatal(e)
	}
	return row
}
func TestBalanceReservationSettlementAndTenantIsolation(t *testing.T) {
	s := sqliteWorkbench(t)
	ctx := workbenchContext(1, "alice", types.TenantRoleContributor)
	seedBalance(t, s, 1, "alice", 1000)
	seedBalance(t, s, 2, "alice", 9000)
	hold, e := s.reserveBalance(ctx, 600)
	if e != nil {
		t.Fatal(e)
	}
	if _, e = s.reserveBalance(ctx, 500); e == nil {
		t.Fatal("overbooking accepted")
	}
	u := Usage{ID: "usage", TenantID: 1, OwnerID: "alice", CallID: hold.ID, Points: 200, Metadata: JSON{}}
	if e = s.settleBalance(ctx, hold, u); e != nil {
		t.Fatal(e)
	}
	if e = s.settleBalance(ctx, hold, u); e != nil {
		t.Fatal(e)
	}
	got := accountBalance(t, s, 1, "alice")
	if got.Balance != 800 || got.Reserved != 0 {
		t.Fatalf("duplicate debit or hold leak: %+v", got)
	}
	if got = accountBalance(t, s, 2, "alice"); got.Balance != 9000 || got.Reserved != 0 {
		t.Fatal("cross-tenant debit")
	}
	var count int64
	s.db.Model(&Transaction{}).Count(&count)
	if count != 1 {
		t.Fatalf("ledger entries=%d", count)
	}
}
func TestExpiredReservationIsReleasedOnlyOnce(t *testing.T) {
	s := sqliteWorkbench(t)
	ctx := workbenchContext(1, "alice", types.TenantRoleContributor)
	seedBalance(t, s, 1, "alice", 1000)
	h, e := s.reserveBalance(ctx, 600)
	if e != nil {
		t.Fatal(e)
	}
	s.db.Model(h).Update("expires_at", time.Now().UTC().Add(-time.Second))
	for i := 0; i < 2; i++ {
		if e = s.releaseExpiredBalances(ctx); e != nil {
			t.Fatal(e)
		}
	}
	row := accountBalance(t, s, 1, "alice")
	if row.Balance != 1000 || row.Reserved != 0 {
		t.Fatalf("bad release: %+v", row)
	}
	if e = s.settleBalance(ctx, h, Usage{}); !errors.Is(e, ErrConflict) {
		t.Fatalf("expired hold settled: %v", e)
	}
}
func TestMeteredChatUsesProviderUsageAndReleasesFailure(t *testing.T) {
	for _, failed := range []bool{false, true} {
		t.Run(map[bool]string{false: "success", true: "failure"}[failed], func(t *testing.T) {
			s := sqliteWorkbench(t)
			ctx := workbenchContext(1, "alice", types.TenantRoleContributor)
			seedBalance(t, s, 1, "alice", 1000000)
			model := billingChatStub{response: &types.ChatResponse{Content: "response", Usage: types.TokenUsage{PromptTokens: 12, CompletionTokens: 7, TotalTokens: 19}}}
			if failed {
				model.response = nil
				model.err = errors.New("provider rejected")
			}
			_, err := s.WrapChat(model).Chat(ctx, []chat.Message{{Role: "user", Content: "hello"}}, &chat.ChatOptions{MaxTokens: 20})
			if (err != nil) != failed {
				t.Fatalf("unexpected failure: %v", err)
			}
			want := int64(1000000)
			if !failed {
				want -= 190
			}
			got := accountBalance(t, s, 1, "alice")
			if got.Balance != want || got.Reserved != 0 {
				t.Fatalf("bad settlement: %+v", got)
			}
			var rows []Usage
			s.db.Find(&rows)
			if len(rows) != 1 {
				t.Fatal("usage missing")
			}
		})
	}
}

func TestMeteredTimeoutPreservesAttributionAndReleasesBalance(t *testing.T) {
	s := sqliteWorkbench(t)
	ctx := types.WithLLMCallMetadata(workbenchContext(1, "alice", types.TenantRoleContributor), "digital_human_distillation", "private-fingerprint")
	seedBalance(t, s, 1, "alice", 1000000)
	_, err := s.WrapChat(billingChatStub{err: context.DeadlineExceeded}).Chat(ctx, nil, &chat.ChatOptions{MaxTokens: 20})
	if !errors.Is(err, context.DeadlineExceeded) {
		t.Fatalf("lost timeout: %v", err)
	}
	var log OperationLog
	if err := s.db.Where("kind = ?", "system").Take(&log).Error; err != nil {
		t.Fatal(err)
	}
	if log.Metadata["phase"] != "model" || log.Metadata["purpose"] != "digital_human_distillation" || log.Metadata["model_id"] != "model" || log.Metadata["elapsed_ms"] == nil || log.Metadata["timeout_budget_ms"] == nil {
		t.Fatalf("missing failure detail: %+v", log.Metadata)
	}
	if account := accountBalance(t, s, 1, "alice"); account.Reserved != 0 || account.Balance != 1000000 {
		t.Fatal("timeout charged an unreported call or leaked reservation")
	}
}
func TestMeteredStreamSettlementAndCancellation(t *testing.T) {
	for _, mode := range []string{"success", "nil", "cancel", "incomplete"} {
		t.Run(mode, func(t *testing.T) {
			s := sqliteWorkbench(t)
			base := workbenchContext(1, "alice", types.TenantRoleContributor)
			ctx, cancel := context.WithCancel(base)
			defer cancel()
			seedBalance(t, s, 1, "alice", 1000000)
			upstream := make(chan types.StreamResponse, 3)
			model := billingChatStub{stream: upstream}
			switch mode {
			case "success":
				upstream <- types.StreamResponse{Content: "hello"}
				upstream <- types.StreamResponse{Done: true, Usage: &types.TokenUsage{TotalTokens: 21}}
				close(upstream)
			case "incomplete":
				upstream <- types.StreamResponse{Content: "partial"}
				close(upstream)
			case "nil":
				model.stream = nil
			}
			stream, e := s.WrapChat(model).ChatStream(ctx, nil, &chat.ChatOptions{MaxTokens: 20})
			if mode == "nil" {
				if e == nil {
					t.Fatal("nil stream accepted")
				}
			} else {
				if e != nil {
					t.Fatal(e)
				}
				if mode == "cancel" {
					cancel()
				}
				var final types.StreamResponse
				timeout := time.After(2 * time.Second)
				done := false
				for !done {
					select {
					case row, ok := <-stream:
						if !ok {
							done = true
						} else {
							final = row
						}
					case <-timeout:
						t.Fatal("stream stalled")
					}
				}
				if mode == "success" && (!final.Done || final.FinishReason == types.FinishReasonIncomplete) {
					t.Fatal("success lost")
				}
				if mode == "incomplete" && final.FinishReason != types.FinishReasonIncomplete {
					t.Fatal("missing failure")
				}
			}
			got := accountBalance(t, s, 1, "alice")
			if got.Reserved != 0 {
				t.Fatal("stream leaked balance hold")
			}
			if mode == "success" && got.Balance != 999790 {
				t.Fatalf("usage lost: %+v", got)
			}
		})
	}
}