199 lines
7.3 KiB
Go
199 lines
7.3 KiB
Go
package store
|
||
|
||
import (
|
||
"database/sql"
|
||
"errors"
|
||
"fmt"
|
||
"time"
|
||
)
|
||
|
||
// ledger.go: reserve-before-call / settle-after money discipline (порт
|
||
// семантики vojo Reserve/Settle/ReleaseReservation на SQLite). Потолки — на
|
||
// книгу (суммарно) и на день (Р7); проверяются по committed + reserved под
|
||
// immediate-транзакцией write-пула, так что конкурентные вызовы не
|
||
// проскакивают потолок больше чем на одну максимальную резервацию.
|
||
|
||
// ReserveResult is the outcome of a pre-call admission check.
|
||
type ReserveResult int
|
||
|
||
const (
|
||
ReserveOK ReserveResult = iota
|
||
ReserveDeniedBook // per-book USD ceiling hit
|
||
ReserveDeniedDay // daily USD ceiling hit
|
||
)
|
||
|
||
// Reservation is the handle Settle/Release need to undo the admission.
|
||
type Reservation struct {
|
||
BookID string
|
||
Date string // UTC day the reservation was booked under
|
||
Estimate float64
|
||
}
|
||
|
||
// Ceilings are the admission limits. Zero means "no limit" for that axis —
|
||
// fail-fast in config forbids an all-zero pair for real runs.
|
||
type Ceilings struct {
|
||
BookUSD float64
|
||
DayUSD float64
|
||
}
|
||
|
||
func todayUTC() string { return time.Now().UTC().Format("2006-01-02") }
|
||
|
||
// Reserve books estimate USD against the ceilings BEFORE the call. On success
|
||
// the estimate is added to reserved_usd; Settle converts it to committed
|
||
// spend, Release returns it. Denials are results, not errors.
|
||
func (s *Store) Reserve(bookID string, estimate float64, c Ceilings) (Reservation, ReserveResult, error) {
|
||
ctx, cancel := opContext()
|
||
defer cancel()
|
||
day := todayUTC()
|
||
res := Reservation{BookID: bookID, Date: day, Estimate: estimate}
|
||
|
||
tx, err := s.w.BeginTx(ctx, nil) // write pool: immediate tx, serialized
|
||
if err != nil {
|
||
return res, ReserveOK, err
|
||
}
|
||
defer tx.Rollback()
|
||
|
||
var bookTotal, dayTotal float64
|
||
if err := tx.QueryRowContext(ctx,
|
||
`SELECT COALESCE(SUM(committed_usd + reserved_usd), 0) FROM spend WHERE book_id = ?`, bookID,
|
||
).Scan(&bookTotal); err != nil {
|
||
return res, ReserveOK, err
|
||
}
|
||
if err := tx.QueryRowContext(ctx,
|
||
`SELECT COALESCE(SUM(committed_usd + reserved_usd), 0) FROM spend WHERE date = ?`, day,
|
||
).Scan(&dayTotal); err != nil {
|
||
return res, ReserveOK, err
|
||
}
|
||
if c.BookUSD > 0 && bookTotal+estimate > c.BookUSD {
|
||
return res, ReserveDeniedBook, nil
|
||
}
|
||
if c.DayUSD > 0 && dayTotal+estimate > c.DayUSD {
|
||
return res, ReserveDeniedDay, nil
|
||
}
|
||
|
||
if _, err := tx.ExecContext(ctx,
|
||
`INSERT INTO spend (book_id, date, reserved_usd) VALUES (?, ?, ?)
|
||
ON CONFLICT (book_id, date) DO UPDATE SET reserved_usd = reserved_usd + excluded.reserved_usd`,
|
||
bookID, day, estimate); err != nil {
|
||
return res, ReserveOK, err
|
||
}
|
||
if err := tx.Commit(); err != nil {
|
||
return res, ReserveOK, err
|
||
}
|
||
return res, ReserveOK, nil
|
||
}
|
||
|
||
// Release frees a reservation whose call produced no billable spend (transport
|
||
// exhaustion, terminal 4xx before a 2xx). MAX(0, …) guards a double-release.
|
||
func (s *Store) Release(res Reservation) error {
|
||
ctx, cancel := opContext()
|
||
defer cancel()
|
||
_, err := s.w.ExecContext(ctx,
|
||
`UPDATE spend SET reserved_usd = MAX(0, reserved_usd - ?) WHERE book_id = ? AND date = ?`,
|
||
res.Estimate, res.BookID, res.Date)
|
||
return err
|
||
}
|
||
|
||
// Checkpoint is one persisted raw LLM response (см. migrate.go v1).
|
||
type Checkpoint struct {
|
||
RequestHash string
|
||
JobID int64
|
||
ChunkIdx int
|
||
Attempt int
|
||
Stage string
|
||
Role string
|
||
ModelRequested string
|
||
ModelActual string
|
||
ResponseText string
|
||
UsageJSON string
|
||
CostUSD float64
|
||
FinishReason string
|
||
ProviderRequestID string
|
||
}
|
||
|
||
// SettleWithCheckpoint atomically (ONE transaction, one file) converts the
|
||
// reservation into committed spend AND persists the raw response. Это
|
||
// закрытие дыры «settle прошёл — kill -9 — чекпоинт не записан = двойная
|
||
// оплата» (implementation-notes §3.3): после рестарта либо есть и списание, и
|
||
// чекпоинт (вызов не повторится), либо нет ни того ни другого (recovery
|
||
// снимет резерв, вызов повторится и оплатится один раз). Неустранимая потеря
|
||
// — только in-flight вызов, что и есть честные «≤1 вызов» приёмки.
|
||
//
|
||
// Idempotent per request_hash: a duplicate settle for the same hash books
|
||
// nothing and keeps the original checkpoint.
|
||
func (s *Store) SettleWithCheckpoint(res Reservation, cost float64, cp Checkpoint) error {
|
||
ctx, cancel := opContext()
|
||
defer cancel()
|
||
|
||
tx, err := s.w.BeginTx(ctx, nil)
|
||
if err != nil {
|
||
return err
|
||
}
|
||
defer tx.Rollback()
|
||
|
||
ins, err := tx.ExecContext(ctx, `
|
||
INSERT INTO checkpoints (
|
||
request_hash, job_id, chunk_idx, attempt, stage, role,
|
||
model_requested, model_actual, response_text, usage_json,
|
||
cost_usd, finish_reason, provider_request_id
|
||
) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
|
||
ON CONFLICT (request_hash) DO NOTHING`,
|
||
cp.RequestHash, cp.JobID, cp.ChunkIdx, cp.Attempt, cp.Stage, cp.Role,
|
||
cp.ModelRequested, cp.ModelActual, cp.ResponseText, cp.UsageJSON,
|
||
cp.CostUSD, cp.FinishReason, cp.ProviderRequestID)
|
||
if err != nil {
|
||
return fmt.Errorf("store: checkpoint insert: %w", err)
|
||
}
|
||
inserted, _ := ins.RowsAffected()
|
||
if inserted == 0 {
|
||
// Duplicate settle (retried caller): release the reservation, book no
|
||
// new spend — the original settle already did.
|
||
if _, err := tx.ExecContext(ctx,
|
||
`UPDATE spend SET reserved_usd = MAX(0, reserved_usd - ?) WHERE book_id = ? AND date = ?`,
|
||
res.Estimate, res.BookID, res.Date); err != nil {
|
||
return err
|
||
}
|
||
return tx.Commit()
|
||
}
|
||
|
||
if _, err := tx.ExecContext(ctx, `
|
||
INSERT INTO spend (book_id, date, committed_usd, reserved_usd) VALUES (?, ?, ?, 0)
|
||
ON CONFLICT (book_id, date) DO UPDATE SET
|
||
committed_usd = committed_usd + excluded.committed_usd,
|
||
reserved_usd = MAX(0, reserved_usd - ?)`,
|
||
res.BookID, res.Date, cost, res.Estimate); err != nil {
|
||
return fmt.Errorf("store: settle spend: %w", err)
|
||
}
|
||
return tx.Commit()
|
||
}
|
||
|
||
// GetCheckpoint returns the persisted response for a request hash, if any —
|
||
// the resume path: a hit means the call is NOT repeated or re-billed.
|
||
func (s *Store) GetCheckpoint(requestHash string) (*Checkpoint, error) {
|
||
ctx, cancel := opContext()
|
||
defer cancel()
|
||
cp := Checkpoint{RequestHash: requestHash}
|
||
err := s.r.QueryRowContext(ctx, `
|
||
SELECT job_id, chunk_idx, attempt, stage, role, model_requested, model_actual,
|
||
response_text, usage_json, cost_usd, finish_reason, provider_request_id
|
||
FROM checkpoints WHERE request_hash = ?`, requestHash).Scan(
|
||
&cp.JobID, &cp.ChunkIdx, &cp.Attempt, &cp.Stage, &cp.Role, &cp.ModelRequested, &cp.ModelActual,
|
||
&cp.ResponseText, &cp.UsageJSON, &cp.CostUSD, &cp.FinishReason, &cp.ProviderRequestID)
|
||
if errors.Is(err, sql.ErrNoRows) {
|
||
return nil, nil
|
||
}
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
return &cp, nil
|
||
}
|
||
|
||
// SpentUSD reports (committed, reserved) for a book across all days.
|
||
func (s *Store) SpentUSD(bookID string) (committed, reserved float64, err error) {
|
||
ctx, cancel := opContext()
|
||
defer cancel()
|
||
err = s.r.QueryRowContext(ctx,
|
||
`SELECT COALESCE(SUM(committed_usd),0), COALESCE(SUM(reserved_usd),0) FROM spend WHERE book_id = ?`,
|
||
bookID).Scan(&committed, &reserved)
|
||
return
|
||
}
|