Compare commits

...

2 Commits

Author SHA1 Message Date
ParthSareen
90dd5b3a70 feat(cli): add usage command 2026-07-23 12:58:14 -07:00
ParthSareen
4855c61358 feat(api): add account usage endpoint 2026-07-23 12:58:08 -07:00
7 changed files with 695 additions and 3 deletions

View File

@@ -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
}

View File

@@ -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 {

View File

@@ -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"`

View File

@@ -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,

View File

@@ -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

View File

@@ -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
View 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)
}
}