wehub-resource-sync
f99010fae1
Desktop Artifacts / Desktop Build (Linux) (push) Waiting to run
Desktop Artifacts / Desktop Build (Windows) (push) Waiting to run
Desktop Artifacts / Desktop Build (Linux (arm64)) (push) Waiting to run
Desktop Artifacts (macOS) / Desktop Build (macOS (aarch64)) (push) Waiting to run
Desktop Artifacts (macOS) / Desktop Build (macOS (x86_64)) (push) Waiting to run
CI / lint (push) Failing after 1s
CI / frontend (push) Failing after 1s
CI / scripts (push) Failing after 1s
CI / Go Test (ubuntu-latest) (push) Failing after 0s
CI / frontend-node-25 (push) Failing after 1s
CI / docs (push) Failing after 0s
CI / coverage (push) Failing after 0s
CI / e2e (push) Failing after 0s
Docker / build-and-push (push) Failing after 1s
CI / integration (push) Failing after 4m43s
CI / Go Test (windows-latest) (push) Has been cancelled
CI / Desktop Unit Tests (Windows) (push) Has been cancelled
518 lines
16 KiB
Go
518 lines
16 KiB
Go
package postgres
|
|
|
|
import (
|
|
"context"
|
|
"database/sql"
|
|
"database/sql/driver"
|
|
"errors"
|
|
"io"
|
|
"strings"
|
|
"sync"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/stretchr/testify/assert"
|
|
"github.com/stretchr/testify/require"
|
|
|
|
"go.kenn.io/agentsview/internal/db"
|
|
"go.kenn.io/agentsview/internal/export"
|
|
)
|
|
|
|
type usageProbeDriver struct{}
|
|
|
|
type usageProbeConn struct {
|
|
state *usageProbeState
|
|
}
|
|
|
|
type usageProbeRows struct {
|
|
columns []string
|
|
values [][]driver.Value
|
|
next int
|
|
}
|
|
|
|
type usageProbeState struct {
|
|
mu sync.Mutex
|
|
queries []string
|
|
}
|
|
|
|
var (
|
|
usageProbeRegisterOnce sync.Once
|
|
usageProbeStatesMu sync.Mutex
|
|
usageProbeStates = map[string]*usageProbeState{}
|
|
)
|
|
|
|
func newUsageProbeDB(
|
|
t *testing.T, state *usageProbeState,
|
|
) *sql.DB {
|
|
t.Helper()
|
|
usageProbeRegisterOnce.Do(func() {
|
|
sql.Register("agentsview_usage_probe", usageProbeDriver{})
|
|
})
|
|
name := t.Name()
|
|
usageProbeStatesMu.Lock()
|
|
usageProbeStates[name] = state
|
|
usageProbeStatesMu.Unlock()
|
|
t.Cleanup(func() {
|
|
usageProbeStatesMu.Lock()
|
|
delete(usageProbeStates, name)
|
|
usageProbeStatesMu.Unlock()
|
|
})
|
|
|
|
pg, err := sql.Open("agentsview_usage_probe", name)
|
|
require.NoError(t, err, "open usage probe db")
|
|
t.Cleanup(func() { pg.Close() })
|
|
return pg
|
|
}
|
|
|
|
func (usageProbeDriver) Open(name string) (driver.Conn, error) {
|
|
usageProbeStatesMu.Lock()
|
|
state := usageProbeStates[name]
|
|
usageProbeStatesMu.Unlock()
|
|
return &usageProbeConn{state: state}, nil
|
|
}
|
|
|
|
func (c *usageProbeConn) Prepare(string) (driver.Stmt, error) {
|
|
return nil, errors.New("prepare not implemented")
|
|
}
|
|
|
|
func (c *usageProbeConn) Close() error { return nil }
|
|
|
|
func (c *usageProbeConn) Begin() (driver.Tx, error) {
|
|
return nil, errors.New("begin not implemented")
|
|
}
|
|
|
|
func (c *usageProbeConn) QueryContext(
|
|
_ context.Context, query string, _ []driver.NamedValue,
|
|
) (driver.Rows, error) {
|
|
c.state.mu.Lock()
|
|
c.state.queries = append(c.state.queries, query)
|
|
c.state.mu.Unlock()
|
|
|
|
normalized := strings.ToLower(query)
|
|
if strings.Contains(normalized, "from model_pricing") {
|
|
return &usageProbeRows{
|
|
columns: []string{
|
|
"model_pattern",
|
|
"input_per_mtok",
|
|
"output_per_mtok",
|
|
"cache_creation_per_mtok",
|
|
"cache_read_per_mtok",
|
|
"updated_at",
|
|
},
|
|
values: [][]driver.Value{{
|
|
"claude-sonnet", 3.0, 15.0, 3.75, 0.3, "2026-06-08",
|
|
}},
|
|
}, nil
|
|
}
|
|
if strings.Contains(normalized, "from source_archives") {
|
|
return &usageProbeRows{
|
|
columns: []string{"source_archive_id", "source_archive_salt"},
|
|
values: [][]driver.Value{{"probe-archive", "probe-salt"}},
|
|
}, nil
|
|
}
|
|
if strings.Contains(normalized, "from source_project_identity_observations") {
|
|
return &usageProbeRows{
|
|
columns: []string{
|
|
"project",
|
|
"machine",
|
|
"root_path",
|
|
"git_remote",
|
|
"git_remote_name",
|
|
"worktree_name",
|
|
"worktree_root_path",
|
|
"observed_at",
|
|
"normalized_remote",
|
|
"key_source",
|
|
"key",
|
|
},
|
|
}, nil
|
|
}
|
|
if strings.Contains(normalized, "select project, cwd") &&
|
|
strings.Contains(normalized, "from sessions") {
|
|
return &usageProbeRows{
|
|
columns: []string{"project", "cwd"},
|
|
}, nil
|
|
}
|
|
if strings.Contains(normalized, "select id from sessions") {
|
|
return &usageProbeRows{
|
|
columns: []string{"id"},
|
|
values: [][]driver.Value{
|
|
{"kimi:project-hash:session-uuid"},
|
|
{"openclaw:project-hash:session-uuid"},
|
|
},
|
|
}, nil
|
|
}
|
|
if strings.Contains(normalized, "select count(*)") &&
|
|
strings.Contains(normalized, "from sessions s") {
|
|
return &usageProbeRows{
|
|
columns: []string{"count"},
|
|
values: [][]driver.Value{
|
|
{int64(1)},
|
|
},
|
|
}, nil
|
|
}
|
|
if strings.Contains(normalized, "from (") &&
|
|
strings.Contains(normalized, "from messages") {
|
|
ts := time.Date(2024, 6, 15, 10, 0, 0, 0, time.UTC)
|
|
return &usageProbeRows{
|
|
columns: []string{
|
|
"session_id",
|
|
"message_ordinal",
|
|
"usage_source",
|
|
"ts",
|
|
"model",
|
|
"token_usage",
|
|
"input_tokens",
|
|
"output_tokens",
|
|
"cache_creation_input_tokens",
|
|
"cache_read_input_tokens",
|
|
"reasoning_tokens",
|
|
"cost_usd",
|
|
"claude_message_id",
|
|
"claude_request_id",
|
|
"source_uuid",
|
|
"usage_dedup_key",
|
|
"project",
|
|
"agent",
|
|
},
|
|
values: [][]driver.Value{
|
|
usageProbeUsageRow("s-parent", "proj-a", "claude", ts),
|
|
usageProbeUsageRow("s-fork", "proj-b", "codex", ts.Add(time.Minute)),
|
|
},
|
|
}, nil
|
|
}
|
|
return nil, errors.New("unexpected usage query")
|
|
}
|
|
|
|
func usageProbeUsageRow(
|
|
sessionID, project, agent string, ts time.Time,
|
|
) []driver.Value {
|
|
return []driver.Value{
|
|
sessionID,
|
|
int64(0),
|
|
"message",
|
|
ts,
|
|
"claude-sonnet",
|
|
`{"input_tokens":100,"output_tokens":50}`,
|
|
int64(0),
|
|
int64(0),
|
|
int64(0),
|
|
int64(0),
|
|
int64(0),
|
|
nil,
|
|
"msg-dup",
|
|
"req-dup",
|
|
"",
|
|
"",
|
|
project,
|
|
agent,
|
|
}
|
|
}
|
|
|
|
func (r *usageProbeRows) Columns() []string { return r.columns }
|
|
|
|
func (r *usageProbeRows) Close() error { return nil }
|
|
|
|
func (r *usageProbeRows) Next(dest []driver.Value) error {
|
|
if r.next >= len(r.values) {
|
|
return io.EOF
|
|
}
|
|
copy(dest, r.values[r.next])
|
|
r.next++
|
|
return nil
|
|
}
|
|
|
|
func TestPGGetDailyUsageReturnsDedupedSessionCounts(t *testing.T) {
|
|
store := &Store{
|
|
pg: newUsageProbeDB(t, &usageProbeState{}),
|
|
}
|
|
|
|
result, err := store.GetDailyUsage(context.Background(), db.UsageFilter{
|
|
From: "2024-06-15",
|
|
To: "2024-06-15",
|
|
})
|
|
require.NoError(t, err, "GetDailyUsage")
|
|
|
|
assert.Equal(t, 1, result.SessionCounts.Total)
|
|
countsByDisplay := make(map[string]int, len(result.Projects))
|
|
for key, project := range result.Projects {
|
|
countsByDisplay[project.DisplayLabel] = result.SessionCounts.ByProject[key]
|
|
assert.NotContains(t, key, project.DisplayLabel)
|
|
}
|
|
assert.Equal(t, map[string]int{"proj-a": 1}, countsByDisplay)
|
|
assert.Equal(t, 1, result.SessionCounts.ByAgent["claude"])
|
|
assert.NotContains(t, countsByDisplay, "proj-b")
|
|
assert.Zero(t, result.SessionCounts.ByAgent["codex"])
|
|
}
|
|
|
|
func TestPGUsageDedupTokenForRowFallsBackToSourceUUIDWhenClaudePairIncomplete(t *testing.T) {
|
|
got, ok := pgUsageDedupTokenForRow(
|
|
"message",
|
|
"claude-code",
|
|
"msg-dup",
|
|
"",
|
|
"source-dup",
|
|
"",
|
|
)
|
|
require.True(t, ok, "expected source_uuid fallback key")
|
|
assert.Equal(t, pgUsageDedupToken{
|
|
kind: "source",
|
|
value: "claude-code:source-dup",
|
|
}, got)
|
|
}
|
|
|
|
func TestPGUsageAmountsPreserveSessionSummaryUsageEventTokens(t *testing.T) {
|
|
rawInput := db.MaxPlausibleTokens + 250_000
|
|
rawOutput := db.MaxPlausibleTokens + 500_000
|
|
resolver := export.NewPricingResolver([]export.EffectivePricingRow{{
|
|
ModelPattern: "gpt-5.4",
|
|
Rates: export.ModelRates{
|
|
InputPerMTok: 1.0, OutputPerMTok: 2.0,
|
|
},
|
|
}})
|
|
|
|
inTok, outTok, _, _, cost, _ := pgDailyUsageAmounts(
|
|
pgDailyUsageScanRow{
|
|
usageSource: "session",
|
|
model: "gpt-5.4",
|
|
inputTokens: rawInput,
|
|
outputTokens: rawOutput,
|
|
},
|
|
resolver,
|
|
)
|
|
assert.Equal(t, rawInput, inTok, "daily input")
|
|
assert.Equal(t, rawOutput, outTok, "daily output")
|
|
wantCost := (float64(rawInput)*1.0 + float64(rawOutput)*2.0) / 1_000_000
|
|
assert.InDelta(t, wantCost, cost, 1e-9, "daily cost")
|
|
|
|
cost, priced, contributes := pgSessionRowCost(pgUsageScanRow{
|
|
usageSource: "session",
|
|
model: "gpt-5.4",
|
|
inputTokens: rawInput,
|
|
outputTokens: rawOutput,
|
|
}, resolver)
|
|
require.True(t, priced, "priced")
|
|
require.True(t, contributes, "contributes")
|
|
assert.InDelta(t, wantCost, cost, 1e-9, "session cost")
|
|
}
|
|
|
|
func TestPGUsageRowQueryPushesDateBoundsIntoUnion(t *testing.T) {
|
|
pb := ¶mBuilder{}
|
|
query := pgUsageRowQuery(pb, db.UsageFilter{
|
|
From: "2024-06-01",
|
|
To: "2024-06-30",
|
|
ExcludeAutomated: true,
|
|
})
|
|
|
|
normalized := strings.ToLower(query)
|
|
assert.NotContains(t, normalized, "and u.ts >=")
|
|
assert.NotContains(t, normalized, "and u.ts <=")
|
|
assert.NotContains(t, normalized, " or ")
|
|
assert.NotContains(t, normalized, "display_name")
|
|
assert.NotContains(t, normalized, "first_message")
|
|
assert.NotContains(t, normalized, "cost_status")
|
|
assert.NotContains(t, normalized, "cost_source")
|
|
assert.Contains(t, normalized, "u.reasoning_tokens")
|
|
assert.NotContains(t, normalized, "user_message_count")
|
|
assert.NotContains(t, normalized, "session_activity_at")
|
|
assert.NotContains(t, normalized, " as started_at")
|
|
assert.NotContains(t, normalized, "u.machine")
|
|
assert.Contains(t, normalized, "message_timestamp_rows as materialized")
|
|
assert.Contains(t, normalized, "usage_event_timestamp_rows as materialized")
|
|
assert.Contains(t, normalized, "from message_timestamp_rows m\njoin sessions s")
|
|
assert.Contains(t, normalized, "from usage_event_timestamp_rows ue\njoin sessions s")
|
|
assert.Contains(t, normalized, "m.timestamp is not null")
|
|
assert.Contains(t, normalized, "ue.occurred_at is not null")
|
|
assert.Contains(t, normalized, "m.timestamp is null")
|
|
assert.Contains(t, normalized, "ue.occurred_at is null")
|
|
assert.Contains(t, normalized, "m.timestamp >= $1::timestamptz")
|
|
assert.Contains(t, normalized, "ue.occurred_at >= $1::timestamptz")
|
|
assert.Contains(t, normalized, "s.started_at >= $1::timestamptz")
|
|
assert.Contains(t, normalized, "m.timestamp <= $2::timestamptz")
|
|
assert.Contains(t, normalized, "ue.occurred_at <= $2::timestamptz")
|
|
assert.Contains(t, normalized, "s.started_at <= $2::timestamptz")
|
|
require.Len(t, pb.args, 2)
|
|
assert.Equal(t, "2024-05-31T10:00:00Z", pb.args[0])
|
|
assert.Equal(t, "2024-07-01T13:59:59Z", pb.args[1])
|
|
}
|
|
|
|
func TestPGBoundedDailyUsageRowsCTEProjectsReasoningTokens(t *testing.T) {
|
|
pb := ¶mBuilder{}
|
|
f := db.UsageFilter{
|
|
From: "2024-06-01",
|
|
To: "2024-06-30",
|
|
}
|
|
query := pgDailyUsageRowsSQLForBounds(pb, f, pgUsageBoundsForFilter(pb, f))
|
|
|
|
normalized := strings.ToLower(query)
|
|
assert.Contains(t, normalized, "usage_event_timestamp_rows as materialized")
|
|
assert.Contains(t, normalized, "ue.cache_read_input_tokens,\n\t\tue.reasoning_tokens,\n\t\tue.cost_usd")
|
|
assert.Contains(t, normalized, "from usage_event_timestamp_rows ue\njoin sessions s")
|
|
}
|
|
|
|
// TestPGGetUsageMatchingSessionCountUsesSessionQuery exercises the
|
|
// unbounded (no From/To) code path, which counts sessions directly with
|
|
// EXISTS subqueries built from the same relaxed matching eligibility
|
|
// predicates as the bounded branch. Bounded filters take the
|
|
// timestamped-CTE path (see TestPGMatchingUsageRowsSQLForBoundsRelaxesTokenEligibility
|
|
// for that SQL shape) and are not exercised by this probe-mock test.
|
|
func TestPGGetUsageMatchingSessionCountUsesSessionQuery(t *testing.T) {
|
|
state := &usageProbeState{}
|
|
store := &Store{
|
|
pg: newUsageProbeDB(t, state),
|
|
}
|
|
|
|
count, err := store.GetUsageMatchingSessionCount(context.Background(), db.UsageFilter{
|
|
Agent: "copilot",
|
|
Model: "gpt-5.3-codex",
|
|
})
|
|
require.NoError(t, err, "GetUsageMatchingSessionCount")
|
|
assert.Equal(t, 1, count)
|
|
|
|
state.mu.Lock()
|
|
queries := append([]string(nil), state.queries...)
|
|
state.mu.Unlock()
|
|
require.NotEmpty(t, queries)
|
|
|
|
last := strings.ToLower(queries[len(queries)-1])
|
|
assert.Contains(t, last, "select count(*)")
|
|
assert.Contains(t, last, "from sessions s")
|
|
assert.Contains(t, last, "exists (")
|
|
assert.Contains(t, last, "from messages m")
|
|
assert.Contains(t, last, "from usage_events ue")
|
|
assert.Contains(t, last, "s.agent = ")
|
|
assert.Contains(t, last, "m.model = ")
|
|
// The message EXISTS uses the relaxed matching eligibility: assistant
|
|
// rows without requiring a model name, so empty-model Copilot
|
|
// assistant messages match the same way they do on the bounded path.
|
|
assert.Contains(t, last, "m.role = 'assistant'")
|
|
assert.NotContains(t, last, "m.model != ''")
|
|
}
|
|
|
|
// TestPGMatchingUsageRowsSQLForBoundsRelaxesTokenEligibility asserts the
|
|
// bounded matching-session query relaxes token eligibility (no
|
|
// m.token_usage check) and model-presence eligibility (m.role = 'assistant'
|
|
// instead of m.model != ”, since some Copilot assistant messages parse
|
|
// before a model name is known) while filtering Model/ExcludeModel
|
|
// directly on the bounded message/event row, matching the normal bounded
|
|
// path instead of folding in a session-wide model-match EXISTS clause.
|
|
// Mirrors TestPGUsageRowQueryPushesDateBoundsIntoUnion's direct-call style
|
|
// with no live DB and no probe mock.
|
|
func TestPGMatchingUsageRowsSQLForBoundsRelaxesTokenEligibility(t *testing.T) {
|
|
pb := ¶mBuilder{}
|
|
f := db.UsageFilter{
|
|
From: "2024-06-01",
|
|
To: "2024-06-30",
|
|
Model: "gpt-5.3-codex",
|
|
}
|
|
bounds := pgUsageBoundsForFilter(pb, f)
|
|
query := pgMatchingUsageRowsSQLForBounds(pb, f, bounds)
|
|
|
|
normalized := strings.ToLower(query)
|
|
assert.Contains(t, normalized, "message_timestamp_rows as materialized")
|
|
assert.Contains(t, normalized, "usage_event_timestamp_rows as materialized")
|
|
assert.Contains(t, normalized, "m.role = 'assistant'")
|
|
assert.NotContains(t, normalized, "m.model != ''")
|
|
assert.NotContains(t, normalized, "m.token_usage != ''")
|
|
// Model is filtered on the bounded row directly, not via a
|
|
// session-wide EXISTS: each of the four branches (message-timestamp
|
|
// source, event-timestamp source, message fallback, event fallback)
|
|
// applies its own m.model/ue.model comparison, no EXISTS subqueries.
|
|
assert.Equal(t, 0, strings.Count(normalized, "exists ("))
|
|
assert.Equal(t, 2, strings.Count(normalized, "m.model = "))
|
|
assert.Equal(t, 2, strings.Count(normalized, "ue.model = "))
|
|
}
|
|
|
|
func TestPGTopSessionsUsageRowQueryUsesNarrowScan(t *testing.T) {
|
|
pb := ¶mBuilder{}
|
|
query := pgTopSessionsUsageRowQuery(pb, db.UsageFilter{
|
|
From: "2024-06-01",
|
|
To: "2024-06-30",
|
|
})
|
|
|
|
normalized := strings.ToLower(query)
|
|
assert.NotContains(t, normalized, "display_name")
|
|
assert.NotContains(t, normalized, "first_message")
|
|
assert.NotContains(t, normalized, "cost_status")
|
|
assert.NotContains(t, normalized, "cost_source")
|
|
assert.Contains(t, normalized, "u.reasoning_tokens")
|
|
assert.NotContains(t, normalized, "user_message_count")
|
|
assert.NotContains(t, normalized, "session_activity_at")
|
|
assert.NotContains(t, normalized, " as started_at")
|
|
assert.NotContains(t, normalized, "u.machine")
|
|
assert.Contains(t, normalized, "m.timestamp is not null")
|
|
assert.Contains(t, normalized, "ue.occurred_at is not null")
|
|
assert.Contains(t, normalized, "m.timestamp is null")
|
|
assert.Contains(t, normalized, "ue.occurred_at is null")
|
|
assert.Contains(t, normalized, "m.timestamp >= $1::timestamptz")
|
|
assert.Contains(t, normalized, "ue.occurred_at >= $1::timestamptz")
|
|
assert.Contains(t, normalized,
|
|
"m.timestamp is null\n\tand s.started_at >= $1::timestamptz")
|
|
assert.Contains(t, normalized,
|
|
"ue.occurred_at is null\n\tand s.started_at >= $1::timestamptz")
|
|
assert.Contains(t, normalized, "m.timestamp <= $2::timestamptz")
|
|
assert.Contains(t, normalized, "ue.occurred_at <= $2::timestamptz")
|
|
require.Len(t, pb.args, 2)
|
|
assert.Equal(t, "2024-05-31T10:00:00Z", pb.args[0])
|
|
assert.Equal(t, "2024-07-01T13:59:59Z", pb.args[1])
|
|
}
|
|
|
|
func TestPGSessionRowCostIncludesReasoningOnlyRows(t *testing.T) {
|
|
resolver := export.NewPricingResolver(
|
|
[]export.EffectivePricingRow{{
|
|
ModelPattern: "reasoning-model",
|
|
Rates: export.ModelRates{
|
|
OutputPerMTok: 20,
|
|
Source: export.PricingRowSourceFetched,
|
|
},
|
|
}},
|
|
)
|
|
|
|
cost, priced, contributes := pgSessionRowCost(pgUsageScanRow{
|
|
usageSource: "provider",
|
|
model: "reasoning-model",
|
|
reasoningTokens: 25,
|
|
}, resolver)
|
|
|
|
assert.True(t, contributes)
|
|
assert.True(t, priced)
|
|
assert.InDelta(t, 0.0005, cost, 0.0000001)
|
|
block, err := resolver.BuildBlock()
|
|
require.NoError(t, err)
|
|
require.Contains(t, block.Models, "reasoning-model")
|
|
assert.Equal(t, export.CostSourceComputed,
|
|
block.Models["reasoning-model"].CostSource)
|
|
}
|
|
|
|
func TestPGUsageAmountsIncludeMessageReasoningTokens(t *testing.T) {
|
|
resolver := export.NewPricingResolver(
|
|
[]export.EffectivePricingRow{{
|
|
ModelPattern: "gpt-5.4",
|
|
Rates: export.ModelRates{
|
|
InputPerMTok: 1,
|
|
OutputPerMTok: 2,
|
|
},
|
|
}},
|
|
)
|
|
row := pgDailyUsageScanRow{
|
|
usageSource: "message",
|
|
model: "gpt-5.4",
|
|
tokenJSON: `{"input_tokens":1000,"output_tokens":0,` +
|
|
`"reasoning_tokens":500}`,
|
|
}
|
|
|
|
inTok, outTok, _, _, cost, _ := pgDailyUsageAmounts(row, resolver)
|
|
assert.Equal(t, 1000, inTok)
|
|
assert.Zero(t, outTok)
|
|
assert.InDelta(t, 0.002, cost, 1e-12)
|
|
|
|
sessionCost, priced, contributes := pgSessionRowCost(pgUsageScanRow{
|
|
usageSource: "message",
|
|
model: "gpt-5.4",
|
|
tokenJSON: row.tokenJSON,
|
|
}, resolver)
|
|
assert.True(t, priced)
|
|
assert.True(t, contributes)
|
|
assert.InDelta(t, 0.002, sessionCost, 1e-12)
|
|
}
|