mirror of
https://github.com/ollama/ollama.git
synced 2026-07-23 17:21:27 -05:00
Compare commits
2 Commits
v0.32.2-rc
...
parth-usag
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
90dd5b3a70 | ||
|
|
4855c61358 |
@@ -510,3 +510,13 @@ func (c *Client) Whoami(ctx context.Context) (*UserResponse, error) {
|
||||
}
|
||||
return &resp, nil
|
||||
}
|
||||
|
||||
// Usage returns the authenticated user's recent activity and included-usage
|
||||
// limits.
|
||||
func (c *Client) Usage(ctx context.Context) (*UsageResponse, error) {
|
||||
var resp UsageResponse
|
||||
if err := c.do(ctx, http.MethodGet, "/api/usage", nil, &resp); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return &resp, nil
|
||||
}
|
||||
|
||||
@@ -9,6 +9,7 @@ import (
|
||||
"net/url"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
func TestClientFromEnvironment(t *testing.T) {
|
||||
@@ -51,6 +52,68 @@ func TestClientFromEnvironment(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestClientUsage(t *testing.T) {
|
||||
startsAt := time.Date(2026, time.June, 29, 0, 0, 0, 0, time.UTC)
|
||||
endsAt := time.Date(2026, time.July, 27, 0, 0, 0, 0, time.UTC)
|
||||
|
||||
ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
if r.Method != http.MethodGet {
|
||||
t.Errorf("method = %q, want GET", r.Method)
|
||||
}
|
||||
if r.URL.Path != "/api/usage" {
|
||||
t.Errorf("path = %q, want /api/usage", r.URL.Path)
|
||||
}
|
||||
|
||||
if err := json.NewEncoder(w).Encode(UsageResponse{
|
||||
Activity: UsageActivity{
|
||||
Cost: "0.00709",
|
||||
Period: UsagePeriod{
|
||||
Type: "last_4_weeks",
|
||||
StartingAt: startsAt,
|
||||
EndingAt: endsAt,
|
||||
},
|
||||
Models: []UsageActivityModel{{
|
||||
Name: "qwen3-coder:480b",
|
||||
RequestCount: 1,
|
||||
Cost: "0.00709",
|
||||
}},
|
||||
},
|
||||
Limits: UsageLimits{
|
||||
Session: UsageLimit{Usage: 0.006, Models: []UsageLimitModel{{Name: "qwen3-coder:480b", RequestCount: 2}}},
|
||||
Weekly: UsageLimit{Usage: 0.002, Models: []UsageLimitModel{{Name: "web search", RequestCount: 1}}},
|
||||
},
|
||||
}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}))
|
||||
defer ts.Close()
|
||||
|
||||
base, err := url.Parse(ts.URL)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
got, err := NewClient(base, ts.Client()).Usage(t.Context())
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if got.Activity.Cost != "0.00709" {
|
||||
t.Errorf("activity cost = %q, want 0.00709", got.Activity.Cost)
|
||||
}
|
||||
if got.Activity.Period.Type != "last_4_weeks" || !got.Activity.Period.StartingAt.Equal(startsAt) || !got.Activity.Period.EndingAt.Equal(endsAt) {
|
||||
t.Errorf("period = %#v, want last four weeks from %v to %v", got.Activity.Period, startsAt, endsAt)
|
||||
}
|
||||
if len(got.Activity.Models) != 1 || got.Activity.Models[0].Name != "qwen3-coder:480b" || got.Activity.Models[0].RequestCount != 1 || got.Activity.Models[0].Cost != "0.00709" {
|
||||
t.Errorf("activity models = %#v, want qwen activity", got.Activity.Models)
|
||||
}
|
||||
if got.Limits.Session.Usage != 0.006 || len(got.Limits.Session.Models) != 1 || got.Limits.Session.Models[0].Name != "qwen3-coder:480b" {
|
||||
t.Errorf("session limit = %#v, want qwen usage at 0.006", got.Limits.Session)
|
||||
}
|
||||
if got.Limits.Weekly.Usage != 0.002 || len(got.Limits.Weekly.Models) != 1 || got.Limits.Weekly.Models[0].Name != "web search" {
|
||||
t.Errorf("weekly limit = %#v, want web search usage at 0.002", got.Limits.Weekly)
|
||||
}
|
||||
}
|
||||
|
||||
// testError represents an internal error type with status code and message
|
||||
// this is used since the error response from the server is not a standard error struct
|
||||
type testError struct {
|
||||
|
||||
45
api/types.go
45
api/types.go
@@ -978,6 +978,51 @@ type UserResponse struct {
|
||||
Plan string `json:"plan,omitempty"`
|
||||
}
|
||||
|
||||
// UsageResponse reports recent activity and included-usage limits.
|
||||
type UsageResponse struct {
|
||||
Activity UsageActivity `json:"activity"`
|
||||
Limits UsageLimits `json:"limits"`
|
||||
}
|
||||
|
||||
// UsageActivity reports usage activity over a period.
|
||||
type UsageActivity struct {
|
||||
Cost string `json:"cost"`
|
||||
Period UsagePeriod `json:"period"`
|
||||
Models []UsageActivityModel `json:"models"`
|
||||
}
|
||||
|
||||
// UsagePeriod describes the time window the usage covers.
|
||||
type UsagePeriod struct {
|
||||
Type string `json:"type"`
|
||||
StartingAt time.Time `json:"starting_at"`
|
||||
EndingAt time.Time `json:"ending_at"`
|
||||
}
|
||||
|
||||
// UsageActivityModel reports per-model activity within a period.
|
||||
type UsageActivityModel struct {
|
||||
Name string `json:"name"`
|
||||
RequestCount int `json:"request_count"`
|
||||
Cost string `json:"cost"`
|
||||
}
|
||||
|
||||
// UsageLimits reports included usage for the current session and week.
|
||||
type UsageLimits struct {
|
||||
Session UsageLimit `json:"session"`
|
||||
Weekly UsageLimit `json:"weekly"`
|
||||
}
|
||||
|
||||
// UsageLimit reports the consumed fraction of an included-usage limit.
|
||||
type UsageLimit struct {
|
||||
Usage float64 `json:"usage"`
|
||||
Models []UsageLimitModel `json:"models"`
|
||||
}
|
||||
|
||||
// UsageLimitModel reports a model's contribution to an included-usage limit.
|
||||
type UsageLimitModel struct {
|
||||
Name string `json:"name"`
|
||||
RequestCount int `json:"request_count"`
|
||||
}
|
||||
|
||||
// Tensor describes the metadata for a given tensor.
|
||||
type Tensor struct {
|
||||
Name string `json:"name"`
|
||||
|
||||
122
cmd/cmd.go
122
cmd/cmd.go
@@ -27,8 +27,10 @@ import (
|
||||
"strings"
|
||||
"sync/atomic"
|
||||
"syscall"
|
||||
"text/tabwriter"
|
||||
"time"
|
||||
|
||||
"github.com/charmbracelet/lipgloss"
|
||||
"github.com/containerd/console"
|
||||
"github.com/mattn/go-runewidth"
|
||||
"github.com/olekukonko/tablewriter"
|
||||
@@ -977,6 +979,116 @@ func SignoutHandler(cmd *cobra.Command, args []string) error {
|
||||
return nil
|
||||
}
|
||||
|
||||
func UsageHandler(cmd *cobra.Command, args []string) error {
|
||||
out := cmd.OutOrStdout()
|
||||
|
||||
client, err := api.ClientFromEnvironment()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
status, err := client.CloudStatusExperimental(cmd.Context())
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if status.Cloud.Disabled {
|
||||
fmt.Fprintln(out, "Ollama Cloud is disabled; usage is unavailable")
|
||||
return nil
|
||||
}
|
||||
|
||||
usage, err := client.Usage(cmd.Context())
|
||||
if err != nil {
|
||||
var aErr api.AuthorizationError
|
||||
if errors.As(err, &aErr) && aErr.StatusCode == http.StatusUnauthorized {
|
||||
fmt.Fprintln(out, "You need to be signed in to Ollama to view usage.")
|
||||
fmt.Fprintln(out)
|
||||
if aErr.SigninURL != "" {
|
||||
_ = browser.OpenURL(aErr.SigninURL)
|
||||
fmt.Fprintf(out, ConnectInstructions, aErr.SigninURL)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
return err
|
||||
}
|
||||
|
||||
fmt.Fprintln(out, usageHeading(out, "Usage"))
|
||||
details := tabwriter.NewWriter(out, 0, 4, 2, ' ', 0)
|
||||
fmt.Fprintf(details, " Period\t%s to %s\n", usage.Activity.Period.StartingAt.Format("2006-01-02"), usage.Activity.Period.EndingAt.Format("2006-01-02"))
|
||||
fmt.Fprintf(details, " Spend\t$%s\n", usage.Activity.Cost)
|
||||
if err := details.Flush(); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
hasLimits := usage.Limits.Session.Usage > 0 || len(usage.Limits.Session.Models) > 0 ||
|
||||
usage.Limits.Weekly.Usage > 0 || len(usage.Limits.Weekly.Models) > 0
|
||||
if len(usage.Activity.Models) == 0 && !hasLimits {
|
||||
fmt.Fprintln(out)
|
||||
fmt.Fprintln(out, "No usage recorded for this period.")
|
||||
return nil
|
||||
}
|
||||
|
||||
if usage.Activity.Cost != "0.00000" && len(usage.Activity.Models) > 0 {
|
||||
fmt.Fprintln(out)
|
||||
fmt.Fprintln(out, usageHeading(out, "Activity"))
|
||||
table := tabwriter.NewWriter(out, 0, 4, 2, ' ', 0)
|
||||
fmt.Fprintln(table, " Model\tRequests\tSpend")
|
||||
for _, m := range usage.Activity.Models {
|
||||
fmt.Fprintf(table, " %s\t%d\t$%s\n", usageModelName(m.Name), m.RequestCount, m.Cost)
|
||||
}
|
||||
if err := table.Flush(); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
|
||||
if err := writeUsageLimit(out, "Session", usage.Limits.Session); err != nil {
|
||||
return err
|
||||
}
|
||||
if err := writeUsageLimit(out, "Weekly", usage.Limits.Weekly); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
var usageTitleStyle = lipgloss.NewStyle().Bold(true)
|
||||
|
||||
func usageHeading(out io.Writer, text string) string {
|
||||
if f, ok := out.(*os.File); ok && term.IsTerminal(int(f.Fd())) {
|
||||
return usageTitleStyle.Render(text)
|
||||
}
|
||||
return text
|
||||
}
|
||||
|
||||
func usageModelName(name string) string {
|
||||
switch name {
|
||||
case "web search":
|
||||
return "Web Search"
|
||||
case "web fetch":
|
||||
return "Web Fetch"
|
||||
default:
|
||||
return name
|
||||
}
|
||||
}
|
||||
|
||||
func writeUsageLimit(out io.Writer, name string, limit api.UsageLimit) error {
|
||||
if limit.Usage == 0 && len(limit.Models) == 0 {
|
||||
return nil
|
||||
}
|
||||
|
||||
fmt.Fprintln(out)
|
||||
fmt.Fprintln(out, usageHeading(out, name))
|
||||
table := tabwriter.NewWriter(out, 0, 4, 2, ' ', 0)
|
||||
fmt.Fprintf(table, " Used\t%.1f%%\n", limit.Usage*100)
|
||||
if len(limit.Models) > 0 {
|
||||
fmt.Fprintln(table, " Model\tRequests")
|
||||
}
|
||||
for _, m := range limit.Models {
|
||||
fmt.Fprintf(table, " %s\t%d\n", usageModelName(m.Name), m.RequestCount)
|
||||
}
|
||||
return table.Flush()
|
||||
}
|
||||
|
||||
func PushHandler(cmd *cobra.Command, args []string) error {
|
||||
client, err := api.ClientFromEnvironment()
|
||||
if err != nil {
|
||||
@@ -2450,6 +2562,14 @@ func NewCLI() *cobra.Command {
|
||||
RunE: SignoutHandler,
|
||||
}
|
||||
|
||||
usageCmd := &cobra.Command{
|
||||
Use: "usage",
|
||||
Short: "Show your ollama.com usage",
|
||||
Args: cobra.ExactArgs(0),
|
||||
PreRunE: checkServerHeartbeat,
|
||||
RunE: UsageHandler,
|
||||
}
|
||||
|
||||
listCmd := &cobra.Command{
|
||||
Use: "list",
|
||||
Aliases: []string{"ls"},
|
||||
@@ -2514,6 +2634,7 @@ func NewCLI() *cobra.Command {
|
||||
stopCmd,
|
||||
pullCmd,
|
||||
pushCmd,
|
||||
usageCmd,
|
||||
listCmd,
|
||||
psCmd,
|
||||
copyCmd,
|
||||
@@ -2566,6 +2687,7 @@ func NewCLI() *cobra.Command {
|
||||
loginCmd,
|
||||
signoutCmd,
|
||||
logoutCmd,
|
||||
usageCmd,
|
||||
listCmd,
|
||||
psCmd,
|
||||
copyCmd,
|
||||
|
||||
186
cmd/cmd_test.go
186
cmd/cmd_test.go
@@ -12,6 +12,7 @@ import (
|
||||
"path/filepath"
|
||||
"reflect"
|
||||
"strings"
|
||||
"sync/atomic"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
@@ -1398,6 +1399,191 @@ func TestListHandler(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestUsageHandler(t *testing.T) {
|
||||
startsAt := time.Date(2026, time.June, 29, 0, 0, 0, 0, time.UTC)
|
||||
endsAt := time.Date(2026, time.July, 27, 0, 0, 0, 0, time.UTC)
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
cloudDisabled bool
|
||||
statusCode int
|
||||
response any
|
||||
want string
|
||||
wantErr string
|
||||
}{
|
||||
{
|
||||
name: "team activity",
|
||||
statusCode: http.StatusOK,
|
||||
response: api.UsageResponse{
|
||||
Activity: api.UsageActivity{
|
||||
Cost: "12.34000",
|
||||
Period: api.UsagePeriod{
|
||||
Type: "last_4_weeks",
|
||||
StartingAt: startsAt,
|
||||
EndingAt: endsAt,
|
||||
},
|
||||
Models: []api.UsageActivityModel{
|
||||
{Name: "gpt-oss:120b", RequestCount: 42, Cost: "10.25000"},
|
||||
{Name: "qwen3-coder:480b", RequestCount: 7, Cost: "2.09000"},
|
||||
},
|
||||
},
|
||||
Limits: api.UsageLimits{
|
||||
Session: api.UsageLimit{Models: []api.UsageLimitModel{}},
|
||||
Weekly: api.UsageLimit{Models: []api.UsageLimitModel{}},
|
||||
},
|
||||
},
|
||||
want: "Usage\n" +
|
||||
" Period 2026-06-29 to 2026-07-27\n" +
|
||||
" Spend $12.34000\n\n" +
|
||||
"Activity\n" +
|
||||
" Model Requests Spend\n" +
|
||||
" gpt-oss:120b 42 $10.25000\n" +
|
||||
" qwen3-coder:480b 7 $2.09000\n",
|
||||
},
|
||||
{
|
||||
name: "personal limits",
|
||||
statusCode: http.StatusOK,
|
||||
response: api.UsageResponse{
|
||||
Activity: api.UsageActivity{
|
||||
Cost: "0.00000",
|
||||
Period: api.UsagePeriod{Type: "last_4_weeks", StartingAt: startsAt, EndingAt: endsAt},
|
||||
Models: []api.UsageActivityModel{{Name: "web fetch", RequestCount: 1, Cost: "0.00000"}},
|
||||
},
|
||||
Limits: api.UsageLimits{
|
||||
Session: api.UsageLimit{Usage: 0.006, Models: []api.UsageLimitModel{{Name: "qwen3-coder:480b", RequestCount: 2}, {Name: "web search", RequestCount: 1}}},
|
||||
Weekly: api.UsageLimit{Usage: 0.002, Models: []api.UsageLimitModel{{Name: "qwen3-coder:480b", RequestCount: 2}, {Name: "web search", RequestCount: 1}}},
|
||||
},
|
||||
},
|
||||
want: "Usage\n" +
|
||||
" Period 2026-06-29 to 2026-07-27\n" +
|
||||
" Spend $0.00000\n\n" +
|
||||
"Session\n" +
|
||||
" Used 0.6%\n" +
|
||||
" Model Requests\n" +
|
||||
" qwen3-coder:480b 2\n" +
|
||||
" Web Search 1\n\n" +
|
||||
"Weekly\n" +
|
||||
" Used 0.2%\n" +
|
||||
" Model Requests\n" +
|
||||
" qwen3-coder:480b 2\n" +
|
||||
" Web Search 1\n",
|
||||
},
|
||||
{
|
||||
name: "no usage",
|
||||
statusCode: http.StatusOK,
|
||||
response: api.UsageResponse{
|
||||
Activity: api.UsageActivity{
|
||||
Cost: "0.00000",
|
||||
Period: api.UsagePeriod{Type: "last_4_weeks", StartingAt: startsAt, EndingAt: endsAt},
|
||||
Models: []api.UsageActivityModel{},
|
||||
},
|
||||
Limits: api.UsageLimits{
|
||||
Session: api.UsageLimit{Models: []api.UsageLimitModel{}},
|
||||
Weekly: api.UsageLimit{Models: []api.UsageLimitModel{}},
|
||||
},
|
||||
},
|
||||
want: "Usage\n" +
|
||||
" Period 2026-06-29 to 2026-07-27\n" +
|
||||
" Spend $0.00000\n\n" +
|
||||
"No usage recorded for this period.\n",
|
||||
},
|
||||
{
|
||||
name: "cloud disabled",
|
||||
cloudDisabled: true,
|
||||
want: "Ollama Cloud is disabled; usage is unavailable\n",
|
||||
},
|
||||
{
|
||||
name: "not signed in",
|
||||
statusCode: http.StatusUnauthorized,
|
||||
response: map[string]string{"error": "unauthorized"},
|
||||
want: "You need to be signed in to Ollama to view usage.\n\n",
|
||||
},
|
||||
{
|
||||
name: "suspended account",
|
||||
statusCode: http.StatusForbidden,
|
||||
response: map[string]string{"error": "account suspended"},
|
||||
wantErr: "account suspended",
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
var usageRequests atomic.Int32
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
if r.Method != http.MethodGet {
|
||||
t.Errorf("unexpected request to %s %s", r.Method, r.URL.Path)
|
||||
http.NotFound(w, r)
|
||||
return
|
||||
}
|
||||
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
switch r.URL.Path {
|
||||
case "/api/status":
|
||||
if err := json.NewEncoder(w).Encode(api.StatusResponse{Cloud: api.CloudStatus{Disabled: tt.cloudDisabled}}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
case "/api/usage":
|
||||
usageRequests.Add(1)
|
||||
w.WriteHeader(tt.statusCode)
|
||||
if err := json.NewEncoder(w).Encode(tt.response); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
default:
|
||||
t.Errorf("unexpected request to %s %s", r.Method, r.URL.Path)
|
||||
http.NotFound(w, r)
|
||||
}
|
||||
}))
|
||||
defer server.Close()
|
||||
|
||||
t.Setenv("OLLAMA_HOST", server.URL)
|
||||
|
||||
cmd := &cobra.Command{}
|
||||
cmd.SetContext(t.Context())
|
||||
var out bytes.Buffer
|
||||
cmd.SetOut(&out)
|
||||
|
||||
err := UsageHandler(cmd, nil)
|
||||
if tt.wantErr != "" {
|
||||
if err == nil || !strings.Contains(err.Error(), tt.wantErr) {
|
||||
t.Fatalf("error = %v, want %q", err, tt.wantErr)
|
||||
}
|
||||
} else if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if got := out.String(); got != tt.want {
|
||||
t.Errorf("unexpected output (-want +got):\n%s", cmp.Diff(tt.want, got))
|
||||
}
|
||||
|
||||
wantUsageRequests := int32(1)
|
||||
if tt.cloudDisabled {
|
||||
wantUsageRequests = 0
|
||||
}
|
||||
if got := usageRequests.Load(); got != wantUsageRequests {
|
||||
t.Errorf("usage requests = %d, want %d", got, wantUsageRequests)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestUsageModelName(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
want string
|
||||
}{
|
||||
{name: "web search", want: "Web Search"},
|
||||
{name: "web fetch", want: "Web Fetch"},
|
||||
{name: "qwen3-coder:480b", want: "qwen3-coder:480b"},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
if got := usageModelName(tt.name); got != tt.want {
|
||||
t.Errorf("usageModelName(%q) = %q, want %q", tt.name, got, tt.want)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestCreateHandler(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
|
||||
@@ -56,11 +56,17 @@ import (
|
||||
|
||||
const signinURLStr = "https://ollama.com/connect?name=%s&key=%s"
|
||||
|
||||
var (
|
||||
accountBaseURL = "https://ollama.com"
|
||||
usageBaseURL = "https://ollama.com"
|
||||
)
|
||||
|
||||
const (
|
||||
cloudErrRemoteInferenceUnavailable = "remote model is unavailable"
|
||||
cloudErrRemoteModelDetailsUnavailable = "remote model details are unavailable"
|
||||
cloudErrWebSearchUnavailable = "web search is unavailable"
|
||||
cloudErrWebFetchUnavailable = "web fetch is unavailable"
|
||||
cloudErrUsageUnavailable = "usage is unavailable"
|
||||
copilotChatUserAgentPrefix = "GitHubCopilotChat/"
|
||||
)
|
||||
|
||||
@@ -1882,7 +1888,7 @@ func (s *Server) GenerateRoutes() (http.Handler, error) {
|
||||
r.DELETE("/api/delete", s.DeleteHandler)
|
||||
|
||||
r.POST("/api/me", s.WhoamiHandler)
|
||||
|
||||
r.GET("/api/usage", s.UsageHandler)
|
||||
r.POST("/api/signout", s.SignoutHandler)
|
||||
// deprecated
|
||||
r.DELETE("/api/user/keys/:encodedKey", s.SignoutHandler)
|
||||
@@ -2174,7 +2180,7 @@ func (s *Server) webExperimentalProxyHandler(c *gin.Context, proxyPath, disabled
|
||||
|
||||
func (s *Server) WhoamiHandler(c *gin.Context) {
|
||||
// todo allow other hosts
|
||||
u, err := url.Parse("https://ollama.com")
|
||||
u, err := url.Parse(accountBaseURL)
|
||||
if err != nil {
|
||||
slog.Error(err.Error())
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": "URL parse error"})
|
||||
@@ -2226,6 +2232,53 @@ func (s *Server) WhoamiHandler(c *gin.Context) {
|
||||
c.JSON(http.StatusOK, user)
|
||||
}
|
||||
|
||||
func (s *Server) UsageHandler(c *gin.Context) {
|
||||
if internalcloud.Disabled() {
|
||||
c.JSON(http.StatusForbidden, gin.H{"error": internalcloud.DisabledError(cloudErrUsageUnavailable)})
|
||||
return
|
||||
}
|
||||
|
||||
// todo allow other hosts
|
||||
u, err := url.Parse(usageBaseURL)
|
||||
if err != nil {
|
||||
slog.Error(err.Error())
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": "URL parse error"})
|
||||
return
|
||||
}
|
||||
|
||||
client := api.NewClient(u, http.DefaultClient)
|
||||
usage, err := client.Usage(c)
|
||||
if err != nil {
|
||||
var authErr api.AuthorizationError
|
||||
if errors.As(err, &authErr) && authErr.StatusCode == http.StatusUnauthorized {
|
||||
sURL := authErr.SigninURL
|
||||
if sURL == "" {
|
||||
var sErr error
|
||||
sURL, sErr = signinURL()
|
||||
if sErr != nil {
|
||||
slog.Error(sErr.Error())
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": "error getting authorization details"})
|
||||
return
|
||||
}
|
||||
}
|
||||
c.JSON(http.StatusUnauthorized, gin.H{"error": "unauthorized", "signin_url": sURL})
|
||||
return
|
||||
}
|
||||
|
||||
var statusErr api.StatusError
|
||||
if errors.As(err, &statusErr) && statusErr.StatusCode == http.StatusForbidden {
|
||||
c.JSON(http.StatusForbidden, gin.H{"error": statusErr.ErrorMessage})
|
||||
return
|
||||
}
|
||||
|
||||
slog.Error(err.Error())
|
||||
c.JSON(http.StatusServiceUnavailable, gin.H{"error": "account unavailable"})
|
||||
return
|
||||
}
|
||||
|
||||
c.JSON(http.StatusOK, usage)
|
||||
}
|
||||
|
||||
func (s *Server) SignoutHandler(c *gin.Context) {
|
||||
pubKey, err := auth.GetPublicKey()
|
||||
if err != nil {
|
||||
@@ -2237,7 +2290,7 @@ func (s *Server) SignoutHandler(c *gin.Context) {
|
||||
encKey := base64.RawURLEncoding.EncodeToString([]byte(pubKey))
|
||||
|
||||
// todo allow other hosts
|
||||
u, err := url.Parse("https://ollama.com")
|
||||
u, err := url.Parse(accountBaseURL)
|
||||
if err != nil {
|
||||
slog.Error(err.Error())
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": "URL parse error"})
|
||||
|
||||
213
server/routes_usage_test.go
Normal file
213
server/routes_usage_test.go
Normal file
@@ -0,0 +1,213 @@
|
||||
package server
|
||||
|
||||
import (
|
||||
"crypto/ed25519"
|
||||
"crypto/rand"
|
||||
"encoding/json"
|
||||
"encoding/pem"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"sync/atomic"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/ollama/ollama/api"
|
||||
internalcloud "github.com/ollama/ollama/internal/cloud"
|
||||
"golang.org/x/crypto/ssh"
|
||||
)
|
||||
|
||||
func TestUsageHandler(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
setTestHome(t, t.TempDir())
|
||||
t.Setenv("OLLAMA_NO_CLOUD", "")
|
||||
t.Setenv("OLLAMA_AUTH", "1")
|
||||
writeTestPrivateKey(t)
|
||||
|
||||
startsAt := time.Date(2026, time.June, 29, 0, 0, 0, 0, time.UTC)
|
||||
endsAt := time.Date(2026, time.July, 27, 0, 0, 0, 0, time.UTC)
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
upstreamStatus int
|
||||
upstreamBody any
|
||||
wantStatus int
|
||||
wantError string
|
||||
wantSigninURL string
|
||||
}{
|
||||
{
|
||||
name: "success",
|
||||
upstreamStatus: http.StatusOK,
|
||||
upstreamBody: api.UsageResponse{
|
||||
Activity: api.UsageActivity{
|
||||
Cost: "0.00709",
|
||||
Period: api.UsagePeriod{
|
||||
Type: "last_4_weeks",
|
||||
StartingAt: startsAt,
|
||||
EndingAt: endsAt,
|
||||
},
|
||||
Models: []api.UsageActivityModel{{Name: "qwen3-coder:480b", RequestCount: 1, Cost: "0.00709"}},
|
||||
},
|
||||
Limits: api.UsageLimits{
|
||||
Session: api.UsageLimit{Usage: 0.006, Models: []api.UsageLimitModel{{Name: "qwen3-coder:480b", RequestCount: 2}}},
|
||||
Weekly: api.UsageLimit{Usage: 0.002, Models: []api.UsageLimitModel{{Name: "web search", RequestCount: 1}}},
|
||||
},
|
||||
},
|
||||
wantStatus: http.StatusOK,
|
||||
},
|
||||
{
|
||||
name: "not signed in",
|
||||
upstreamStatus: http.StatusUnauthorized,
|
||||
upstreamBody: map[string]string{"error": "unauthorized", "signin_url": "https://ollama.com/connect/test"},
|
||||
wantStatus: http.StatusUnauthorized,
|
||||
wantError: "unauthorized",
|
||||
wantSigninURL: "https://ollama.com/connect/test",
|
||||
},
|
||||
{
|
||||
name: "suspended account",
|
||||
upstreamStatus: http.StatusForbidden,
|
||||
upstreamBody: map[string]string{"error": "account suspended"},
|
||||
wantStatus: http.StatusForbidden,
|
||||
wantError: "account suspended",
|
||||
},
|
||||
{
|
||||
name: "upstream unavailable",
|
||||
upstreamStatus: http.StatusInternalServerError,
|
||||
upstreamBody: map[string]string{"error": "internal error"},
|
||||
wantStatus: http.StatusServiceUnavailable,
|
||||
wantError: "account unavailable",
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
if r.Method != http.MethodGet || r.URL.Path != "/api/usage" {
|
||||
t.Errorf("unexpected upstream request to %s %s", r.Method, r.URL.Path)
|
||||
http.NotFound(w, r)
|
||||
return
|
||||
}
|
||||
if got := r.Header.Get("Authorization"); got == "" {
|
||||
t.Error("upstream request is missing Authorization header")
|
||||
}
|
||||
if got := r.URL.Query().Get("ts"); got == "" {
|
||||
t.Error("upstream request is missing authentication timestamp")
|
||||
}
|
||||
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
w.WriteHeader(tt.upstreamStatus)
|
||||
if err := json.NewEncoder(w).Encode(tt.upstreamBody); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}))
|
||||
defer upstream.Close()
|
||||
|
||||
originalBaseURL := usageBaseURL
|
||||
usageBaseURL = upstream.URL
|
||||
t.Cleanup(func() { usageBaseURL = originalBaseURL })
|
||||
|
||||
server := &Server{}
|
||||
router, err := server.GenerateRoutes()
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
request := httptest.NewRequest(http.MethodGet, "/api/usage", nil)
|
||||
response := httptest.NewRecorder()
|
||||
router.ServeHTTP(response, request)
|
||||
|
||||
if response.Code != tt.wantStatus {
|
||||
t.Fatalf("status = %d, want %d (%s)", response.Code, tt.wantStatus, response.Body.String())
|
||||
}
|
||||
|
||||
if tt.wantStatus == http.StatusOK {
|
||||
var usage api.UsageResponse
|
||||
if err := json.Unmarshal(response.Body.Bytes(), &usage); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if usage.Activity.Cost != "0.00709" || len(usage.Activity.Models) != 1 || usage.Activity.Models[0].Name != "qwen3-coder:480b" || usage.Limits.Session.Usage != 0.006 || usage.Limits.Weekly.Usage != 0.002 {
|
||||
t.Errorf("usage = %#v, want upstream usage", usage)
|
||||
}
|
||||
return
|
||||
}
|
||||
|
||||
var body struct {
|
||||
Error string `json:"error"`
|
||||
SigninURL string `json:"signin_url"`
|
||||
}
|
||||
if err := json.Unmarshal(response.Body.Bytes(), &body); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if body.Error != tt.wantError {
|
||||
t.Errorf("error = %q, want %q", body.Error, tt.wantError)
|
||||
}
|
||||
if body.SigninURL != tt.wantSigninURL {
|
||||
t.Errorf("signin URL = %q, want %q", body.SigninURL, tt.wantSigninURL)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func writeTestPrivateKey(t *testing.T) {
|
||||
t.Helper()
|
||||
|
||||
_, privateKey, err := ed25519.GenerateKey(rand.Reader)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
privateKeyBlock, err := ssh.MarshalPrivateKey(privateKey, "")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
home, err := os.UserHomeDir()
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
keyPath := filepath.Join(home, ".ollama", "id_ed25519")
|
||||
if err := os.MkdirAll(filepath.Dir(keyPath), 0o755); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := os.WriteFile(keyPath, pem.EncodeToMemory(privateKeyBlock), 0o600); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestUsageHandlerCloudDisabled(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
setTestHome(t, t.TempDir())
|
||||
t.Setenv("OLLAMA_NO_CLOUD", "1")
|
||||
|
||||
var upstreamRequests atomic.Int32
|
||||
upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
upstreamRequests.Add(1)
|
||||
http.Error(w, "unexpected request", http.StatusInternalServerError)
|
||||
}))
|
||||
defer upstream.Close()
|
||||
|
||||
originalBaseURL := usageBaseURL
|
||||
usageBaseURL = upstream.URL
|
||||
t.Cleanup(func() { usageBaseURL = originalBaseURL })
|
||||
|
||||
server := &Server{}
|
||||
router, err := server.GenerateRoutes()
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
request := httptest.NewRequest(http.MethodGet, "/api/usage", nil)
|
||||
response := httptest.NewRecorder()
|
||||
router.ServeHTTP(response, request)
|
||||
|
||||
if response.Code != http.StatusForbidden {
|
||||
t.Fatalf("status = %d, want %d (%s)", response.Code, http.StatusForbidden, response.Body.String())
|
||||
}
|
||||
if got, want := response.Body.String(), `{"error":"`+internalcloud.DisabledError(cloudErrUsageUnavailable)+`"}`; got != want {
|
||||
t.Fatalf("body = %q, want %q", got, want)
|
||||
}
|
||||
if got := upstreamRequests.Load(); got != 0 {
|
||||
t.Fatalf("upstream requests = %d, want 0", got)
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user