mirror of
https://github.com/mudler/LocalAI.git
synced 2026-05-24 00:26:34 -04:00
feat(usage): add /api/auth/usage/sources and admin variant
Self endpoint filters legacy server-side; admin endpoint includes
legacy and accepts user_id + api_key_id filters. Response includes
buckets, totals.{by_source, by_key, grand_total}, and a truncated
flag set when the per-key roll-up was capped at 200.
Refs: #9862
Signed-off-by: Ettore Di Giacinto <mudler@localai.io>
This commit is contained in:
@@ -789,6 +789,30 @@ func RegisterAuthRoutes(e *echo.Echo, app *application.Application) {
|
||||
})
|
||||
})
|
||||
|
||||
// GET /api/auth/usage/sources - caller's per-source breakdown (no legacy)
|
||||
e.GET("/api/auth/usage/sources", func(c echo.Context) error {
|
||||
user := auth.GetUser(c)
|
||||
if user == nil {
|
||||
return c.JSON(http.StatusUnauthorized, map[string]string{"error": "not authenticated"})
|
||||
}
|
||||
|
||||
period := c.QueryParam("period")
|
||||
if period == "" {
|
||||
period = "month"
|
||||
}
|
||||
|
||||
buckets, totals, err := auth.GetUserUsageBySource(db, user.ID, period)
|
||||
if err != nil {
|
||||
return c.JSON(http.StatusInternalServerError, map[string]string{"error": "failed to get usage"})
|
||||
}
|
||||
|
||||
return c.JSON(http.StatusOK, map[string]any{
|
||||
"buckets": buckets,
|
||||
"totals": totals,
|
||||
"truncated": false,
|
||||
})
|
||||
})
|
||||
|
||||
// Admin endpoints
|
||||
adminMw := auth.RequireAdmin()
|
||||
|
||||
@@ -1104,6 +1128,27 @@ func RegisterAuthRoutes(e *echo.Echo, app *application.Application) {
|
||||
})
|
||||
}, adminMw)
|
||||
|
||||
// GET /api/auth/admin/usage/sources - all users' per-source breakdown (admin only)
|
||||
e.GET("/api/auth/admin/usage/sources", func(c echo.Context) error {
|
||||
period := c.QueryParam("period")
|
||||
if period == "" {
|
||||
period = "month"
|
||||
}
|
||||
userID := c.QueryParam("user_id")
|
||||
apiKeyID := c.QueryParam("api_key_id")
|
||||
|
||||
buckets, totals, truncated, err := auth.GetAllUsageBySource(db, period, userID, apiKeyID)
|
||||
if err != nil {
|
||||
return c.JSON(http.StatusInternalServerError, map[string]string{"error": "failed to get usage"})
|
||||
}
|
||||
|
||||
return c.JSON(http.StatusOK, map[string]any{
|
||||
"buckets": buckets,
|
||||
"totals": totals,
|
||||
"truncated": truncated,
|
||||
})
|
||||
}, adminMw)
|
||||
|
||||
// --- Invite management endpoints ---
|
||||
|
||||
// POST /api/auth/admin/invites - create invite (admin only)
|
||||
|
||||
@@ -286,6 +286,42 @@ func newTestAuthApp(db *gorm.DB, appConfig *config.ApplicationConfig) *echo.Echo
|
||||
return c.JSON(http.StatusOK, map[string]string{"message": "user deleted"})
|
||||
}, adminMw)
|
||||
|
||||
// GET /api/auth/usage/sources - mirror of production handler
|
||||
e.GET("/api/auth/usage/sources", func(c echo.Context) error {
|
||||
user := auth.GetUser(c)
|
||||
if user == nil {
|
||||
return c.JSON(http.StatusUnauthorized, map[string]string{"error": "not authenticated"})
|
||||
}
|
||||
period := c.QueryParam("period")
|
||||
if period == "" {
|
||||
period = "month"
|
||||
}
|
||||
buckets, totals, err := auth.GetUserUsageBySource(db, user.ID, period)
|
||||
if err != nil {
|
||||
return c.JSON(http.StatusInternalServerError, map[string]string{"error": "failed to get usage"})
|
||||
}
|
||||
return c.JSON(http.StatusOK, map[string]any{
|
||||
"buckets": buckets, "totals": totals, "truncated": false,
|
||||
})
|
||||
})
|
||||
|
||||
// GET /api/auth/admin/usage/sources - mirror of production handler
|
||||
e.GET("/api/auth/admin/usage/sources", func(c echo.Context) error {
|
||||
period := c.QueryParam("period")
|
||||
if period == "" {
|
||||
period = "month"
|
||||
}
|
||||
userID := c.QueryParam("user_id")
|
||||
apiKeyID := c.QueryParam("api_key_id")
|
||||
buckets, totals, truncated, err := auth.GetAllUsageBySource(db, period, userID, apiKeyID)
|
||||
if err != nil {
|
||||
return c.JSON(http.StatusInternalServerError, map[string]string{"error": "failed to get usage"})
|
||||
}
|
||||
return c.JSON(http.StatusOK, map[string]any{
|
||||
"buckets": buckets, "totals": totals, "truncated": truncated,
|
||||
})
|
||||
}, adminMw)
|
||||
|
||||
// Regular API endpoint for testing
|
||||
e.POST("/v1/chat/completions", func(c echo.Context) error {
|
||||
return c.String(http.StatusOK, "ok")
|
||||
@@ -931,4 +967,110 @@ var _ = Describe("Auth Routes", Label("auth"), func() {
|
||||
Expect(providers).To(ContainElement(auth.ProviderGitHub))
|
||||
})
|
||||
})
|
||||
|
||||
Describe("GET /api/auth/usage/sources", func() {
|
||||
It("returns only the caller's data, never legacy", func() {
|
||||
app := newTestAuthApp(db, appConfig)
|
||||
|
||||
alice := createRouteTestUser(db, "alice@example.com", auth.RoleUser)
|
||||
aliceToken, err := auth.CreateSession(db, alice.ID, "")
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
|
||||
keyID := "k-alice"
|
||||
now := time.Now()
|
||||
Expect(auth.RecordUsage(db, &auth.UsageRecord{
|
||||
UserID: alice.ID, Source: auth.UsageSourceAPIKey,
|
||||
APIKeyID: &keyID, APIKeyName: "alice-key",
|
||||
Model: "gpt-4", TotalTokens: 100, CreatedAt: now,
|
||||
})).To(Succeed())
|
||||
Expect(auth.RecordUsage(db, &auth.UsageRecord{
|
||||
UserID: alice.ID, Source: auth.UsageSourceWeb,
|
||||
Model: "gpt-4", TotalTokens: 50, CreatedAt: now,
|
||||
})).To(Succeed())
|
||||
Expect(auth.RecordUsage(db, &auth.UsageRecord{
|
||||
UserID: "legacy-api-key", Source: auth.UsageSourceLegacy,
|
||||
Model: "gpt-4", TotalTokens: 30, CreatedAt: now,
|
||||
})).To(Succeed())
|
||||
|
||||
rec := doAuthRequest(app, http.MethodGet, "/api/auth/usage/sources?period=month", nil, withSession(aliceToken))
|
||||
Expect(rec.Code).To(Equal(http.StatusOK))
|
||||
|
||||
var resp struct {
|
||||
Buckets []auth.UsageBucket `json:"buckets"`
|
||||
Totals auth.SourceTotals `json:"totals"`
|
||||
Truncated bool `json:"truncated"`
|
||||
}
|
||||
Expect(json.Unmarshal(rec.Body.Bytes(), &resp)).To(Succeed())
|
||||
_, hasLegacy := resp.Totals.BySource[auth.UsageSourceLegacy]
|
||||
Expect(hasLegacy).To(BeFalse())
|
||||
Expect(resp.Totals.GrandTotal.Tokens).To(Equal(int64(150)))
|
||||
Expect(resp.Truncated).To(BeFalse())
|
||||
})
|
||||
|
||||
It("returns 401 when unauthenticated", func() {
|
||||
app := newTestAuthApp(db, appConfig)
|
||||
|
||||
// Without a session cookie or bearer token, the global auth middleware
|
||||
// should refuse the request before our handler runs.
|
||||
rec := doAuthRequest(app, http.MethodGet, "/api/auth/usage/sources?period=month", nil)
|
||||
Expect(rec.Code).To(Equal(http.StatusUnauthorized))
|
||||
})
|
||||
})
|
||||
|
||||
Describe("GET /api/auth/admin/usage/sources", func() {
|
||||
It("returns 403 for non-admin", func() {
|
||||
app := newTestAuthApp(db, appConfig)
|
||||
|
||||
alice := createRouteTestUser(db, "alice@example.com", auth.RoleUser)
|
||||
aliceToken, _ := auth.CreateSession(db, alice.ID, "")
|
||||
|
||||
rec := doAuthRequest(app, http.MethodGet, "/api/auth/admin/usage/sources?period=month", nil, withSession(aliceToken))
|
||||
Expect(rec.Code).To(Equal(http.StatusForbidden))
|
||||
})
|
||||
|
||||
It("returns legacy bucket for admin and applies api_key_id filter", func() {
|
||||
app := newTestAuthApp(db, appConfig)
|
||||
|
||||
admin := createRouteTestUser(db, "admin@example.com", auth.RoleAdmin)
|
||||
adminToken, _ := auth.CreateSession(db, admin.ID, "")
|
||||
|
||||
k1 := "k1"
|
||||
k2 := "k2"
|
||||
now := time.Now()
|
||||
Expect(auth.RecordUsage(db, &auth.UsageRecord{UserID: "alice", Source: auth.UsageSourceAPIKey, APIKeyID: &k1, APIKeyName: "ci", Model: "gpt-4", TotalTokens: 10, CreatedAt: now})).To(Succeed())
|
||||
Expect(auth.RecordUsage(db, &auth.UsageRecord{UserID: "alice", Source: auth.UsageSourceAPIKey, APIKeyID: &k2, APIKeyName: "lap", Model: "gpt-4", TotalTokens: 20, CreatedAt: now})).To(Succeed())
|
||||
Expect(auth.RecordUsage(db, &auth.UsageRecord{UserID: "legacy-api-key", Source: auth.UsageSourceLegacy, Model: "gpt-4", TotalTokens: 5, CreatedAt: now})).To(Succeed())
|
||||
|
||||
rec := doAuthRequest(app, http.MethodGet,
|
||||
"/api/auth/admin/usage/sources?period=month&api_key_id=k2", nil, withSession(adminToken))
|
||||
Expect(rec.Code).To(Equal(http.StatusOK))
|
||||
|
||||
var resp struct {
|
||||
Totals auth.SourceTotals `json:"totals"`
|
||||
Truncated bool `json:"truncated"`
|
||||
}
|
||||
Expect(json.Unmarshal(rec.Body.Bytes(), &resp)).To(Succeed())
|
||||
Expect(resp.Totals.GrandTotal.Tokens).To(Equal(int64(20)))
|
||||
})
|
||||
|
||||
It("includes legacy in by_source for admin with no filter", func() {
|
||||
app := newTestAuthApp(db, appConfig)
|
||||
|
||||
admin := createRouteTestUser(db, "admin@example.com", auth.RoleAdmin)
|
||||
adminToken, _ := auth.CreateSession(db, admin.ID, "")
|
||||
|
||||
now := time.Now()
|
||||
Expect(auth.RecordUsage(db, &auth.UsageRecord{UserID: "legacy-api-key", Source: auth.UsageSourceLegacy, Model: "gpt-4", TotalTokens: 7, CreatedAt: now})).To(Succeed())
|
||||
|
||||
rec := doAuthRequest(app, http.MethodGet, "/api/auth/admin/usage/sources?period=month", nil, withSession(adminToken))
|
||||
Expect(rec.Code).To(Equal(http.StatusOK))
|
||||
|
||||
var resp struct {
|
||||
Totals auth.SourceTotals `json:"totals"`
|
||||
}
|
||||
Expect(json.Unmarshal(rec.Body.Bytes(), &resp)).To(Succeed())
|
||||
Expect(resp.Totals.BySource).To(HaveKey(auth.UsageSourceLegacy))
|
||||
Expect(resp.Totals.BySource[auth.UsageSourceLegacy].Tokens).To(Equal(int64(7)))
|
||||
})
|
||||
})
|
||||
})
|
||||
|
||||
Reference in New Issue
Block a user