From dcb21f52da730d3e17518cc545ede9713524e9b8 Mon Sep 17 00:00:00 2001 From: Vasil Date: Sat, 24 Oct 2026 02:20:01 +0200 Subject: [PATCH 1/2] tokenmeter CLI and Go client usage summary, top, by-meta and export with totals in USD, fetching windows the server caps in parts; token new and hash for TOKENMETER_TOKENS; event send; health, now with the version in /healthz. Contexts keep the token in the environment, a 0600 file or the system keychain. pkg/client wraps the gRPC API. Released as archives with signed checksums, the coyl/tokenmeter-cli image and a Homebrew formula; an end-to-end job runs the CLI against a built service with and without ClickHouse. --- .github/workflows/docker-publish.yml | 26 +- .github/workflows/release.yml | 31 +- .github/workflows/test.yml | 24 +- .goreleaser.yaml | 63 ++++ CHANGELOG.md | 10 + Dockerfile.cli | 17 + README.md | 48 ++- SECURITY.md | 10 + cmd/server/main.go | 1 + cmd/tokenmeter/conn.go | 153 +++++++++ cmd/tokenmeter/context_cmd.go | 232 ++++++++++++++ cmd/tokenmeter/contexts.go | 234 ++++++++++++++ cmd/tokenmeter/errors.go | 88 ++++++ cmd/tokenmeter/event_cmd.go | 97 ++++++ cmd/tokenmeter/http_cmd.go | 89 ++++++ cmd/tokenmeter/main.go | 146 +++++++++ cmd/tokenmeter/main_test.go | 255 +++++++++++++++ cmd/tokenmeter/output.go | 144 +++++++++ cmd/tokenmeter/token_cmd.go | 104 ++++++ cmd/tokenmeter/usage_cmd.go | 454 +++++++++++++++++++++++++++ go.mod | 8 +- go.sum | 15 +- internal/server/http.go | 5 +- pkg/client/client.go | 203 ++++++++++++ pkg/client/client_test.go | 118 +++++++ scripts/cli-e2e.sh | 119 +++++++ 26 files changed, 2679 insertions(+), 15 deletions(-) create mode 100644 .goreleaser.yaml create mode 100644 Dockerfile.cli create mode 100644 cmd/tokenmeter/conn.go create mode 100644 cmd/tokenmeter/context_cmd.go create mode 100644 cmd/tokenmeter/contexts.go create mode 100644 cmd/tokenmeter/errors.go create mode 100644 cmd/tokenmeter/event_cmd.go create mode 100644 cmd/tokenmeter/http_cmd.go create mode 100644 cmd/tokenmeter/main.go create mode 100644 cmd/tokenmeter/main_test.go create mode 100644 cmd/tokenmeter/output.go create mode 100644 cmd/tokenmeter/token_cmd.go create mode 100644 cmd/tokenmeter/usage_cmd.go create mode 100644 pkg/client/client.go create mode 100644 pkg/client/client_test.go create mode 100755 scripts/cli-e2e.sh diff --git a/.github/workflows/docker-publish.yml b/.github/workflows/docker-publish.yml index 0324b98..1960b1e 100644 --- a/.github/workflows/docker-publish.yml +++ b/.github/workflows/docker-publish.yml @@ -1,7 +1,8 @@ name: Publish Docker image -# Every vX.Y.Z release becomes a signed multi-arch image coyl/light-tokenmeter:X.Y.Z -# with SBOM and provenance; X.Y follows the newest release of its line, latest the newest of all. +# Every vX.Y.Z release becomes signed multi-arch images coyl/light-tokenmeter:X.Y.Z (the +# service) and coyl/tokenmeter-cli:X.Y.Z (the CLI) with SBOM and provenance; X.Y follows the +# newest release of its line, latest the newest of all. # Called by release.yml. Needs the secrets DOCKER_USERNAME and DOCKER_PASSWORD. on: workflow_call: @@ -18,6 +19,18 @@ permissions: jobs: publish: runs-on: ubuntu-latest + strategy: + # Each image signs itself; one leg failing must not cancel another + # between its push and its signature. + fail-fast: false + matrix: + include: + - image: coyl/light-tokenmeter + dockerfile: Dockerfile + scope: service + - image: coyl/tokenmeter-cli + dockerfile: Dockerfile.cli + scope: cli steps: - uses: actions/checkout@11d5960a326750d5838078e36cf38b85af677262 # v4.4.0 with: @@ -48,7 +61,7 @@ jobs: id: meta uses: docker/metadata-action@c299e40c65443455700f0fdfc63efafe5b349051 # v5.10.0 with: - images: coyl/light-tokenmeter + images: ${{ matrix.image }} tags: | type=semver,pattern={{version}} type=semver,pattern={{major}}.{{minor}},enable=${{ steps.newest.outputs.line }} @@ -60,13 +73,14 @@ jobs: uses: docker/build-push-action@10e90e3645eae34f1e60eeb005ba3a3d33f178e8 # v6.19.2 with: context: . + file: ${{ matrix.dockerfile }} platforms: linux/amd64,linux/arm64 push: true tags: ${{ steps.meta.outputs.tags }} labels: ${{ steps.meta.outputs.labels }} build-args: VERSION=${{ steps.meta.outputs.version }} - cache-from: type=gha - cache-to: type=gha,mode=max + cache-from: type=gha,scope=${{ matrix.scope }} + cache-to: type=gha,mode=max,scope=${{ matrix.scope }} provenance: mode=max sbom: true @@ -80,7 +94,7 @@ jobs: - name: Sign the image run: | for i in 1 2 3; do - cosign sign --yes --recursive "coyl/light-tokenmeter@${{ steps.build.outputs.digest }}" && exit 0 + cosign sign --yes --recursive "${{ matrix.image }}@${{ steps.build.outputs.digest }}" && exit 0 [ "$i" -lt 3 ] && sleep $((i * 15)) done exit 1 diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml index 48c8ecb..4a32b33 100644 --- a/.github/workflows/release.yml +++ b/.github/workflows/release.yml @@ -47,6 +47,35 @@ jobs: gh release create "$GITHUB_REF_NAME" --draft --title "light-tokenmeter $v" --notes-file notes.md --verify-tag fi + # The CLI's archives and signed checksums, onto the draft. + binaries: + needs: draft + runs-on: ubuntu-latest + # Every action here is pinned by commit: the job can sign under the + # release identity and write to the release. + permissions: + contents: write + id-token: write # keyless cosign signature of checksums.txt + steps: + - uses: actions/checkout@11d5960a326750d5838078e36cf38b85af677262 # v4.4.0 + with: + fetch-depth: 0 + - uses: actions/setup-go@40f1582b2485089dde7abd97c1529aa768e1baff # v5.6.0 + with: + go-version: '1.27.3' + cache: true + - uses: sigstore/cosign-installer@6f9f17788090df1f26f669e9d70d6ae9567deba6 # v4.1.2 + with: + cosign-release: v3.1.3 + - uses: anchore/sbom-action/download-syft@3ad7283483fc7af8ff2b4ea19663c2d5ca935e26 # v0.24.2, syft for the SBOMs + - uses: goreleaser/goreleaser-action@e435ccd777264be153ace6237001ef4d979d3a7a # v6.4.0 + with: + version: '~> v2' + args: release --clean + env: + GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} + HOMEBREW_TAP_TOKEN: ${{ secrets.HOMEBREW_TAP_TOKEN }} + images: needs: draft permissions: @@ -58,7 +87,7 @@ jobs: DOCKER_PASSWORD: ${{ secrets.DOCKER_PASSWORD }} publish: - needs: images + needs: [images, binaries] runs-on: ubuntu-latest permissions: contents: write diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index e5e52b0..32d4a96 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -79,6 +79,18 @@ jobs: CUR: ${{ github.ref_type == 'tag' && github.ref_name || '' }} # a release tag is not its own predecessor run: ./scripts/upgrade-test.sh + # The tokenmeter CLI against a built service, with ClickHouse and without. + cli-e2e: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + - uses: actions/setup-go@v5 + with: + go-version: '1.27.3' + cache: true + - name: CLI end to end + run: ./scripts/cli-e2e.sh + docker: runs-on: ubuntu-latest steps: @@ -89,5 +101,13 @@ jobs: with: context: . push: false - cache-from: type=gha - cache-to: type=gha,mode=max + cache-from: type=gha,scope=service + cache-to: type=gha,mode=max,scope=service + - name: Build CLI image (no push) + uses: docker/build-push-action@v6 + with: + context: . + file: Dockerfile.cli + push: false + cache-from: type=gha,scope=cli + cache-to: type=gha,mode=max,scope=cli diff --git a/.goreleaser.yaml b/.goreleaser.yaml new file mode 100644 index 0000000..d84ad9c --- /dev/null +++ b/.goreleaser.yaml @@ -0,0 +1,63 @@ +# The tokenmeter CLI on the GitHub release that release.yml drafted: archives, +# a signed checksums.txt and SBOMs. The service itself ships as an image. +# The Homebrew tap switches on only when HOMEBREW_TAP_TOKEN exists. +version: 2 +project_name: tokenmeter + +builds: + - id: cli + main: ./cmd/tokenmeter + binary: tokenmeter + env: [CGO_ENABLED=0] + goos: [darwin, linux] + goarch: [amd64, arm64] + flags: [-trimpath] + ldflags: ["-s -w -X main.version={{ .Version }}"] + +archives: + - id: cli + ids: [cli] + formats: [tar.gz] + name_template: "tokenmeter_{{ .Version }}_{{ .Os }}_{{ .Arch }}" + files: [LICENSE, NOTICE] + +checksum: + name_template: checksums.txt + +# One SPDX SBOM per archive, built by syft. +sboms: + - artifacts: archive + +# A keyless signature of checksums.txt covers every archive listed in it. +signs: + - cmd: cosign + artifacts: checksum + signature: "${artifact}.sigstore.json" + args: ["sign-blob", "--yes", "--bundle=${signature}", "${artifact}"] + +# The release, its notes and its publishing belong to release.yml. +release: + draft: true + use_existing_draft: true + mode: keep-existing + replace_existing_artifacts: true + +changelog: + disable: true + +brews: + - name: tokenmeter + ids: [cli] + homepage: https://github.com/ai-process/light-tokenmeter + description: Usage reports and tokens for light-tokenmeter + license: Apache-2.0 + repository: + owner: ai-process + name: homebrew-tap + token: '{{ if isEnvSet "HOMEBREW_TAP_TOKEN" }}{{ .Env.HOMEBREW_TAP_TOKEN }}{{ end }}' + skip_upload: '{{ if isEnvSet "HOMEBREW_TAP_TOKEN" }}false{{ else }}true{{ end }}' + install: | + bin.install "tokenmeter" + generate_completions_from_executable(bin/"tokenmeter", "completion") + test: | + system "#{bin}/tokenmeter", "version" diff --git a/CHANGELOG.md b/CHANGELOG.md index 5edd1fa..77eac32 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -8,6 +8,16 @@ All notable changes to light-tokenmeter are listed here. The format follows ## [Unreleased] +### Added +- The `tokenmeter` CLI: `usage summary|top|by-meta|export` with totals in USD and windows that + the server caps fetched in parts, `token new|hash` for `TOKENMETER_TOKENS`, `event send`, + `health`, contexts with the token from an environment variable, a 0600 file or the system + keychain, `-o json|csv`, exit codes by error class and completion for bash, zsh and fish. + Released as archives with a signed `checksums.txt`, the `coyl/tokenmeter-cli` image and a + Homebrew formula (published once the tap token is set). +- `pkg/client`, a Go client for the gRPC API. +- `/healthz` reports the service's `version`. + ## [1.0.0] - 2026-10-23 The first stable release, together with Keystone. The Model Switchboard. (KTMS) 1.0 and diff --git a/Dockerfile.cli b/Dockerfile.cli new file mode 100644 index 0000000..aa18f2f --- /dev/null +++ b/Dockerfile.cli @@ -0,0 +1,17 @@ +FROM golang:1.27.3-alpine3.23 AS builder +RUN apk add --no-cache 'ca-certificates>=20261020-r0' +ARG TARGETOS=linux +ARG TARGETARCH +ARG VERSION=dev +WORKDIR /app +COPY go.mod go.sum ./ +RUN --mount=type=cache,target=/go/pkg/mod go mod download +COPY . . +RUN --mount=type=cache,target=/go/pkg/mod --mount=type=cache,target=/root/.cache/go-build \ + CGO_ENABLED=0 GOOS=${TARGETOS} GOARCH=${TARGETARCH} go build -trimpath -ldflags="-s -w -X main.version=${VERSION}" -o /out/tokenmeter ./cmd/tokenmeter + +FROM gcr.io/distroless/static-debian12:nonroot +COPY --from=builder /out/tokenmeter /tokenmeter +COPY --from=builder /etc/ssl/certs/ca-certificates.crt /etc/ssl/certs/ca-certificates.crt +USER 65532:65532 +ENTRYPOINT ["/tokenmeter"] diff --git a/README.md b/README.md index b992245..c8bb9bc 100644 --- a/README.md +++ b/README.md @@ -26,6 +26,7 @@ who is worthy. - [Installation](#installation) - [Sending events](#sending-events) - [Finding the greedy and the worthy](#finding-the-greedy-and-the-worthy) +- [Command-line tool](#command-line-tool) - [API reference](#api-reference) - [API stability](#api-stability) - [Configuration](#configuration) @@ -214,6 +215,45 @@ grpcurl -plaintext -d '{"meta_key":"job_id","meta_value":"job-42","group_by_meta gRPC reflection is on, so `grpcurl` needs no `.proto` files (with auth on, pass a token with `-H 'authorization: Bearer …'`). +## Command-line tool + +`tokenmeter` reads the reports without hand-built query strings or dividing by 1e9, makes tokens +without `openssl`, and sends a test event. It talks gRPC to `:9090`; `health` uses HTTP `:8080`. + +```bash +# a release archive (darwin/linux, amd64/arm64, signed checksums), or: +go install github.com/ai-process/light-tokenmeter/cmd/tokenmeter@latest +docker run --rm -e TOKENMETER_ADDR -e TOKENMETER_TOKEN coyl/tokenmeter-cli:latest usage summary +``` + +```bash +export TOKENMETER_ADDR=localhost:9090 TOKENMETER_TOKEN=… # a read token for reports +tokenmeter usage summary --last 7d --group-by project # totals in USD; --daily per day +tokenmeter usage summary --from 2026-09-01 --to 2026-10-01 --model gpt-4o --group-by user +tokenmeter usage top --by cost --dimension action --limit 5 +tokenmeter usage by-meta --key job_id --value job-42 --group-by task_id +tokenmeter usage export --last 30d --format csv --out usage.csv # daily rows, every dimension +tokenmeter token new --name ktms --scope write # the token, then its entry +tokenmeter event send --project p1 --user u1 --action check --vendor openai --model gpt-4o-mini +tokenmeter health +tokenmeter completion zsh > "${fpath[1]}/_tokenmeter" # also bash and fish +``` + +- Windows are `--last` (`24h`, `7d`, `4w`; 7 days by default) or `--from`/`--to` in RFC 3339 or + `YYYY-MM-DD`, all UTC. A window the server caps at 10,000 rows (`truncated`) is fetched in + smaller parts and added up, so totals are complete; only an hour that is still capped is + reported as incomplete. +- `-o json` keeps the API's fields (`cost_nanousd` as a string), `-o csv` gives exact dollars. +- Contexts work as in `ktms`: `tokenmeter context set prod --addr tm.example.com:443 --tls + --token keychain`, then `tokenmeter context store-token prod` reads the token without echo. + They live in `~/.config/tokenmeter/config.yaml` (`TOKENMETER_CONFIG`); a token source is + `env:VAR`, `file:PATH` or `keychain`. Flags win over `TOKENMETER_ADDR`, `TOKENMETER_TLS`, + `TOKENMETER_HTTP` and `TOKENMETER_TOKEN`, which win over the context. +- Exit codes: 0 done, 1 failed, 2 bad usage, 4 not authenticated or not allowed, 5 not found, + 6 rejected as invalid, 7 unavailable (retry later). +- [`pkg/client`](pkg/client) is the Go client underneath: `Record` in batches of 500, `Summary` + with the window splitting, `ByMeta`. + ## API reference ### HTTP (`:8080`) @@ -223,7 +263,7 @@ gRPC reflection is on, so `grpcurl` needs no `.proto` files (with auth on, pass | `POST` | `/v1/usage/events` | Record one event. | | `POST` | `/v1/usage/events:batch` | Record up to 500 events. | | `GET` | `/v1/usage/summary` | Aggregated usage, see above. | -| `GET` | `/healthz` | Liveness. | +| `GET` | `/healthz` | Liveness, with the service's `version`. | | `GET` | `/readyz` | Readiness (database ping). | Errors from `/v1/usage/*` are JSON `{"error": "", "code": "", "details": ""}`, @@ -266,6 +306,10 @@ releases fix bugs. Not public, and free to change in any release: log message text (other than the prefix above), the `internal/` Go packages, the outbox's lease and retry timings, and resource numbers. +The `tokenmeter` CLI (commands, flags, exit codes, `-o json|csv` output) and `pkg/client` are new +in 1.1 and Experimental until 1.2: they may still change in a minor release, with a changelog +note. + ### Compatibility with Keystone. The Model Switchboard. (KTMS) and ktms-admin | light-tokenmeter | Keystone (KTMS) | ktms-admin | Authentication | @@ -345,6 +389,8 @@ The service stores only SHA-256 hashes. Each entry is `name:scope[+scope]:sha256 are separated by commas or newlines, and in a file `#` starts a comment: ```bash +tokenmeter token new --name ktms --scope write # prints the token, then ktms:write: +# or by hand: TOKEN=$(openssl rand -hex 32) # give this to the client printf %s "$TOKEN" | openssl dgst -sha256 -r | cut -d' ' -f1 # put this in the config TOKENMETER_TOKENS="ktms:write:,console:read:" diff --git a/SECURITY.md b/SECURITY.md index d460336..858ab07 100644 --- a/SECURITY.md +++ b/SECURITY.md @@ -25,6 +25,16 @@ docker buildx imagetools inspect coyl/light-tokenmeter:${V#v} --format '{{ json docker buildx imagetools inspect coyl/light-tokenmeter:${V#v} --format '{{ json .Provenance }}' ``` +The CLI image `coyl/tokenmeter-cli` is signed the same way. The CLI's archives on the GitHub +release (`checksums.txt` is signed; each archive has an `.sbom.json`), from 1.1.0 on: + +```bash +cosign verify-blob checksums.txt --bundle checksums.txt.sigstore.json \ + --certificate-identity "https://github.com/ai-process/light-tokenmeter/.github/workflows/release.yml@refs/tags/$V" \ + --certificate-oidc-issuer https://token.actions.githubusercontent.com +sha256sum --ignore-missing -c checksums.txt +``` + Every pull request and every release runs [`govulncheck`](https://go.dev/doc/security/vuln/): a known vulnerability that light-tokenmeter's code can reach fails the build, and a release does not go out. diff --git a/cmd/server/main.go b/cmd/server/main.go index 52a50e2..bf2a5ef 100644 --- a/cmd/server/main.go +++ b/cmd/server/main.go @@ -60,6 +60,7 @@ func main() { // Before migrations, which can be long: a scrape must not see the service down. metrics.BuildInfo.WithLabelValues(version).Set(1) + server.Version = version if cfg.ClickHouseEnabled { metrics.ClickHouseEnabled.Set(1) } diff --git a/cmd/tokenmeter/conn.go b/cmd/tokenmeter/conn.go new file mode 100644 index 0000000..4bdcd7a --- /dev/null +++ b/cmd/tokenmeter/conn.go @@ -0,0 +1,153 @@ +package main + +import ( + "crypto/tls" + "errors" + "fmt" + "io" + "net" + "net/http" + "os" + "strings" + + "github.com/ai-process/light-tokenmeter/pkg/client" +) + +// target is where and how to reach light-tokenmeter. Flags win over +// environment variables, which win over the selected context. +type target struct { + context string + addr string + tls bool + tlsSet bool + http string + tokenFile string + + ctx *Context // the selected context, nil without one + ctxName string + warn io.Writer + resolved bool +} + +// resolve fills addr, tls and http from the environment and the context, on +// first use, so commands that need no service never read the contexts file. +func (t *target) resolve() error { + if t.resolved { + return nil + } + t.resolved = true + asked := firstNonEmpty(t.context, os.Getenv("TOKENMETER_CONTEXT")) + cs, err := loadContexts() + var noHome contextsUnreachable + if errors.As(err, &noHome) && asked == "" { + // No home directory (cron, systemd): flags and environment still work. + cs, err = &Contexts{}, nil + } + if err != nil { + return err + } + name := firstNonEmpty(asked, cs.Current) + if name != "" { + c, ok := cs.Contexts[name] + if !ok { + return fmt.Errorf("no context %q (tokenmeter context list)", name) + } + t.ctx, t.ctxName = c, name + } + c := t.ctx + if c == nil { + c = &Context{} + } + t.addr = firstNonEmpty(t.addr, os.Getenv("TOKENMETER_ADDR"), c.Addr, "localhost:9090") + if !t.tlsSet { + if v := os.Getenv("TOKENMETER_TLS"); v != "" { + t.tls = v == "true" || v == "1" + } else { + t.tls = c.TLS + } + } + t.http = firstNonEmpty(t.http, os.Getenv("TOKENMETER_HTTP"), c.HTTP) + if t.http == "" { + host, _, err := net.SplitHostPort(t.addr) + if err != nil { + host = t.addr + } + scheme := "http" + if t.tls { + scheme = "https" + } + t.http = scheme + "://" + net.JoinHostPort(host, "8080") + } + t.http = strings.TrimRight(t.http, "/") + return nil +} + +// token reads the bearer token from --token-file, TOKENMETER_TOKEN or the +// context, never from a flag value. Empty means none: the service may run +// without auth, and then it refuses nothing. +func (t *target) token() (string, error) { + if err := t.resolve(); err != nil { + return "", err + } + if t.tokenFile != "" { + return readKeyFile(t.tokenFile) + } + if k := os.Getenv("TOKENMETER_TOKEN"); k != "" { + return k, nil + } + if t.ctx != nil && !t.ctx.Token.empty() { + k, err := t.ctx.Token.read(t.ctxName) + if err != nil { + return "", fmt.Errorf("token of context %s: %w", t.ctxName, err) + } + return k, nil + } + return "", nil +} + +func (t *target) client() (*client.Client, error) { + tok, err := t.token() + if err != nil { + return nil, err + } + if tok != "" && !t.tls && !loopback(t.addr) { + fmt.Fprintf(t.warn, "tokenmeter: warning: sending the token to %s without TLS; pass --tls unless this network is private\n", t.addr) + } + return client.New(client.Options{Addr: t.addr, Token: tok, TLS: t.tls}) +} + +// httpBase is the HTTP API's base URL. +func (t *target) httpBase() (string, error) { + if err := t.resolve(); err != nil { + return "", err + } + return t.http, nil +} + +func (t *target) httpClient() *http.Client { + return &http.Client{Transport: &http.Transport{ + Proxy: http.ProxyFromEnvironment, + TLSClientConfig: &tls.Config{MinVersion: tls.VersionTLS12}, + }} +} + +func loopback(addr string) bool { + host, _, err := net.SplitHostPort(addr) + if err != nil { + host = addr + } + if host == "localhost" { + return true + } + ip := net.ParseIP(host) + return ip != nil && ip.IsLoopback() +} + +func firstNonEmpty(vs ...string) string { + for _, v := range vs { + if v != "" { + return v + } + } + return "" +} diff --git a/cmd/tokenmeter/context_cmd.go b/cmd/tokenmeter/context_cmd.go new file mode 100644 index 0000000..d313f8e --- /dev/null +++ b/cmd/tokenmeter/context_cmd.go @@ -0,0 +1,232 @@ +package main + +import ( + "bufio" + "errors" + "fmt" + "os" + "path/filepath" + "strings" + + "github.com/spf13/cobra" + "golang.org/x/term" +) + +func (a *app) contextCmd() *cobra.Command { + cmd := &cobra.Command{ + Use: "context", + Short: "Name services and where their tokens live (~/.config/tokenmeter/config.yaml)", + } + complete := func(_ *cobra.Command, args []string, _ string) ([]string, cobra.ShellCompDirective) { + if len(args) > 0 { + return nil, cobra.ShellCompDirectiveNoFileComp + } + cs, err := loadContexts() + if err != nil { + return nil, cobra.ShellCompDirectiveError + } + return cs.names(), cobra.ShellCompDirectiveNoFileComp + } + + var format string + list := &cobra.Command{ + Use: "list", Short: "List contexts; * marks the current one", Args: cobra.NoArgs, + RunE: func(*cobra.Command, []string) error { + cs, err := loadContexts() + if err != nil { + return err + } + p, err := newPrinter(a.stdout, format) + if err != nil { + return err + } + if p.json() { + return p.data(cs) + } + var rows [][]string + for _, n := range cs.names() { + c, cur := cs.Contexts[n], " " + if n == cs.Current { + cur = "*" + } + rows = append(rows, []string{cur, n, c.Addr, yesNo(c.TLS), c.HTTP, c.Token.String()}) + } + return p.rows([]string{"", "NAME", "ADDR", "TLS", "HTTP", "TOKEN"}, rows) + }, + } + outputFlag(list, &format) + + use := &cobra.Command{ + Use: "use NAME", Short: "Make a context the current one", Args: cobra.ExactArgs(1), ValidArgsFunction: complete, + RunE: func(_ *cobra.Command, args []string) error { + cs, err := loadContexts() + if err != nil { + return err + } + if cs.Contexts[args[0]] == nil { + return fmt.Errorf("no context %q", args[0]) + } + cs.Current = args[0] + return cs.save() + }, + } + + var addr, httpURL, token string + var useTLS bool + set := &cobra.Command{ + Use: "set NAME [--addr HOST:PORT] [--tls] [--http URL] [--token SOURCE]", + Short: "Create or change a context", + Long: `Create or change a context. A token SOURCE is env:VAR, file:PATH or keychain +(macOS Keychain or the Linux secret service; fill it with tokenmeter context +store-token). The first context created becomes the current one.`, + Args: cobra.ExactArgs(1), ValidArgsFunction: complete, + RunE: func(cmd *cobra.Command, args []string) error { + cs, err := loadContexts() + if err != nil { + return err + } + c := cs.Contexts[args[0]] + if c == nil { + c = &Context{} + cs.Contexts[args[0]] = c + } + f := cmd.Flags() + if f.Changed("addr") { + c.Addr = addr + } + if f.Changed("tls") { + c.TLS = useTLS + } + if f.Changed("http") { + c.HTTP = httpURL + } + if f.Changed("token") { + if c.Token, err = parseKeySource(token); err != nil { + return usageError{err} + } + } + if cs.Current == "" { + cs.Current = args[0] + } + return cs.save() + }, + } + // Local flags shadow the global connection flags: here they describe the context. + set.Flags().StringVar(&addr, "addr", "", "gRPC API address host:port") + set.Flags().BoolVar(&useTLS, "tls", false, "use TLS") + set.Flags().StringVar(&httpURL, "http", "", "HTTP API base URL") + set.Flags().StringVar(&token, "token", "", "token source: env:VAR, file:PATH or keychain") + + del := &cobra.Command{ + Use: "delete NAME", Short: "Remove a context (keychain entries stay)", Args: cobra.ExactArgs(1), ValidArgsFunction: complete, + RunE: func(_ *cobra.Command, args []string) error { + cs, err := loadContexts() + if err != nil { + return err + } + if cs.Contexts[args[0]] == nil { + return fmt.Errorf("no context %q", args[0]) + } + delete(cs.Contexts, args[0]) + if cs.Current == args[0] { + cs.Current = "" + } + return cs.save() + }, + } + + var toFile bool + store := &cobra.Command{ + Use: "store-token NAME [--file]", + Short: "Read a token from stdin and keep it in the keychain (or a 0600 file)", + Long: `Read a token from stdin, without echo on a terminal, and keep it in the system +keychain; --file keeps it in ~/.config/tokenmeter/NAME.token (mode 0600) +instead. The context's token source is updated to match.`, + Args: cobra.ExactArgs(1), ValidArgsFunction: complete, + RunE: func(_ *cobra.Command, args []string) error { + name := args[0] + cs, err := loadContexts() + if err != nil { + return err + } + c := cs.Contexts[name] + if c == nil { + return fmt.Errorf("no context %q (tokenmeter context set %s first)", name, name) + } + plain, err := a.readSecret("Token: ") + if err != nil { + return err + } + var src KeySource + if toFile { + dir, err := configDir() + if err != nil { + return err + } + if err := os.MkdirAll(dir, 0o700); err != nil { + return err + } + path := filepath.Join(dir, safeName(name)+".token") + if err := writePrivate(path, []byte(plain+"\n")); err != nil { + return err + } + src = KeySource{File: path} + } else { + if err := keychainSet(name, plain); err != nil { + return fmt.Errorf("%w (pass --file to keep the token in a file instead)", err) + } + src = KeySource{Keychain: true} + } + c.Token = src + if err := cs.save(); err != nil { + return err + } + fmt.Fprintf(a.stdout, "Stored the token of %s in %s.\n", name, src) + return nil + }, + } + store.Flags().BoolVar(&toFile, "file", false, "keep the token in a 0600 file instead of the keychain") + + cmd.AddCommand(list, use, set, del, store) + return cmd +} + +// readSecret reads one line from stdin, without echo on a terminal. +func (a *app) readSecret(prompt string) (string, error) { + var line string + if f, ok := a.stdin.(*os.File); ok && term.IsTerminal(int(f.Fd())) { + fmt.Fprint(a.stderr, prompt) + b, err := term.ReadPassword(int(f.Fd())) + fmt.Fprintln(a.stderr) + if err != nil { + return "", err + } + line = string(b) + } else { + s, err := bufio.NewReader(a.stdin).ReadString('\n') + if err != nil && s == "" { + return "", errors.New("nothing on stdin") + } + line = s + } + if line = strings.TrimSpace(line); line == "" { + return "", errors.New("empty input") + } + return line, nil +} + +func safeName(s string) string { + return strings.Map(func(r rune) rune { + if r == '/' || r == '\\' || r == os.PathSeparator { + return '_' + } + return r + }, s) +} + +func yesNo(b bool) string { + if b { + return "yes" + } + return "no" +} diff --git a/cmd/tokenmeter/contexts.go b/cmd/tokenmeter/contexts.go new file mode 100644 index 0000000..4873e8c --- /dev/null +++ b/cmd/tokenmeter/contexts.go @@ -0,0 +1,234 @@ +package main + +import ( + "bytes" + "errors" + "fmt" + "os" + "os/exec" + "path/filepath" + "runtime" + "sort" + "strings" + + "go.yaml.in/yaml/v3" +) + +// Contexts is ~/.config/tokenmeter/config.yaml: named services and where each +// one's token lives. Tokens themselves are never written to this file. +type Contexts struct { + Current string `yaml:"current_context,omitempty" json:"current_context,omitempty"` + Contexts map[string]*Context `yaml:"contexts,omitempty" json:"contexts,omitempty"` +} + +// Context is one light-tokenmeter. +type Context struct { + Addr string `yaml:"addr,omitempty" json:"addr,omitempty"` // gRPC host:port + TLS bool `yaml:"tls,omitempty" json:"tls,omitempty"` + HTTP string `yaml:"http,omitempty" json:"http,omitempty"` // base URL of /v1 and the probes + Token KeySource `yaml:"token,omitempty" json:"token,omitempty"` // bearer token, read or write scope +} + +// KeySource says where a token is read from; exactly one field is set. +type KeySource struct { + Env string `yaml:"env,omitempty" json:"env,omitempty"` + File string `yaml:"file,omitempty" json:"file,omitempty"` + Keychain bool `yaml:"keychain,omitempty" json:"keychain,omitempty"` +} + +func (k KeySource) empty() bool { return k == KeySource{} } + +func (k KeySource) String() string { + switch { + case k.Env != "": + return "env:" + k.Env + case k.File != "": + return "file:" + k.File + case k.Keychain: + return "keychain" + } + return "-" +} + +// parseKeySource reads env:VAR, file:PATH or keychain. +func parseKeySource(s string) (KeySource, error) { + switch { + case strings.HasPrefix(s, "env:") && len(s) > 4: + return KeySource{Env: s[4:]}, nil + case strings.HasPrefix(s, "file:") && len(s) > 5: + return KeySource{File: s[5:]}, nil + case s == "keychain": + return KeySource{Keychain: true}, nil + } + return KeySource{}, fmt.Errorf("key source %q: want env:VAR, file:PATH or keychain", s) +} + +// read returns the token; account names the keychain entry. +func (k KeySource) read(account string) (string, error) { + switch { + case k.Env != "": + if v := os.Getenv(k.Env); v != "" { + return v, nil + } + return "", fmt.Errorf("%s is not set", k.Env) + case k.File != "": + return readKeyFile(k.File) + case k.Keychain: + return keychainGet(account) + } + return "", errors.New("no token source") +} + +func readKeyFile(path string) (string, error) { + b, err := os.ReadFile(expandHome(path)) + if err != nil { + return "", err + } + k := strings.TrimSpace(string(b)) + if k == "" { + return "", fmt.Errorf("%s is empty", path) + } + return k, nil +} + +func configDir() (string, error) { + if d := os.Getenv("XDG_CONFIG_HOME"); d != "" { + return filepath.Join(d, "tokenmeter"), nil + } + home, err := os.UserHomeDir() + if err != nil { + return "", err + } + return filepath.Join(home, ".config", "tokenmeter"), nil +} + +// contextsUnreachable is a contexts file with no place to be, as without $HOME. +type contextsUnreachable struct{ error } + +func contextsPath() (string, error) { + if p := os.Getenv("TOKENMETER_CONFIG"); p != "" { + return p, nil + } + d, err := configDir() + if err != nil { + return "", contextsUnreachable{err} + } + return filepath.Join(d, "config.yaml"), nil +} + +// loadContexts returns an empty set when the file does not exist yet. +func loadContexts() (*Contexts, error) { + p, err := contextsPath() + if err != nil { + return nil, err + } + b, err := os.ReadFile(p) + if errors.Is(err, os.ErrNotExist) { + return &Contexts{Contexts: map[string]*Context{}}, nil + } + if err != nil { + return nil, err + } + c := &Contexts{} + if err := yaml.Unmarshal(b, c); err != nil { + return nil, fmt.Errorf("%s: %w", p, err) + } + if c.Contexts == nil { + c.Contexts = map[string]*Context{} + } + return c, nil +} + +func (c *Contexts) save() error { + p, err := contextsPath() + if err != nil { + return err + } + if err := os.MkdirAll(filepath.Dir(p), 0o700); err != nil { + return err + } + b, err := yaml.Marshal(c) + if err != nil { + return err + } + return writePrivate(p, b) +} + +func (c *Contexts) names() []string { + out := make([]string, 0, len(c.Contexts)) + for n := range c.Contexts { + out = append(out, n) + } + sort.Strings(out) + return out +} + +// writePrivate writes a file only its owner can read, replacing it whole. +func writePrivate(path string, b []byte) error { + tmp := path + ".tmp" + if err := os.WriteFile(tmp, b, 0o600); err != nil { + return err + } + return os.Rename(tmp, path) +} + +func expandHome(p string) string { + if p == "~" || strings.HasPrefix(p, "~/") { + if home, err := os.UserHomeDir(); err == nil { + return filepath.Join(home, p[1:]) + } + } + return p +} + +// The system keychain through its own tools: security on macOS, secret-tool +// (libsecret) on Linux. The token goes over stdin, never on a command line. +const keychainService = "tokenmeter" + +func keychainGet(account string) (string, error) { + var cmd *exec.Cmd + switch runtime.GOOS { + case "darwin": + cmd = exec.Command("security", "find-generic-password", "-s", keychainService, "-a", account, "-w") + case "linux": + cmd = exec.Command("secret-tool", "lookup", "service", keychainService, "account", account) + default: + return "", fmt.Errorf("no keychain support on %s", runtime.GOOS) + } + var out, errb bytes.Buffer + cmd.Stdout, cmd.Stderr = &out, &errb + if err := cmd.Run(); err != nil { + return "", fmt.Errorf("keychain entry %s/%s: %v %s", keychainService, account, err, strings.TrimSpace(errb.String())) + } + k := strings.TrimSpace(out.String()) + if k == "" { + return "", fmt.Errorf("keychain entry %s/%s is empty", keychainService, account) + } + return k, nil +} + +func keychainSet(account, key string) error { + var cmd *exec.Cmd + switch runtime.GOOS { + case "darwin": + // security -i reads commands from stdin, keeping the token out of ps. + cmd = exec.Command("security", "-i") + cmd.Stdin = strings.NewReader(fmt.Sprintf("add-generic-password -U -s %s -a %s -w %s\n", + keychainService, quoteSecurity(account), quoteSecurity(key))) + case "linux": + cmd = exec.Command("secret-tool", "store", "--label=tokenmeter "+account, "service", keychainService, "account", account) + cmd.Stdin = strings.NewReader(key) + default: + return fmt.Errorf("no keychain support on %s", runtime.GOOS) + } + var errb bytes.Buffer + cmd.Stdout, cmd.Stderr = &errb, &errb + if err := cmd.Run(); err != nil { + return fmt.Errorf("keychain: %v %s", err, strings.TrimSpace(errb.String())) + } + return nil +} + +func quoteSecurity(s string) string { + return `"` + strings.NewReplacer(`\`, `\\`, `"`, `\"`).Replace(s) + `"` +} diff --git a/cmd/tokenmeter/errors.go b/cmd/tokenmeter/errors.go new file mode 100644 index 0000000..120d34a --- /dev/null +++ b/cmd/tokenmeter/errors.go @@ -0,0 +1,88 @@ +package main + +import ( + "context" + "errors" + "fmt" + "net" + "net/http" + + "google.golang.org/grpc/codes" + "google.golang.org/grpc/status" +) + +// Exit codes, by what the caller should do about the error. +const ( + exitOK = 0 + exitFailed = 1 // anything else + exitUsage = 2 // bad flags or arguments + exitAuth = 4 // no token, a wrong token or one without the scope + exitNotFound = 5 + exitRejected = 6 // the service refused the request as it stands; retrying will not help + exitRetry = 7 // unavailable, throttled or timed out; retry later +) + +const exitCodesHelp = `Exit codes: 0 done, 1 failed, 2 bad usage, +4 not authenticated or not allowed, 5 not found, 6 rejected as invalid, +7 unavailable or throttled (retry later).` + +type usageError struct{ error } + +// httpError is a non-2xx answer from the HTTP API. +type httpError struct { + status int + body string +} + +func (e httpError) Error() string { + return fmt.Sprintf("HTTP %d: %s", e.status, e.body) +} + +func exitCode(err error) int { + if err == nil { + return exitOK + } + var u usageError + if errors.As(err, &u) { + return exitUsage + } + var h httpError + if errors.As(err, &h) { + switch { + case h.status == http.StatusUnauthorized || h.status == http.StatusForbidden: + return exitAuth + case h.status == http.StatusNotFound: + return exitNotFound + case h.status == http.StatusTooManyRequests || h.status >= 500: + return exitRetry + case h.status >= 400: + return exitRejected + } + return exitFailed + } + var netErr net.Error + if errors.As(err, &netErr) || errors.Is(err, context.DeadlineExceeded) { + return exitRetry + } + if s, ok := status.FromError(err); ok { + switch s.Code() { + case codes.Unauthenticated, codes.PermissionDenied: + return exitAuth + case codes.NotFound: + return exitNotFound + case codes.InvalidArgument, codes.FailedPrecondition, codes.AlreadyExists, codes.OutOfRange: + return exitRejected + case codes.Unavailable, codes.ResourceExhausted, codes.DeadlineExceeded, codes.Aborted: + return exitRetry + } + } + return exitFailed +} + +// errText drops the "rpc error: code = ... desc =" wrapping of gRPC errors. +func errText(err error) string { + if s, ok := status.FromError(err); ok && s.Code() != codes.OK && s.Code() != codes.Unknown { + return fmt.Sprintf("%s: %s", s.Code(), s.Message()) + } + return err.Error() +} diff --git a/cmd/tokenmeter/event_cmd.go b/cmd/tokenmeter/event_cmd.go new file mode 100644 index 0000000..97a9886 --- /dev/null +++ b/cmd/tokenmeter/event_cmd.go @@ -0,0 +1,97 @@ +package main + +import ( + "crypto/rand" + "fmt" + "math" + "strconv" + "time" + + "github.com/spf13/cobra" + "google.golang.org/protobuf/types/known/timestamppb" + + usagev1 "github.com/ai-process/light-tokenmeter/gen/usage/v1" +) + +func (a *app) eventCmd() *cobra.Command { + ev := &usagev1.UsageEvent{} + var st, cost, at string + var meta map[string]string + send := &cobra.Command{ + Use: "send --project P --user U --action A --vendor V --model M [--status S] [...]", + Short: "Record one usage event, to check that writing works", + Long: "Record one usage event with a new event_id (or --id). It needs a token with the write scope when auth is on.", + Args: cobra.NoArgs, + RunE: func(*cobra.Command, []string) error { + s, ok := statuses[st] + if !ok { + return usageError{fmt.Errorf("--status %q: want success, error, cancelled or timeout", st)} + } + ev.Status = s + if ev.EventId == "" { + ev.EventId = newUUID() + } + if cost != "" { + usd, err := strconv.ParseFloat(cost, 64) + if err != nil || !(usd >= 0 && usd < math.MaxInt64/1e9) { + return usageError{fmt.Errorf("--cost-usd %q: want an amount of 0 or more", cost)} + } + ev.CostNanousd = int64(math.Round(usd * 1e9)) + } + if at != "" { + t, err := time.Parse(time.RFC3339, at) + if err != nil { + return usageError{fmt.Errorf("--at %q: want RFC 3339", at)} + } + ev.OccurredAt = timestamppb.New(t) + } + ev.Meta = meta + c, err := a.t.client() + if err != nil { + return err + } + defer c.Close() + ctx, cancel := a.ctx() + defer cancel() + res, err := c.Service().RecordUsageEvent(ctx, &usagev1.RecordUsageEventRequest{Event: ev}) + if err != nil { + return err + } + if res.GetCreated() { + fmt.Fprintf(a.stdout, "Recorded event %s.\n", ev.EventId) + } else { + fmt.Fprintf(a.stdout, "Event %s was already recorded; nothing changed.\n", ev.EventId) + } + return nil + }, + } + f := send.Flags() + f.StringVar(&ev.EventId, "id", "", "event_id (UUID; default a new one)") + f.StringVar(&ev.ProjectId, "project", "", "project_id") + f.StringVar(&ev.UserId, "user", "", "user_id") + f.StringVar(&ev.ActionId, "action", "", "action_id") + f.StringVar(&ev.Vendor, "vendor", "", "vendor, e.g. openai") + f.StringVar(&ev.Model, "model", "", "model") + f.StringVar(&st, "status", "success", "success, error, cancelled or timeout") + f.Int64Var(&ev.TokensInput, "tokens-in", 0, "input tokens") + f.Int64Var(&ev.TokensOutput, "tokens-out", 0, "output tokens") + f.Int64Var(&ev.DurationMs, "duration-ms", 0, "latency") + f.StringVar(&cost, "cost-usd", "", "cost in USD") + f.StringVar(&at, "at", "", "occurred_at, RFC 3339 (default now)") + f.StringToStringVar(&meta, "meta", nil, "meta key=value, repeatable") + for _, req := range []string{"project", "user", "action", "vendor", "model"} { + _ = send.MarkFlagRequired(req) + } + cmd := &cobra.Command{Use: "event", Short: "Usage events"} + cmd.AddCommand(send) + return cmd +} + +// newUUID is a random (version 4) UUID. +func newUUID() string { + b := make([]byte, 16) + _, _ = rand.Read(b) + b[6] = b[6]&0x0f | 0x40 + b[8] = b[8]&0x3f | 0x80 + return fmt.Sprintf("%x-%x-%x-%x-%x", b[0:4], b[4:6], b[6:8], b[8:10], b[10:]) +} diff --git a/cmd/tokenmeter/http_cmd.go b/cmd/tokenmeter/http_cmd.go new file mode 100644 index 0000000..61c6639 --- /dev/null +++ b/cmd/tokenmeter/http_cmd.go @@ -0,0 +1,89 @@ +package main + +import ( + "encoding/json" + "fmt" + "io" + "net/http" + "strings" + + "github.com/spf13/cobra" +) + +func (a *app) healthCmd() *cobra.Command { + var format string + cmd := &cobra.Command{ + Use: "health", + Short: "Check /healthz and /readyz and show the service's version", + Args: cobra.NoArgs, + RunE: func(*cobra.Command, []string) error { + p, err := newPrinter(a.stdout, format) + if err != nil { + return err + } + base, err := a.t.httpBase() + if err != nil { + return err + } + ctx, cancel := a.ctx() + defer cancel() + type probe struct { + Path string `json:"path"` + Status int `json:"status"` + Body map[string]any `json:"body,omitempty"` + Problem string `json:"error,omitempty"` + } + var probes []probe + ready := true + for _, path := range []string{"/healthz", "/readyz"} { + pr := probe{Path: path} + req, err := http.NewRequestWithContext(ctx, http.MethodGet, base+path, nil) + if err != nil { + return err + } + res, err := a.t.httpClient().Do(req) + if err != nil { + pr.Problem, ready = err.Error(), false + } else { + raw, _ := io.ReadAll(io.LimitReader(res.Body, 1<<16)) + res.Body.Close() + pr.Status = res.StatusCode + _ = json.Unmarshal(raw, &pr.Body) + ready = ready && res.StatusCode == http.StatusOK + } + probes = append(probes, pr) + } + if p.json() { + if err := p.data(map[string]any{"http": base, "probes": probes}); err != nil { + return err + } + } else { + var rows [][]string + for _, pr := range probes { + st, detail := strings.TrimSpace(fmt.Sprint(pr.Status)), pr.Problem + if pr.Problem == "" { + detail = fmt.Sprint(pr.Body["status"]) + if v, ok := pr.Body["version"]; ok { + detail += ", version " + fmt.Sprint(v) + } + if d, ok := pr.Body["database"]; ok { + detail += ", database " + fmt.Sprint(d) + } + } else { + st = "-" + } + rows = append(rows, []string{base + pr.Path, st, detail}) + } + if err := p.rows([]string{"PROBE", "STATUS", "DETAIL"}, rows); err != nil { + return err + } + } + if !ready { + return codeError{exitRetry} + } + return nil + }, + } + outputFlag(cmd, &format) + return cmd +} diff --git a/cmd/tokenmeter/main.go b/cmd/tokenmeter/main.go new file mode 100644 index 0000000..048caca --- /dev/null +++ b/cmd/tokenmeter/main.go @@ -0,0 +1,146 @@ +// Command tokenmeter reads and writes light-tokenmeter: usage summaries, +// top spenders, usage by your own metadata and exports, test events, the +// service's health, and new tokens for TOKENMETER_TOKENS. +package main + +import ( + "context" + "errors" + "fmt" + "io" + "os" + "strings" + "time" + + "github.com/spf13/cobra" +) + +// version is set at release time with -ldflags "-X main.version=...". +var version = "dev" + +func main() { + os.Exit(run(os.Args[1:], os.Stdin, os.Stdout, os.Stderr)) +} + +// app is one invocation: the streams and the connection flags. +type app struct { + stdin io.Reader + stdout, stderr io.Writer + t target + timeout time.Duration +} + +func run(args []string, stdin io.Reader, stdout, stderr io.Writer) int { + a := &app{stdin: stdin, stdout: stdout, stderr: stderr} + a.t.warn = stderr + root := a.root() + root.SetArgs(args) + root.SetIn(stdin) + root.SetOut(stdout) + root.SetErr(stderr) + err := root.Execute() + if err == nil { + return exitOK + } + var c codeError + if errors.As(err, &c) { + return c.code + } + fmt.Fprintln(stderr, "tokenmeter:", errText(err)) + if isUsage(err) { + return exitUsage + } + return exitCode(err) +} + +// codeError ends the command with a code and nothing more to print. +type codeError struct{ code int } + +func (e codeError) Error() string { return fmt.Sprintf("exit %d", e.code) } + +// isUsage tells cobra's own errors from the API's. Argument and flag errors +// are wrapped as usageError by strict; these two come from cobra's dispatch. +func isUsage(err error) bool { + var u usageError + if errors.As(err, &u) { + return true + } + msg := err.Error() + return strings.HasPrefix(msg, "unknown command") || strings.HasPrefix(msg, "required flag") +} + +// strict makes argument errors usage errors, and a command group called +// without a known subcommand one too instead of help with exit 0. +func strict(cmd *cobra.Command) { + if cmd.HasSubCommands() && cmd.Run == nil && cmd.RunE == nil { + cmd.Args = cobra.NoArgs + cmd.RunE = func(c *cobra.Command, _ []string) error { + return usageError{fmt.Errorf("%s needs a subcommand (%s --help)", c.CommandPath(), c.CommandPath())} + } + } + if v := cmd.Args; v != nil { + cmd.Args = func(c *cobra.Command, args []string) error { + if err := v(c, args); err != nil { + return usageError{err} + } + return nil + } + } + for _, c := range cmd.Commands() { + strict(c) + } +} + +func (a *app) root() *cobra.Command { + root := &cobra.Command{ + Use: "tokenmeter", + Short: "Read and write light-tokenmeter", + Long: `tokenmeter reads usage from light-tokenmeter and sends it test events. + +Connection: --context, or --addr (TOKENMETER_ADDR, default localhost:9090, the +gRPC API) and --tls (TOKENMETER_TLS=true); the HTTP API at --http +(TOKENMETER_HTTP, default port 8080 of the same host). The token comes from +--token-file, TOKENMETER_TOKEN or the context: read for reports, write for +events. Without one the call goes out unauthenticated. + +` + exitCodesHelp, + Version: version, + SilenceUsage: true, + SilenceErrors: true, + PersistentPreRun: func(cmd *cobra.Command, _ []string) { + a.t.tlsSet = cmd.Flags().Changed("tls") + }, + } + root.SetVersionTemplate("{{.Version}}\n") + root.SetFlagErrorFunc(func(_ *cobra.Command, err error) error { return usageError{err} }) + pf := root.PersistentFlags() + pf.StringVar(&a.t.context, "context", "", "context from ~/.config/tokenmeter/config.yaml (TOKENMETER_CONTEXT)") + pf.StringVar(&a.t.addr, "addr", "", "gRPC API address host:port (TOKENMETER_ADDR)") + pf.BoolVar(&a.t.tls, "tls", false, "use TLS (TOKENMETER_TLS)") + pf.StringVar(&a.t.http, "http", "", "HTTP API base URL (TOKENMETER_HTTP)") + pf.StringVar(&a.t.tokenFile, "token-file", "", "read the token from this file instead of TOKENMETER_TOKEN") + pf.DurationVar(&a.timeout, "timeout", time.Minute, "time limit for each command") + + root.AddCommand(a.usageCmd(), a.tokenCmd(), a.eventCmd(), a.healthCmd(), a.contextCmd(), &cobra.Command{ + Use: "version", Short: "Print the version", Args: cobra.NoArgs, + RunE: func(*cobra.Command, []string) error { + fmt.Fprintln(a.stdout, version) + return nil + }, + }) + for _, c := range root.Commands() { + strict(c) + } + return root +} + +// ctx bounds one command by --timeout. +func (a *app) ctx() (context.Context, context.CancelFunc) { + return context.WithTimeout(context.Background(), a.timeout) +} + +// outputFlag adds -o to a command that prints data. +func outputFlag(cmd *cobra.Command, dst *string) { + cmd.Flags().StringVarP(dst, "output", "o", "table", "output format: table, json or csv") + _ = cmd.RegisterFlagCompletionFunc("output", cobra.FixedCompletions([]string{"table", "json", "csv"}, cobra.ShellCompDirectiveNoFileComp)) +} diff --git a/cmd/tokenmeter/main_test.go b/cmd/tokenmeter/main_test.go new file mode 100644 index 0000000..afde1d5 --- /dev/null +++ b/cmd/tokenmeter/main_test.go @@ -0,0 +1,255 @@ +package main + +import ( + "bytes" + "crypto/sha256" + "encoding/csv" + "encoding/hex" + "encoding/json" + "net" + "os" + "path/filepath" + "strings" + "testing" + "time" + + "github.com/ai-process/light-tokenmeter/internal/auth" + "github.com/ai-process/light-tokenmeter/internal/repository" + "github.com/ai-process/light-tokenmeter/internal/server" + "github.com/ai-process/light-tokenmeter/internal/service" + "github.com/ai-process/light-tokenmeter/internal/testdb" +) + +func callIn(stdin string, args ...string) (int, string, string) { + var out, errb bytes.Buffer + code := run(args, strings.NewReader(stdin), &out, &errb) + return code, out.String(), errb.String() +} + +func call(args ...string) (int, string, string) { return callIn("", args...) } + +// isolate points the CLI at an empty contexts file and clears the environment. +func isolate(t *testing.T) string { + t.Helper() + p := filepath.Join(t.TempDir(), "config.yaml") + t.Setenv("TOKENMETER_CONFIG", p) + t.Setenv("XDG_CONFIG_HOME", filepath.Dir(p)) + for _, k := range []string{"TOKENMETER_ADDR", "TOKENMETER_TLS", "TOKENMETER_HTTP", "TOKENMETER_TOKEN", "TOKENMETER_CONTEXT"} { + t.Setenv(k, "") + } + return p +} + +func TestUsageErrors(t *testing.T) { + isolate(t) + for _, args := range [][]string{ + {"nonsense"}, {"usage"}, {"usage", "sumary"}, {"usage", "summary", "--status", "ok"}, + {"usage", "summary", "--group-by", "colour"}, {"usage", "summary", "--last", "7x"}, + {"usage", "summary", "--from", "2026-10-02", "--to", "2026-10-01"}, {"usage", "summary", "--last", "1d", "--from", "2026-10-01"}, + {"usage", "top", "--by", "vibes"}, {"usage", "by-meta", "--key", "job_id"}, {"usage", "export", "--format", "xml"}, + {"token", "new", "--name", "x"}, {"token", "new", "--name", "a:b", "--scope", "read"}, {"token", "new", "--name", "x", "--scope", "admin"}, + {"event", "send", "--project", "p"}, {"usage", "summary", "-o", "yaml"}, + } { + if code, _, errs := call(args...); code != exitUsage { + t.Errorf("tokenmeter %v = %d (%s), want %d", args, code, errs, exitUsage) + } + } + for _, args := range [][]string{{"version"}, {"completion", "fish"}, {"--help"}} { + if code, _, errs := call(args...); code != 0 { + t.Errorf("tokenmeter %v = %d: %s", args, code, errs) + } + } +} + +func TestTokenNewAndHash(t *testing.T) { + isolate(t) + code, out, _ := call("token", "new", "--name", "ktms", "--scope", "write,read") + lines := strings.Split(strings.TrimSpace(out), "\n") + if code != 0 || len(lines) != 2 || len(lines[0]) != 64 { + t.Fatalf("token new = %d %q", code, out) + } + tok, line := lines[0], lines[1] + sum := sha256.Sum256([]byte(tok)) + if line != "ktms:write+read:"+hex.EncodeToString(sum[:]) { + t.Fatalf("entry %q", line) + } + if _, err := auth.Parse(line); err != nil { + t.Fatalf("the service does not accept the entry: %v", err) + } + if _, out, _ := callIn(tok+"\n", "token", "hash"); strings.TrimSpace(out) != hex.EncodeToString(sum[:]) { + t.Fatalf("hash = %q", out) + } + if _, out, _ := callIn(tok+"\n", "token", "hash", "--name", "ktms", "--scope", "write,read"); strings.TrimSpace(out) != line { + t.Fatalf("hash entry = %q", out) + } + _, out, _ = call("token", "new", "--name", "c", "--scope", "read", "-o", "json") + var v struct{ Token, Entry string } + if err := json.Unmarshal([]byte(out), &v); err != nil || !strings.HasPrefix(v.Entry, "c:read:") || len(v.Token) != 64 { + t.Fatalf("token new -o json = %s", out) + } +} + +func TestContexts(t *testing.T) { + p := isolate(t) + if code, _, errs := call("context", "set", "prod", "--addr", "tm.example.com:443", "--tls", "--token", "env:PROD_TM"); code != 0 { + t.Fatalf("context set = %d %s", code, errs) + } + if info, err := os.Stat(p); err != nil || info.Mode().Perm() != 0o600 { + t.Fatalf("contexts file: %v", err) + } + t.Setenv("PROD_TM", "tok-env") + tg := target{} + if err := tg.resolve(); err != nil { + t.Fatal(err) + } + if tg.addr != "tm.example.com:443" || !tg.tls || tg.http != "https://tm.example.com:8080" { + t.Fatalf("resolved %+v", tg) + } + if tok, _ := tg.token(); tok != "tok-env" { + t.Fatalf("token = %q", tok) + } + if code, _, errs := callIn("tok-file\n", "context", "store-token", "prod", "--file"); code != 0 { + t.Fatalf("store-token = %d %s", code, errs) + } + tg = target{} + if tok, err := tg.token(); err != nil || tok != "tok-file" { + t.Fatalf("token after store-token = %q %v", tok, err) + } + t.Setenv("TOKENMETER_TOKEN", "tok-flagenv") + tg = target{} + if tok, _ := tg.token(); tok != "tok-flagenv" { + t.Fatalf("the environment does not win over the context: %q", tok) + } + if _, out, _ := call("context", "list"); !strings.Contains(out, "* prod") { + t.Fatalf("context list: %s", out) + } +} + +// serve runs the real gRPC and HTTP servers over a fresh database with these +// token entries, and returns their addresses. +func serve(t *testing.T, tokens string) (grpcAddr, httpBase string) { + t.Helper() + pool := testdb.Postgres(t) + verifier, err := auth.Parse(tokens) + if err != nil { + t.Fatal(err) + } + repo := repository.New(pool) + svc := service.NewUsageService(repo) + free := func() string { + l, err := net.Listen("tcp", "127.0.0.1:0") + if err != nil { + t.Fatal(err) + } + defer l.Close() + return l.Addr().String() + } + grpcAddr, httpAddr := free(), free() + g := server.NewGRPCServerWithService(svc, grpcAddr, verifier) + h := server.NewHTTPServerWithService(svc, repo, httpAddr, verifier) + go func() { _ = g.Start() }() + go func() { _ = h.Start() }() + t.Cleanup(g.Stop) + for _, a := range []string{grpcAddr, httpAddr} { + for i := 0; i < 100; i++ { + if c, err := net.Dial("tcp", a); err == nil { + _ = c.Close() + break + } + time.Sleep(20 * time.Millisecond) + } + } + return grpcAddr, "http://" + httpAddr +} + +func TestAgainstTheService(t *testing.T) { + isolate(t) + _, out, _ := call("token", "new", "--name", "writer", "--scope", "write") + w := strings.Split(strings.TrimSpace(out), "\n") + _, out, _ = call("token", "new", "--name", "reader", "--scope", "read") + r := strings.Split(strings.TrimSpace(out), "\n") + addr, base := serve(t, w[1]+","+r[1]) + t.Setenv("TOKENMETER_ADDR", addr) + t.Setenv("TOKENMETER_HTTP", base) + must := func(want int, args ...string) string { + t.Helper() + code, out, errs := call(args...) + if code != want { + t.Fatalf("tokenmeter %s = %d, want %d\nstdout: %s\nstderr: %s", strings.Join(args, " "), code, want, out, errs) + } + return out + } + + t.Setenv("TOKENMETER_TOKEN", w[0]) + at := time.Now().UTC().Add(-time.Hour).Format(time.RFC3339) + send := func(project, model, cost string, meta ...string) { + args := []string{"event", "send", "--project", project, "--user", "u1", "--action", "a", "--vendor", "openai", + "--model", model, "--tokens-in", "100", "--tokens-out", "50", "--cost-usd", cost, "--at", at} + for _, m := range meta { + args = append(args, "--meta", m) + } + must(0, args...) + } + send("p1", "m1", "0.01", "job_id=j1", "task_id=t1") + send("p1", "m2", "0.02", "job_id=j1", "task_id=t2") + send("p2", "m1", "0.5") + if out := must(0, "event", "send", "--id", "550e8400-e29b-41d4-a716-446655440000", "--project", "p", "--user", "u", "--action", "a", "--vendor", "v", "--model", "m"); !strings.Contains(out, "Recorded") { + t.Fatalf("event send: %s", out) + } + if out := must(0, "event", "send", "--id", "550e8400-e29b-41d4-a716-446655440000", "--project", "p", "--user", "u", "--action", "a", "--vendor", "v", "--model", "m"); !strings.Contains(out, "already") { + t.Fatalf("duplicate event send: %s", out) + } + must(exitAuth, "usage", "summary") // a write token cannot read + + t.Setenv("TOKENMETER_TOKEN", r[0]) + must(exitAuth, "event", "send", "--project", "p", "--user", "u", "--action", "a", "--vendor", "v", "--model", "m") + var sum struct { + Rows []struct { + ProjectID string `json:"project_id"` + CostNanousd string `json:"cost_nanousd"` + RequestCount string `json:"request_count"` + } `json:"rows"` + } + out = must(0, "usage", "summary", "--last", "1d", "--group-by", "project", "-o", "json") + if err := json.Unmarshal([]byte(out), &sum); err != nil { + t.Fatalf("summary json: %v\n%s", err, out) + } + got := map[string]string{} + for _, row := range sum.Rows { + got[row.ProjectID] = row.CostNanousd + "/" + row.RequestCount + } + if got["p1"] != "30000000/2" || got["p2"] != "500000000/1" { + t.Fatalf("summary by project = %v", got) + } + if out := must(0, "usage", "summary", "--last", "1d", "--group-by", "project"); !strings.Contains(out, "$0.0300") || !strings.Contains(out, "TOTAL") { + t.Fatalf("summary table:\n%s", out) + } + if out := must(0, "usage", "top", "--last", "1d", "--by", "cost", "--dimension", "project"); !strings.Contains(strings.Split(out, "\n")[1], "p2") { + t.Fatalf("top:\n%s", out) + } + if out := must(0, "usage", "by-meta", "--key", "job_id", "--value", "j1", "--group-by", "task_id"); !strings.Contains(out, "t1") || !strings.Contains(out, "t2") || !strings.Contains(out, "$0.0200") { + t.Fatalf("by-meta:\n%s", out) + } + + out = must(0, "usage", "export", "--last", "1d", "--group-by", "project") + recs, err := csv.NewReader(strings.NewReader(out)).ReadAll() + if err != nil || len(recs) < 3 || recs[0][0] != "day" { + t.Fatalf("export csv: %v\n%s", err, out) + } + usd := map[string]string{} + for _, rec := range recs[1:] { + usd[rec[1]] = rec[len(rec)-2] + } + if usd["p1"] != "0.03" || usd["p2"] != "0.5" { + t.Fatalf("export cost_usd = %v", usd) + } + if out := must(0, "usage", "export", "--last", "1d", "--format", "json"); !strings.Contains(out, `"rows"`) { + t.Fatalf("export json: %s", out) + } + + if out := must(0, "health"); !strings.Contains(out, "version dev") { + t.Fatalf("health:\n%s", out) + } + t.Setenv("TOKENMETER_TOKEN", "") + must(exitAuth, "usage", "summary") +} diff --git a/cmd/tokenmeter/output.go b/cmd/tokenmeter/output.go new file mode 100644 index 0000000..8bd03c1 --- /dev/null +++ b/cmd/tokenmeter/output.go @@ -0,0 +1,144 @@ +package main + +import ( + "bytes" + "encoding/csv" + "encoding/json" + "fmt" + "io" + "strconv" + "strings" + "text/tabwriter" + + "google.golang.org/protobuf/encoding/protojson" + "google.golang.org/protobuf/proto" +) + +// printer writes a result as a table for people, or as JSON or CSV for +// scripts. JSON carries the API messages as protojson renders them. +type printer struct { + w io.Writer + format string // table | json | csv +} + +func newPrinter(w io.Writer, format string) (*printer, error) { + switch format { + case "", "table": + return &printer{w: w, format: "table"}, nil + case "json", "csv": + return &printer{w: w, format: format}, nil + } + return nil, usageError{fmt.Errorf("-o %q: want table, json or csv", format)} +} + +func (p *printer) json() bool { return p.format == "json" } + +// data writes v, which may hold proto messages, as indented JSON. +func (p *printer) data(v any) error { + b, err := toJSON(v) + if err != nil { + return err + } + var out bytes.Buffer + if err := json.Indent(&out, b, "", " "); err != nil { + return err + } + out.WriteByte('\n') + _, err = p.w.Write(out.Bytes()) + return err +} + +// toJSON marshals proto messages with protojson and everything else with +// encoding/json, nesting freely. +func toJSON(v any) ([]byte, error) { + switch x := v.(type) { + case proto.Message: + return protojson.MarshalOptions{UseProtoNames: true, EmitUnpopulated: true}.Marshal(x) + case []any: + parts := make([]json.RawMessage, len(x)) + for i, e := range x { + b, err := toJSON(e) + if err != nil { + return nil, err + } + parts[i] = b + } + return json.Marshal(parts) + case map[string]any: + obj := make(map[string]json.RawMessage, len(x)) + for k, e := range x { + b, err := toJSON(e) + if err != nil { + return nil, err + } + obj[k] = b + } + return json.Marshal(obj) + } + return json.Marshal(v) +} + +func messages[T proto.Message](ms []T) []any { + out := make([]any, len(ms)) + for i, m := range ms { + out[i] = m + } + return out +} + +// rows writes a table, or CSV with the same header in snake_case. +func (p *printer) rows(header []string, rows [][]string) error { + if p.format == "csv" { + w := csv.NewWriter(p.w) + h := make([]string, len(header)) + for i, c := range header { + h[i] = strings.ToLower(strings.ReplaceAll(strings.TrimSpace(c), " ", "_")) + } + if err := w.Write(h); err != nil { + return err + } + if err := w.WriteAll(rows); err != nil { + return err + } + w.Flush() + return w.Error() + } + tw := tabwriter.NewWriter(p.w, 0, 0, 2, ' ', 0) + fmt.Fprintln(tw, strings.Join(header, "\t")) + for _, r := range rows { + cells := make([]string, len(r)) + for i, c := range r { + cells[i] = orDash(c) + } + fmt.Fprintln(tw, strings.Join(cells, "\t")) + } + return tw.Flush() +} + +// usd is a cost for this format: $1.2345 in a table, exact dollars in CSV. +func (p *printer) usd(nano int64) string { + if p.format == "csv" { + return exactUSD(nano) + } + return fmt.Sprintf("$%.4f", float64(nano)/1e9) +} + +// exactUSD writes nano-USD as dollars without float rounding: 36000000 is 0.036. +func exactUSD(nano int64) string { + sign := "" + if nano < 0 { + sign, nano = "-", -nano + } + frac := strings.TrimRight(fmt.Sprintf("%09d", nano%1_000_000_000), "0") + if frac == "" { + return sign + strconv.FormatInt(nano/1_000_000_000, 10) + } + return sign + strconv.FormatInt(nano/1_000_000_000, 10) + "." + frac +} + +func orDash(s string) string { + if s == "" { + return "-" + } + return s +} diff --git a/cmd/tokenmeter/token_cmd.go b/cmd/tokenmeter/token_cmd.go new file mode 100644 index 0000000..54e8a94 --- /dev/null +++ b/cmd/tokenmeter/token_cmd.go @@ -0,0 +1,104 @@ +package main + +import ( + "crypto/rand" + "crypto/sha256" + "encoding/hex" + "fmt" + "strings" + + "github.com/spf13/cobra" +) + +// entry is a TOKENMETER_TOKENS line: name:scope[+scope]:sha256hex. +func entry(name string, scopes []string, token string) string { + sum := sha256.Sum256([]byte(token)) + return name + ":" + strings.Join(scopes, "+") + ":" + hex.EncodeToString(sum[:]) +} + +func checkScopes(name string, scopes []string) error { + if name == "" || strings.ContainsAny(name, ":,\n#") { + return usageError{fmt.Errorf("--name %q: want a name without : , # or newlines", name)} + } + if len(scopes) == 0 { + return usageError{fmt.Errorf("--scope is required: read, write or both")} + } + for _, s := range scopes { + if s != "read" && s != "write" { + return usageError{fmt.Errorf("scope %q: want read or write", s)} + } + } + return nil +} + +func (a *app) tokenCmd() *cobra.Command { + cmd := &cobra.Command{Use: "token", Short: "Tokens for TOKENMETER_TOKENS"} + + var name, format string + var scopes []string + mint := &cobra.Command{ + Use: "new --name NAME --scope read|write[,...]", + Short: "Make a token; print it once, with its TOKENMETER_TOKENS entry", + Long: `Make a random token for a client and print it once, with the entry to add to +TOKENMETER_TOKENS (or TOKENMETER_TOKENS_FILE): name:scope[+scope]:sha256. The +service keeps only the hash. Nothing is sent anywhere.`, + Args: cobra.NoArgs, + RunE: func(*cobra.Command, []string) error { + if err := checkScopes(name, scopes); err != nil { + return err + } + b := make([]byte, 32) + if _, err := rand.Read(b); err != nil { + return err + } + tok := hex.EncodeToString(b) + line := entry(name, scopes, tok) + if format == "json" { + p := &printer{w: a.stdout, format: "json"} + return p.data(map[string]any{"name": name, "scopes": scopes, "token": tok, "entry": line}) + } + if format != "table" { + return usageError{fmt.Errorf("-o %q: want table or json", format)} + } + fmt.Fprintln(a.stderr, "Token for the client. It is shown only now; store it in your secret manager:") + fmt.Fprintln(a.stdout, tok) + fmt.Fprintln(a.stderr, "Entry for TOKENMETER_TOKENS:") + fmt.Fprintln(a.stdout, line) + return nil + }, + } + mint.Flags().StringVar(&name, "name", "", "who the token is for, logged with each call") + mint.Flags().StringSliceVar(&scopes, "scope", nil, "read, write or read,write") + mint.Flags().StringVarP(&format, "output", "o", "table", "table or json") + _ = mint.MarkFlagRequired("name") + _ = mint.RegisterFlagCompletionFunc("scope", cobra.FixedCompletions([]string{"read", "write"}, cobra.ShellCompDirectiveNoFileComp)) + + var hName string + var hScopes []string + hash := &cobra.Command{ + Use: "hash [--name NAME --scope SCOPES]", + Short: "Hash a token read from stdin; with --name and --scope print its entry", + Args: cobra.NoArgs, + RunE: func(*cobra.Command, []string) error { + tok, err := a.readSecret("Token: ") + if err != nil { + return err + } + if hName == "" && len(hScopes) == 0 { + sum := sha256.Sum256([]byte(tok)) + fmt.Fprintln(a.stdout, hex.EncodeToString(sum[:])) + return nil + } + if err := checkScopes(hName, hScopes); err != nil { + return err + } + fmt.Fprintln(a.stdout, entry(hName, hScopes, tok)) + return nil + }, + } + hash.Flags().StringVar(&hName, "name", "", "the token's name") + hash.Flags().StringSliceVar(&hScopes, "scope", nil, "read, write or read,write") + + cmd.AddCommand(mint, hash) + return cmd +} diff --git a/cmd/tokenmeter/usage_cmd.go b/cmd/tokenmeter/usage_cmd.go new file mode 100644 index 0000000..97dbc17 --- /dev/null +++ b/cmd/tokenmeter/usage_cmd.go @@ -0,0 +1,454 @@ +package main + +import ( + "errors" + "fmt" + "io" + "os" + "sort" + "strconv" + "strings" + "time" + + "github.com/spf13/cobra" + "google.golang.org/protobuf/types/known/timestamppb" + + usagev1 "github.com/ai-process/light-tokenmeter/gen/usage/v1" + "github.com/ai-process/light-tokenmeter/pkg/client" +) + +// dimensions are what a summary can be grouped by, in column order. +var dimensions = []string{"project", "user", "vendor", "model", "action", "status"} + +var statuses = map[string]usagev1.Status{ + "success": usagev1.Status_STATUS_SUCCESS, "error": usagev1.Status_STATUS_ERROR, + "cancelled": usagev1.Status_STATUS_CANCELLED, "timeout": usagev1.Status_STATUS_TIMEOUT, +} + +func statusName(s usagev1.Status) string { + if s == usagev1.Status_STATUS_UNSPECIFIED { + return "" + } + return strings.ToLower(strings.TrimPrefix(s.String(), "STATUS_")) +} + +// query is the window and filters shared by the usage commands. +type query struct { + from, to, last string + project, user, vendor, model, action, st string +} + +func (q *query) flags(cmd *cobra.Command) { + f := cmd.Flags() + f.StringVar(&q.from, "from", "", "start, inclusive: RFC 3339 or YYYY-MM-DD (UTC)") + f.StringVar(&q.to, "to", "", "end, exclusive: RFC 3339 or YYYY-MM-DD (default now)") + f.StringVar(&q.last, "last", "", "window ending now, e.g. 24h, 7d, 4w (default 7d)") + f.StringVar(&q.project, "project", "", "only this project_id") + f.StringVar(&q.user, "user", "", "only this user_id") + f.StringVar(&q.vendor, "vendor", "", "only this vendor") + f.StringVar(&q.model, "model", "", "only this model") + f.StringVar(&q.action, "action", "", "only this action_id") + f.StringVar(&q.st, "status", "", "only success, error, cancelled or timeout") + _ = cmd.RegisterFlagCompletionFunc("status", cobra.FixedCompletions([]string{"success", "error", "cancelled", "timeout"}, cobra.ShellCompDirectiveNoFileComp)) +} + +func parseTime(s string) (time.Time, error) { + if t, err := time.Parse(time.RFC3339, s); err == nil { + return t.UTC(), nil + } + if t, err := time.Parse(time.DateOnly, s); err == nil { + return t, nil + } + return time.Time{}, fmt.Errorf("time %q: want RFC 3339 (2026-10-01T00:00:00Z) or YYYY-MM-DD", s) +} + +// request builds the summary request; groups are dimension names. +func (q *query) request(groups []string) (*usagev1.GetUsageSummaryRequest, error) { + var from, to time.Time + var err error + if q.last != "" && q.from != "" { + return nil, usageError{errors.New("--last and --from do not go together")} + } + to = time.Now().UTC() + if q.to != "" { + if to, err = parseTime(q.to); err != nil { + return nil, usageError{err} + } + } + if q.from != "" { + if from, err = parseTime(q.from); err != nil { + return nil, usageError{err} + } + } else { + d, err := client.ParseDuration(firstNonEmpty(q.last, "7d")) + if err != nil { + return nil, usageError{err} + } + from = to.Add(-d) + } + if !from.Before(to) { + return nil, usageError{fmt.Errorf("the window %s – %s is empty", from.Format(time.RFC3339), to.Format(time.RFC3339))} + } + req := &usagev1.GetUsageSummaryRequest{ + From: timestamppb.New(from), To: timestamppb.New(to), + ProjectId: q.project, UserId: q.user, Vendor: q.vendor, Model: q.model, ActionId: q.action, + } + if q.st != "" { + s, ok := statuses[q.st] + if !ok { + return nil, usageError{fmt.Errorf("--status %q: want success, error, cancelled or timeout", q.st)} + } + req.Status = s + } + for _, g := range groups { + switch g { + case "project": + req.GroupByProjectId = true + case "user": + req.GroupByUserId = true + case "vendor": + req.GroupByVendor = true + case "model": + req.GroupByModel = true + case "action": + req.GroupByActionId = true + case "status": + req.GroupByStatus = true + default: + return nil, usageError{fmt.Errorf("group %q: want %s", g, strings.Join(dimensions, ", "))} + } + } + return req, nil +} + +func dimValue(r *usagev1.UsageSummaryRow, dim string) string { + switch dim { + case "project": + return r.GetProjectId() + case "user": + return r.GetUserId() + case "vendor": + return r.GetVendor() + case "model": + return r.GetModel() + case "action": + return r.GetActionId() + case "status": + return statusName(r.GetStatus()) + } + return "" +} + +// ordered keeps the dimensions asked for in column order. +func ordered(groups []string) []string { + var out []string + for _, d := range dimensions { + for _, g := range groups { + if g == d { + out = append(out, d) + break + } + } + } + return out +} + +// fetch runs the summary and warns when even the smallest window was capped. +func (a *app) fetch(req *usagev1.GetUsageSummaryRequest) ([]*usagev1.UsageSummaryRow, bool, error) { + c, err := a.t.client() + if err != nil { + return nil, false, err + } + defer c.Close() + ctx, cancel := a.ctx() + defer cancel() + rows, truncated, err := c.Summary(ctx, req) + if err != nil { + return nil, false, err + } + if truncated { + fmt.Fprintf(a.stderr, "tokenmeter: warning: an hour of this query still exceeds the server's row cap; totals are incomplete, narrow the filters\n") + } + return rows, truncated, nil +} + +func (a *app) usageCmd() *cobra.Command { + cmd := &cobra.Command{Use: "usage", Short: "Usage reports: summary, top, by-meta, export"} + cmd.AddCommand(a.summaryCmd(), a.topCmd(), a.byMetaCmd(), a.exportCmd()) + return cmd +} + +func (a *app) summaryCmd() *cobra.Command { + var q query + var groups []string + var daily bool + var format string + cmd := &cobra.Command{ + Use: "summary [--last 7d | --from T --to T] [filters] [--group-by DIM,...] [--daily]", + Short: "Requests, tokens and cost over a window, by the dimensions you group by", + Args: cobra.NoArgs, + RunE: func(*cobra.Command, []string) error { + p, err := newPrinter(a.stdout, format) + if err != nil { + return err + } + req, err := q.request(groups) + if err != nil { + return err + } + rows, truncated, err := a.fetch(req) + if err != nil { + return err + } + rows = client.Merge(rows, daily) + if daily { + sort.SliceStable(rows, func(i, j int) bool { return rows[i].GetBucket().GetSeconds() < rows[j].GetBucket().GetSeconds() }) + } else { + sort.SliceStable(rows, func(i, j int) bool { return rows[i].GetCostNanousd() > rows[j].GetCostNanousd() }) + } + if p.json() { + return p.data(map[string]any{"from": req.GetFrom(), "to": req.GetTo(), "rows": messages(rows), "truncated": truncated}) + } + dims := ordered(groups) + var header []string + if daily { + header = append(header, "DAY") + } + for _, d := range dims { + header = append(header, strings.ToUpper(d)) + } + header = append(header, "REQUESTS", "TOKENS IN", "TOKENS OUT", "CACHE READ", "COST USD", "USD PER REQUEST", "AVG MS") + var out [][]string + total := &usagev1.UsageSummaryRow{} + for _, r := range rows { + var line []string + if daily { + line = append(line, r.GetBucket().AsTime().UTC().Format(time.DateOnly)) + } + for _, d := range dims { + line = append(line, dimValue(r, d)) + } + out = append(out, append(line, a.figures(p, r)...)) + total = client.Merge([]*usagev1.UsageSummaryRow{total, stripped(r)}, false)[0] + } + if p.format == "table" && len(rows) > 1 { + line := make([]string, len(header)-7) + if len(line) > 0 { + line[0] = "TOTAL" + } + out = append(out, append(line, a.figures(p, total)...)) + } + return p.rows(header, out) + }, + } + q.flags(cmd) + cmd.Flags().StringSliceVar(&groups, "group-by", nil, "project, user, vendor, model, action, status") + cmd.Flags().BoolVar(&daily, "daily", false, "one row per UTC day") + _ = cmd.RegisterFlagCompletionFunc("group-by", cobra.FixedCompletions(dimensions, cobra.ShellCompDirectiveNoFileComp)) + outputFlag(cmd, &format) + return cmd +} + +// stripped is a row's figures without its dimensions, for totals. +func stripped(r *usagev1.UsageSummaryRow) *usagev1.UsageSummaryRow { + return &usagev1.UsageSummaryRow{TokensInput: r.GetTokensInput(), TokensOutput: r.GetTokensOutput(), TokensTotal: r.GetTokensTotal(), + RequestCount: r.GetRequestCount(), SumDurationMs: r.GetSumDurationMs(), CostNanousd: r.GetCostNanousd(), + CacheReadTokens: r.GetCacheReadTokens(), CacheWriteTokens: r.GetCacheWriteTokens()} +} + +func (a *app) figures(p *printer, r *usagev1.UsageSummaryRow) []string { + perReq, avg := "", "" + if n := r.GetRequestCount(); n > 0 { + perReq = p.usd(r.GetCostNanousd() / n) + avg = strconv.FormatInt(r.GetSumDurationMs()/n, 10) + } + return []string{strconv.FormatInt(r.GetRequestCount(), 10), strconv.FormatInt(r.GetTokensInput(), 10), + strconv.FormatInt(r.GetTokensOutput(), 10), strconv.FormatInt(r.GetCacheReadTokens(), 10), + p.usd(r.GetCostNanousd()), perReq, avg} +} + +func (a *app) topCmd() *cobra.Command { + var q query + var by, dim, format string + var limit int + cmd := &cobra.Command{ + Use: "top [--by cost|tokens|requests] [--dimension DIM] [--limit N]", + Short: "Who or what used the most over a window", + Args: cobra.NoArgs, + RunE: func(*cobra.Command, []string) error { + p, err := newPrinter(a.stdout, format) + if err != nil { + return err + } + measure := map[string]func(*usagev1.UsageSummaryRow) int64{ + "cost": (*usagev1.UsageSummaryRow).GetCostNanousd, + "tokens": (*usagev1.UsageSummaryRow).GetTokensTotal, + "requests": (*usagev1.UsageSummaryRow).GetRequestCount, + }[by] + if measure == nil { + return usageError{fmt.Errorf("--by %q: want cost, tokens or requests", by)} + } + req, err := q.request([]string{dim}) + if err != nil { + return err + } + rows, truncated, err := a.fetch(req) + if err != nil { + return err + } + rows = client.Merge(rows, false) + var sum int64 + for _, r := range rows { + sum += measure(r) + } + sort.SliceStable(rows, func(i, j int) bool { return measure(rows[i]) > measure(rows[j]) }) + if limit > 0 && len(rows) > limit { + rows = rows[:limit] + } + if p.json() { + return p.data(map[string]any{"from": req.GetFrom(), "to": req.GetTo(), "by": by, "dimension": dim, "rows": messages(rows), "truncated": truncated}) + } + var out [][]string + for i, r := range rows { + share := "" + if sum > 0 { + share = fmt.Sprintf("%.1f%%", 100*float64(measure(r))/float64(sum)) + } + out = append(out, []string{strconv.Itoa(i + 1), dimValue(r, dim), strconv.FormatInt(r.GetRequestCount(), 10), + strconv.FormatInt(r.GetTokensTotal(), 10), p.usd(r.GetCostNanousd()), share}) + } + return p.rows([]string{"RANK", strings.ToUpper(dim), "REQUESTS", "TOKENS", "COST USD", "SHARE OF " + strings.ToUpper(by)}, out) + }, + } + q.flags(cmd) + cmd.Flags().StringVar(&by, "by", "cost", "cost, tokens or requests") + cmd.Flags().StringVar(&dim, "dimension", "project", strings.Join(dimensions, ", ")) + cmd.Flags().IntVar(&limit, "limit", 10, "rows to show, 0 for all") + _ = cmd.RegisterFlagCompletionFunc("by", cobra.FixedCompletions([]string{"cost", "tokens", "requests"}, cobra.ShellCompDirectiveNoFileComp)) + _ = cmd.RegisterFlagCompletionFunc("dimension", cobra.FixedCompletions(dimensions, cobra.ShellCompDirectiveNoFileComp)) + outputFlag(cmd, &format) + return cmd +} + +func (a *app) byMetaCmd() *cobra.Command { + req := &usagev1.GetUsageByMetaRequest{} + var format string + cmd := &cobra.Command{ + Use: "by-meta --key KEY --value VALUE [--group-by KEY] [--project ID]", + Short: "What the events tagged meta KEY=VALUE used, by vendor and model", + Args: cobra.NoArgs, + RunE: func(*cobra.Command, []string) error { + p, err := newPrinter(a.stdout, format) + if err != nil { + return err + } + c, err := a.t.client() + if err != nil { + return err + } + defer c.Close() + ctx, cancel := a.ctx() + defer cancel() + res, err := c.ByMeta(ctx, req) + if err != nil { + return err + } + if res.GetTruncated() { + fmt.Fprintln(a.stderr, "tokenmeter: warning: the result exceeds the server's row cap and is incomplete; filter by --project") + } + if p.json() { + return p.data(res) + } + var header []string + if req.GroupByMetaKey != "" { + header = append(header, strings.ToUpper(req.GroupByMetaKey)) + } + header = append(header, "VENDOR", "MODEL", "REQUESTS", "TOKENS IN", "TOKENS OUT", "COST USD") + var out [][]string + for _, r := range res.GetRows() { + var line []string + if req.GroupByMetaKey != "" { + line = append(line, r.GetGroupValue()) + } + out = append(out, append(line, r.GetVendor(), r.GetModel(), strconv.FormatInt(r.GetRequestCount(), 10), + strconv.FormatInt(r.GetTokensInput(), 10), strconv.FormatInt(r.GetTokensOutput(), 10), p.usd(r.GetCostNanousd()))) + } + return p.rows(header, out) + }, + } + cmd.Flags().StringVar(&req.MetaKey, "key", "", "meta key, e.g. job_id") + cmd.Flags().StringVar(&req.MetaValue, "value", "", "its value") + cmd.Flags().StringVar(&req.GroupByMetaKey, "group-by", "", "a second meta key to split by, e.g. task_id") + cmd.Flags().StringVar(&req.ProjectId, "project", "", "only this project_id") + _ = cmd.MarkFlagRequired("key") + _ = cmd.MarkFlagRequired("value") + outputFlag(cmd, &format) + return cmd +} + +func (a *app) exportCmd() *cobra.Command { + var q query + var groups []string + var format, out string + cmd := &cobra.Command{ + Use: "export [--format csv|json] [--out FILE] [window and filters] [--group-by DIM,...]", + Short: "Daily usage rows for a spreadsheet or a script; grouped by every dimension by default", + Args: cobra.NoArgs, + RunE: func(*cobra.Command, []string) error { + if format != "csv" && format != "json" { + return usageError{fmt.Errorf("--format %q: want csv or json", format)} + } + if len(groups) == 0 { + groups = dimensions + } + req, err := q.request(groups) + if err != nil { + return err + } + rows, truncated, err := a.fetch(req) + if err != nil { + return err + } + rows = client.Merge(rows, true) + sort.SliceStable(rows, func(i, j int) bool { return rows[i].GetBucket().GetSeconds() < rows[j].GetBucket().GetSeconds() }) + var w io.Writer = a.stdout + if out != "" { + f, err := os.Create(out) + if err != nil { + return err + } + defer f.Close() + w = f + } + p := &printer{w: w, format: format} + if format == "json" { + return p.data(map[string]any{"from": req.GetFrom(), "to": req.GetTo(), "rows": messages(rows), "truncated": truncated}) + } + dims := ordered(groups) + header := []string{"day"} + for _, d := range dims { + header = append(header, d) + } + header = append(header, "requests", "tokens_input", "tokens_output", "tokens_total", "cache_read_tokens", + "cache_write_tokens", "sum_duration_ms", "cost_usd", "cost_nanousd") + var lines [][]string + for _, r := range rows { + line := []string{r.GetBucket().AsTime().UTC().Format(time.DateOnly)} + for _, d := range dims { + line = append(line, dimValue(r, d)) + } + lines = append(lines, append(line, strconv.FormatInt(r.GetRequestCount(), 10), strconv.FormatInt(r.GetTokensInput(), 10), + strconv.FormatInt(r.GetTokensOutput(), 10), strconv.FormatInt(r.GetTokensTotal(), 10), + strconv.FormatInt(r.GetCacheReadTokens(), 10), strconv.FormatInt(r.GetCacheWriteTokens(), 10), + strconv.FormatInt(r.GetSumDurationMs(), 10), exactUSD(r.GetCostNanousd()), strconv.FormatInt(r.GetCostNanousd(), 10))) + } + return p.rows(header, lines) + }, + } + q.flags(cmd) + cmd.Flags().StringSliceVar(&groups, "group-by", nil, "dimensions to keep (default all)") + cmd.Flags().StringVar(&format, "format", "csv", "csv or json") + cmd.Flags().StringVar(&out, "out", "", "write to this file instead of stdout") + _ = cmd.RegisterFlagCompletionFunc("format", cobra.FixedCompletions([]string{"csv", "json"}, cobra.ShellCompDirectiveNoFileComp)) + return cmd +} diff --git a/go.mod b/go.mod index 0ac4ec3..2f949bc 100644 --- a/go.mod +++ b/go.mod @@ -12,7 +12,10 @@ require ( github.com/pressly/goose/v3 v3.24.1 github.com/prometheus/client_golang v1.25.0 github.com/rs/zerolog v1.34.0 + github.com/spf13/cobra v1.10.1 github.com/stretchr/testify v1.12.1 + go.yaml.in/yaml/v3 v3.0.5 + golang.org/x/term v0.47.0 google.golang.org/grpc v1.83.2 google.golang.org/protobuf v1.36.12 ) @@ -25,6 +28,7 @@ require ( github.com/cespare/xxhash/v2 v2.3.0 // indirect github.com/go-faster/city v1.0.1 // indirect github.com/go-faster/errors v0.7.1 // indirect + github.com/inconshreveable/mousetrap v1.1.0 // indirect github.com/jackc/pgpassfile v1.0.0 // indirect github.com/jackc/pgservicefile v0.0.0-20240606120523-5a60cdf6a761 // indirect github.com/jackc/puddle/v2 v2.2.2 // indirect @@ -42,13 +46,13 @@ require ( github.com/segmentio/asm v1.2.1 // indirect github.com/sethvargo/go-retry v0.3.0 // indirect github.com/shopspring/decimal v1.4.0 // indirect + github.com/spf13/pflag v1.0.9 // indirect go.opentelemetry.io/otel v1.46.0 // indirect go.opentelemetry.io/otel/trace v1.46.0 // indirect go.uber.org/multierr v1.11.0 // indirect - go.yaml.in/yaml/v3 v3.0.5 // indirect golang.org/x/net v0.60.0 // indirect golang.org/x/sync v0.23.0 // indirect - golang.org/x/sys v0.48.0 // indirect + golang.org/x/sys v0.49.0 // indirect golang.org/x/text v0.42.0 // indirect google.golang.org/genproto/googleapis/rpc v0.0.0-20260526163538-3dc84a4a5aaa // indirect ) diff --git a/go.sum b/go.sum index 1032bc3..09987ed 100644 --- a/go.sum +++ b/go.sum @@ -11,6 +11,7 @@ github.com/buger/jsonparser v1.1.1/go.mod h1:6RYKKt7H4d4+iWqouImQ9R2FZql3VbhNgx2 github.com/cespare/xxhash/v2 v2.3.0 h1:UL815xU9SqsFlibzuggzjXhog7bL6oX9BbNZnL2UFvs= github.com/cespare/xxhash/v2 v2.3.0/go.mod h1:VGX0DQ3Q6kWi7AoAeZDth3/j3BFtOZR5XLFGgcrjCOs= github.com/coreos/go-systemd/v22 v22.5.0/go.mod h1:Y58oyj3AT4RCenI/lSvhwexgC+NSVTIJ3seZv2GcEnc= +github.com/cpuguy83/go-md2man/v2 v2.0.6/go.mod h1:oOW0eioCTA6cOiMLiUPZOpcVxMig6NIQQ7OS05n1F4g= github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= github.com/dustin/go-humanize v1.0.1 h1:GzkhY7T5VNhEkwH0PVJgjz+fX1rhBrR7pRT3mDkpeCY= github.com/dustin/go-humanize v1.0.1/go.mod h1:Mu1zIs6XwVuF/gI1OepvI0qD18qycQx+mFykh5fBlto= @@ -37,6 +38,8 @@ github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0= github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo= github.com/hashicorp/golang-lru/v2 v2.0.7 h1:a+bsQ5rvGLjzHuww6tVxozPZFVghXaHOwFs4luLUK2k= github.com/hashicorp/golang-lru/v2 v2.0.7/go.mod h1:QeFd9opnmA6QUJc5vARoKUSoFhyfM2/ZepoAG6RGpeM= +github.com/inconshreveable/mousetrap v1.1.0 h1:wN+x4NVGpMsO7ErUn/mUI3vEoE6Jt13X2s0bqwp9tc8= +github.com/inconshreveable/mousetrap v1.1.0/go.mod h1:vpF70FUmC8bwa3OWnCshd2FqLfsEA9PFc4w1p2J65bw= github.com/jackc/pgpassfile v1.0.0 h1:/6Hmqy13Ss2zCq62VdNG8tM1wchn8zjSGOBJ6icpsIM= github.com/jackc/pgpassfile v1.0.0/go.mod h1:CEx0iS5ambNFdcRtxPj5JhEz+xB6uRky5eyVu/W2HEg= github.com/jackc/pgservicefile v0.0.0-20240606120523-5a60cdf6a761 h1:iCEnooe7UlwOQYpKFhBabPMi4aNAfoODPEFNiAnClxo= @@ -87,12 +90,17 @@ github.com/remyoudompheng/bigfft v0.0.0-20230129092748-24d4a6f8daec/go.mod h1:qq github.com/rs/xid v1.6.0/go.mod h1:7XoLgs4eV+QndskICGsho+ADou8ySMSjJKDIan90Nz0= github.com/rs/zerolog v1.34.0 h1:k43nTLIwcTVQAncfCw4KZ2VY6ukYoZaBPNOE8txlOeY= github.com/rs/zerolog v1.34.0/go.mod h1:bJsvje4Z08ROH4Nhs5iH600c3IkWhwp44iRc54W6wYQ= +github.com/russross/blackfriday/v2 v2.1.0/go.mod h1:+Rmxgy9KzJVeS9/2gXHxylqXiyQDYRxCVz55jmeOWTM= github.com/segmentio/asm v1.2.1 h1:DTNbBqs57ioxAD4PrArqftgypG4/qNpXoJx8TVXxPR0= github.com/segmentio/asm v1.2.1/go.mod h1:BqMnlJP91P8d+4ibuonYZw9mfnzI9HfxselHZr5aAcs= github.com/sethvargo/go-retry v0.3.0 h1:EEt31A35QhrcRZtrYFDTBg91cqZVnFL2navjDrah2SE= github.com/sethvargo/go-retry v0.3.0/go.mod h1:mNX17F0C/HguQMyMyJxcnU471gOZGxCLyYaFyAZraas= github.com/shopspring/decimal v1.4.0 h1:bxl37RwXBklmTi0C79JfXCEBD1cqqHt0bbgBAGFp81k= github.com/shopspring/decimal v1.4.0/go.mod h1:gawqmDU56v4yIKSwfBSFip1HdCCXN8/+DMd9qYNcwME= +github.com/spf13/cobra v1.10.1 h1:lJeBwCfmrnXthfAupyUTzJ/J4Nc1RsHC/mSRU2dll/s= +github.com/spf13/cobra v1.10.1/go.mod h1:7SmJGaTHFVBY0jW4NXGluQoLvhqFQM+6XSKD+P4XaB0= +github.com/spf13/pflag v1.0.9 h1:9exaQaMOCwffKiiiYk6/BndUBv+iRViNW+4lEMi0PvY= +github.com/spf13/pflag v1.0.9/go.mod h1:McXfInJRrz4CZXVZOBLb0bTZqETkiAhM9Iw0y3An2Bg= github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME= github.com/stretchr/testify v1.3.0/go.mod h1:M5WIy9Dh21IEIfnGCwXGc5bZfKNJtfHm1UVUgZn+9EI= github.com/stretchr/testify v1.7.0/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg= @@ -127,8 +135,10 @@ golang.org/x/sync v0.23.0/go.mod h1:sUUOizhqBxiL6pEWpqNLUiaJn1ShEbZ6BBqskPbjZm0= golang.org/x/sys v0.0.0-20220811171246-fbc7d0a398ab/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= golang.org/x/sys v0.6.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= golang.org/x/sys v0.12.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= -golang.org/x/sys v0.48.0 h1:bbX/i/6MgT9BVLM9RT1thmxL04yeTAhbEz4SyadbXoo= -golang.org/x/sys v0.48.0/go.mod h1:hNLxWAXmnKAxqDtdwIYC4bM9oQPEecfsnNMuSxOs3og= +golang.org/x/sys v0.49.0 h1:XbzkgYJHdqh/8m2Uu0W/dQv8nktxx4BFHp1M0gROTXA= +golang.org/x/sys v0.49.0/go.mod h1:hNLxWAXmnKAxqDtdwIYC4bM9oQPEecfsnNMuSxOs3og= +golang.org/x/term v0.47.0 h1:cpyOwQkr2Nzxe2p6KnrTG7Ei2wtzjsuSjhpXXepZ/BM= +golang.org/x/term v0.47.0/go.mod h1:efAfAuG6RKABvg1ArGebRkvX6LXul14SM67o2e4mUCY= golang.org/x/text v0.42.0 h1:JbOZXgfeCPU9gacVtYliJqOhD+zhrEqK4LfdpmlUZqI= golang.org/x/text v0.42.0/go.mod h1:ojzP1Z+2QtioaF8DTtO8K5q7JWVVYwZKenzujK0Zd0E= gonum.org/v1/gonum v0.17.0 h1:VbpOemQlsSMrYmn7T2OUvQ4dqxQXU+ouZFQsZOx50z4= @@ -141,6 +151,7 @@ google.golang.org/protobuf v1.36.12 h1:pJOKDDOyeXErUroCihFAd5LQuwXBSpVnKGrj5o/fw google.golang.org/protobuf v1.36.12/go.mod h1:HTf+CrKn2C3g5S8VImy6tdcUvCska2kB7j23XfzDpco= gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= +gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= modernc.org/gc/v3 v3.0.0-20240107210532-573471604cb6 h1:5D53IMaUuA5InSeMu9eJtlQXS2NxAhyWQvkKEgXZhHI= modernc.org/gc/v3 v3.0.0-20240107210532-573471604cb6/go.mod h1:Qz0X07sNOR1jWYCrJMEnbW/X55x206Q7Vt4mz6/wHp4= modernc.org/libc v1.55.3 h1:AzcW1mhlPNrRtjS5sS+eW2ISCgSOLLNyFzRh/V3Qj/U= diff --git a/internal/server/http.go b/internal/server/http.go index 9f7959a..3358497 100644 --- a/internal/server/http.go +++ b/internal/server/http.go @@ -82,10 +82,13 @@ func (s *HTTPServer) Shutdown(ctx context.Context) error { } // Health check handlers +// Version is the build's release version, reported by /healthz. +var Version = "dev" + func (s *HTTPServer) handleHealthz(w http.ResponseWriter, r *http.Request) { w.Header().Set("Content-Type", "application/json") w.WriteHeader(http.StatusOK) - _, _ = w.Write([]byte(`{"status":"ok"}`)) + _ = json.NewEncoder(w).Encode(map[string]string{"status": "ok", "version": Version}) } func (s *HTTPServer) handleReadyz(w http.ResponseWriter, r *http.Request) { diff --git a/pkg/client/client.go b/pkg/client/client.go new file mode 100644 index 0000000..8418c41 --- /dev/null +++ b/pkg/client/client.go @@ -0,0 +1,203 @@ +// Package client is a Go client for light-tokenmeter's gRPC API: recording +// usage events and reading summaries, with large windows fetched in parts so +// a summary is not cut at the server's row cap. +package client + +import ( + "context" + "crypto/tls" + "errors" + "fmt" + "time" + + "google.golang.org/grpc" + "google.golang.org/grpc/credentials" + "google.golang.org/grpc/credentials/insecure" + "google.golang.org/protobuf/proto" + "google.golang.org/protobuf/types/known/timestamppb" + + usagev1 "github.com/ai-process/light-tokenmeter/gen/usage/v1" +) + +// MaxBatch is the most events one RecordUsageEvents call takes. +const MaxBatch = 500 + +// MinWindow is the shortest window Summary splits a truncated one into. +const MinWindow = time.Hour + +// Options say where light-tokenmeter is and how to reach it. +type Options struct { + Addr string // gRPC host:port + Token string // bearer token; empty when the service runs without auth + TLS bool // use TLS + TLSConfig *tls.Config // optional TLS settings; nil means system roots, TLS 1.2+ + // DialOptions are added after the client's own, for tests and custom transports. + DialOptions []grpc.DialOption +} + +// Client talks to one light-tokenmeter. It is safe for concurrent use. +type Client struct { + conn *grpc.ClientConn + svc usagev1.UsageServiceClient +} + +// New connects lazily: the first call dials. +func New(o Options) (*Client, error) { + if o.Addr == "" { + return nil, errors.New("client: Addr is required") + } + creds := insecure.NewCredentials() + if o.TLS { + cfg := o.TLSConfig + if cfg == nil { + cfg = &tls.Config{MinVersion: tls.VersionTLS12} + } + creds = credentials.NewTLS(cfg) + } + opts := []grpc.DialOption{grpc.WithTransportCredentials(creds)} + if o.Token != "" { + opts = append(opts, grpc.WithPerRPCCredentials(bearer{token: o.Token, secure: o.TLS})) + } + conn, err := grpc.NewClient(o.Addr, append(opts, o.DialOptions...)...) + if err != nil { + return nil, err + } + return &Client{conn: conn, svc: usagev1.NewUsageServiceClient(conn)}, nil +} + +// Close releases the connection. +func (c *Client) Close() error { return c.conn.Close() } + +// Service is the generated client, for calls this package does not wrap. +func (c *Client) Service() usagev1.UsageServiceClient { return c.svc } + +type bearer struct { + token string + secure bool +} + +func (b bearer) GetRequestMetadata(context.Context, ...string) (map[string]string, error) { + return map[string]string{"authorization": "Bearer " + b.token}, nil +} +func (b bearer) RequireTransportSecurity() bool { return b.secure } + +// Record sends events in batches of MaxBatch and counts what was new. A +// refused batch stops it; the batches before it stay recorded. +func (c *Client) Record(ctx context.Context, events ...*usagev1.UsageEvent) (created, duplicates int, err error) { + for len(events) > 0 { + n := min(len(events), MaxBatch) + res, err := c.svc.RecordUsageEvents(ctx, &usagev1.RecordUsageEventsRequest{Events: events[:n]}) + if err != nil { + return created, duplicates, err + } + created += int(res.GetCreatedCount()) + duplicates += int(res.GetDuplicateCount()) + events = events[n:] + } + return created, duplicates, nil +} + +// Summary returns the rows of req's window. When the server caps a window +// (truncated), Summary fetches its halves and merges them, down to MinWindow; +// truncated is true only if a window that short is still capped. +func (c *Client) Summary(ctx context.Context, req *usagev1.GetUsageSummaryRequest) (rows []*usagev1.UsageSummaryRow, truncated bool, err error) { + res, err := c.svc.GetUsageSummary(ctx, req) + if err != nil { + return nil, false, err + } + if !res.GetTruncated() { + return res.GetRows(), false, nil + } + from, to := req.GetFrom().AsTime(), req.GetTo().AsTime() + mid := from.Add(to.Sub(from) / 2).Truncate(time.Hour) + if to.Sub(from) <= MinWindow || !mid.After(from) || !mid.Before(to) { + return res.GetRows(), true, nil + } + left := proto.Clone(req).(*usagev1.GetUsageSummaryRequest) + left.To = timestamppb.New(mid) + right := proto.Clone(req).(*usagev1.GetUsageSummaryRequest) + right.From = timestamppb.New(mid) + a, ta, err := c.Summary(ctx, left) + if err != nil { + return nil, false, err + } + b, tb, err := c.Summary(ctx, right) + if err != nil { + return nil, false, err + } + // A day split at mid shows up in both halves; its parts add up. + return Merge(append(a, b...), true), ta || tb, nil +} + +// ByMeta returns what the events tagged meta_key=meta_value used. +func (c *Client) ByMeta(ctx context.Context, req *usagev1.GetUsageByMetaRequest) (*usagev1.GetUsageByMetaResponse, error) { + return c.svc.GetUsageByMeta(ctx, req) +} + +// Merge sums rows that have the same dimensions, and the same day when +// byDay is set; without it every row covers the whole window. The order is +// that of each key's first row. +func Merge(rows []*usagev1.UsageSummaryRow, byDay bool) []*usagev1.UsageSummaryRow { + type key struct { + day int64 + project, user, vendor, model, action string + status usagev1.Status + } + index := map[key]*usagev1.UsageSummaryRow{} + var out []*usagev1.UsageSummaryRow + for _, r := range rows { + k := key{project: r.GetProjectId(), user: r.GetUserId(), vendor: r.GetVendor(), model: r.GetModel(), + action: r.GetActionId(), status: r.GetStatus()} + if byDay { + k.day = r.GetBucket().GetSeconds() + } + m, ok := index[k] + if !ok { + m = proto.Clone(r).(*usagev1.UsageSummaryRow) + if !byDay { + m.Bucket = nil + } + index[k] = m + out = append(out, m) + continue + } + m.TokensInput += r.GetTokensInput() + m.TokensOutput += r.GetTokensOutput() + m.TokensTotal += r.GetTokensTotal() + m.RequestCount += r.GetRequestCount() + m.SumDurationMs += r.GetSumDurationMs() + m.CostNanousd += r.GetCostNanousd() + m.CacheReadTokens += r.GetCacheReadTokens() + m.CacheWriteTokens += r.GetCacheWriteTokens() + } + return out +} + +// USD converts nano-USD, the API's unit of cost, to dollars. +func USD(nanousd int64) float64 { return float64(nanousd) / 1e9 } + +// Window is [from, to) for a duration ending now, in UTC, e.g. Last(7*24*time.Hour). +func Last(d time.Duration) (from, to time.Time) { + to = time.Now().UTC() + return to.Add(-d), to +} + +// ParseDuration reads Go durations plus d (days) and w (weeks): 7d, 2w, 36h. +func ParseDuration(s string) (time.Duration, error) { + if n := len(s); n > 1 && (s[n-1] == 'd' || s[n-1] == 'w') { + var v int + if _, err := fmt.Sscanf(s[:n-1], "%d", &v); err != nil || v <= 0 || fmt.Sprint(v) != s[:n-1] { + return 0, fmt.Errorf("duration %q: want e.g. 7d, 2w or 36h", s) + } + unit := 24 * time.Hour + if s[n-1] == 'w' { + unit *= 7 + } + return time.Duration(v) * unit, nil + } + d, err := time.ParseDuration(s) + if err != nil || d <= 0 { + return 0, fmt.Errorf("duration %q: want e.g. 7d, 2w or 36h", s) + } + return d, nil +} diff --git a/pkg/client/client_test.go b/pkg/client/client_test.go new file mode 100644 index 0000000..b58c2cf --- /dev/null +++ b/pkg/client/client_test.go @@ -0,0 +1,118 @@ +package client + +import ( + "context" + "net" + "testing" + "time" + + "google.golang.org/grpc" + "google.golang.org/grpc/credentials/insecure" + "google.golang.org/grpc/test/bufconn" + "google.golang.org/protobuf/types/known/timestamppb" + + usagev1 "github.com/ai-process/light-tokenmeter/gen/usage/v1" +) + +// capped answers windows up to two hours with one request per hour, and caps +// anything longer, as the server does past its row limit. +type capped struct { + usagev1.UnimplementedUsageServiceServer + calls int +} + +func (c *capped) GetUsageSummary(_ context.Context, req *usagev1.GetUsageSummaryRequest) (*usagev1.GetUsageSummaryResponse, error) { + c.calls++ + from, to := req.GetFrom().AsTime(), req.GetTo().AsTime() + day := timestamppb.New(from.Truncate(24 * time.Hour)) + if to.Sub(from) > 2*time.Hour { + return &usagev1.GetUsageSummaryResponse{Truncated: true, Rows: []*usagev1.UsageSummaryRow{{Bucket: day, ProjectId: "p", RequestCount: 1}}}, nil + } + h := int64(to.Sub(from) / time.Hour) + return &usagev1.GetUsageSummaryResponse{Rows: []*usagev1.UsageSummaryRow{{Bucket: day, ProjectId: "p", RequestCount: h, CostNanousd: h * 10}}}, nil +} + +func serve(t *testing.T, srv usagev1.UsageServiceServer) *Client { + t.Helper() + lis := bufconn.Listen(1 << 20) + g := grpc.NewServer() + usagev1.RegisterUsageServiceServer(g, srv) + go func() { _ = g.Serve(lis) }() + t.Cleanup(g.Stop) + c, err := New(Options{Addr: "passthrough:///tm", DialOptions: []grpc.DialOption{ + grpc.WithContextDialer(func(context.Context, string) (net.Conn, error) { return lis.Dial() }), + grpc.WithTransportCredentials(insecure.NewCredentials()), + }}) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { _ = c.Close() }) + return c +} + +func TestSummarySplitsCappedWindows(t *testing.T) { + fake := &capped{} + c := serve(t, fake) + from := time.Date(2026, 10, 1, 0, 0, 0, 0, time.UTC) + req := &usagev1.GetUsageSummaryRequest{From: timestamppb.New(from), To: timestamppb.New(from.Add(10 * time.Hour))} + rows, truncated, err := c.Summary(context.Background(), req) + if err != nil { + t.Fatal(err) + } + if truncated || len(rows) != 1 || rows[0].GetRequestCount() != 10 || rows[0].GetCostNanousd() != 100 { + t.Fatalf("rows = %v, truncated = %v", rows, truncated) + } + if fake.calls < 3 { + t.Fatalf("calls = %d: the window was not split", fake.calls) + } + + // A window that stays capped at MinWindow is reported, not looped on. + always := &alwaysCapped{} + rows, truncated, err = serve(t, always).Summary(context.Background(), req) + if err != nil || !truncated || len(rows) == 0 { + t.Fatalf("rows = %v, truncated = %v, err = %v", rows, truncated, err) + } +} + +type alwaysCapped struct { + usagev1.UnimplementedUsageServiceServer +} + +func (alwaysCapped) GetUsageSummary(_ context.Context, req *usagev1.GetUsageSummaryRequest) (*usagev1.GetUsageSummaryResponse, error) { + return &usagev1.GetUsageSummaryResponse{Truncated: true, Rows: []*usagev1.UsageSummaryRow{{Bucket: req.GetFrom(), RequestCount: 1}}}, nil +} + +func TestMerge(t *testing.T) { + d1 := timestamppb.New(time.Date(2026, 10, 1, 0, 0, 0, 0, time.UTC)) + d2 := timestamppb.New(time.Date(2026, 10, 2, 0, 0, 0, 0, time.UTC)) + rows := []*usagev1.UsageSummaryRow{ + {Bucket: d1, ProjectId: "a", RequestCount: 1, CostNanousd: 5}, + {Bucket: d2, ProjectId: "a", RequestCount: 2, CostNanousd: 7}, + {Bucket: d1, ProjectId: "b", RequestCount: 4}, + {Bucket: d1, ProjectId: "a", RequestCount: 8, TokensTotal: 3}, + } + daily := Merge(rows, true) + if len(daily) != 3 || daily[0].GetRequestCount() != 9 || daily[0].GetTokensTotal() != 3 { + t.Fatalf("daily = %v", daily) + } + total := Merge(rows, false) + if len(total) != 2 || total[0].GetRequestCount() != 11 || total[0].GetCostNanousd() != 12 || total[0].GetBucket() != nil { + t.Fatalf("total = %v", total) + } + if rows[0].GetRequestCount() != 1 { + t.Fatal("Merge changed its input") + } +} + +func TestParseDuration(t *testing.T) { + for in, want := range map[string]time.Duration{"7d": 7 * 24 * time.Hour, "2w": 14 * 24 * time.Hour, "36h": 36 * time.Hour, "90m": 90 * time.Minute} { + if got, err := ParseDuration(in); err != nil || got != want { + t.Errorf("ParseDuration(%q) = %v, %v", in, got, err) + } + } + for _, in := range []string{"", "d", "0d", "-1d", "1.5d", "7x", "+3d", "-5h"} { + if _, err := ParseDuration(in); err == nil { + t.Errorf("ParseDuration(%q) accepted", in) + } + } +} diff --git a/scripts/cli-e2e.sh b/scripts/cli-e2e.sh new file mode 100755 index 0000000..4ce6457 --- /dev/null +++ b/scripts/cli-e2e.sh @@ -0,0 +1,119 @@ +#!/bin/sh +# The tokenmeter CLI against a built light-tokenmeter, once on Postgres alone +# and once with ClickHouse: tokens, events, summary, top, by-meta, export. +# Env: IMAGE (a built service image; default builds this checkout). +set -eu + +cd "$(dirname "$0")/.." + +P=tm-cli-e2e-$$ +NET=$P-net +IMAGE=${IMAGE:-} +BUILT="" +TMP=$(mktemp -d) + +cleanup() { + docker rm -f -v "$P-svc" "$P-pg" "$P-ch" >/dev/null 2>&1 || true + docker network rm "$NET" >/dev/null 2>&1 || true + [ -z "$BUILT" ] || docker rmi -f "$BUILT" >/dev/null 2>&1 || true + rm -rf "$TMP" +} +trap cleanup EXIT + +fail() { + echo "cli e2e: $*" >&2 + docker logs "$P-svc" 2>&1 | tail -n 40 >&2 || true + exit 1 +} + +# retry N CMD... runs CMD until it succeeds, once a second. +retry() { + n=$1 + shift + i=0 + until "$@"; do + i=$((i + 1)) + [ "$i" -lt "$n" ] || return 1 + sleep 1 + done +} + +if [ -z "$IMAGE" ]; then + IMAGE="$P-image" + BUILT=$IMAGE + docker build -q -t "$IMAGE" . >/dev/null +fi +go build -o "$TMP/tokenmeter" ./cmd/tokenmeter +PATH="$TMP:$PATH" +export TOKENMETER_CONFIG="$TMP/config.yaml" + +tokenmeter token new --name writer --scope write >"$TMP/w" 2>/dev/null +tokenmeter token new --name reader --scope read >"$TMP/r" 2>/dev/null +WRITE=$(sed -n 1p "$TMP/w") +READ=$(sed -n 1p "$TMP/r") +TOKENS="$(sed -n 2p "$TMP/w"),$(sed -n 2p "$TMP/r")" + +# check MODE runs the scenario against a fresh service; MODE is pg or clickhouse. +check() { + echo "=== $1" + docker rm -f -v "$P-svc" "$P-pg" "$P-ch" >/dev/null 2>&1 || true + docker network rm "$NET" >/dev/null 2>&1 || true + docker network create "$NET" >/dev/null + docker run -d --name "$P-pg" --network "$NET" \ + -e POSTGRES_USER=tokenmeter -e POSTGRES_PASSWORD=tokenmeter -e POSTGRES_DB=tokenmeter postgres:17-alpine >/dev/null + ch="" + if [ "$1" = clickhouse ]; then + docker run -d --name "$P-ch" --network "$NET" -e CLICKHOUSE_USER=tokenmeter -e CLICKHOUSE_PASSWORD=tokenmeter \ + -e CLICKHOUSE_DEFAULT_ACCESS_MANAGEMENT=1 clickhouse/clickhouse-server:24.3-alpine >/dev/null + retry 90 docker exec "$P-ch" wget -q --spider http://127.0.0.1:8123/ping >/dev/null 2>&1 || fail "clickhouse did not start" + ch="-e CLICKHOUSE_ENABLED=true -e CLICKHOUSE_ADDR=$P-ch:9000 -e CLICKHOUSE_DATABASE=tokenmeter -e CLICKHOUSE_USER=tokenmeter -e CLICKHOUSE_PASSWORD=tokenmeter" + fi + # Over TCP: the socket answers while the image's init server still runs. + retry 60 docker exec "$P-pg" pg_isready -h 127.0.0.1 -U tokenmeter -d tokenmeter >/dev/null 2>&1 || fail "postgres did not start" + # shellcheck disable=SC2086 # $ch is a list of flags + docker run -d --name "$P-svc" --network "$NET" -p 127.0.0.1::8080 -p 127.0.0.1::9090 \ + -e DATABASE_URL="postgres://tokenmeter:tokenmeter@$P-pg:5432/tokenmeter?sslmode=disable" \ + -e TOKENMETER_TOKENS="$TOKENS" $ch "$IMAGE" >/dev/null + TOKENMETER_ADDR=$(docker port "$P-svc" 9090/tcp | head -n 1) + TOKENMETER_HTTP="http://$(docker port "$P-svc" 8080/tcp | head -n 1)" + export TOKENMETER_ADDR TOKENMETER_HTTP + retry 90 tokenmeter health >/dev/null 2>&1 || fail "the service did not become ready" + tokenmeter health + + echo "--- event send" + export TOKENMETER_TOKEN=$WRITE + at=$(date -u +%Y-%m-%dT%H:%M:%SZ) + send() { + tokenmeter event send --project "$1" --user u1 --action a --vendor openai --model "$2" \ + --tokens-in 100 --tokens-out 50 --cost-usd "$3" --at "$at" --meta job_id=j1 >/dev/null + } + send p1 m1 0.01 + send p1 m2 0.02 + send p2 m1 0.5 + if tokenmeter usage summary >/dev/null 2>&1; then fail "a write token read a summary"; fi + + echo "--- usage summary" + export TOKENMETER_TOKEN=$READ + # With ClickHouse, reads follow the outbox replication. + three() { tokenmeter usage summary --last 1d -o json | grep -q '"request_count": "3"'; } + retry 30 three || { tokenmeter usage summary --last 1d -o json >&2; fail "the summary misses events"; } + tokenmeter usage summary --last 1d --group-by project + # USD on the CLI is cost_nanousd / 1e9 from the API. + tokenmeter usage summary --last 1d --group-by project -o json | grep -q '"cost_nanousd": "30000000"' || fail "p1 cost_nanousd" + tokenmeter usage export --last 1d --group-by project >"$TMP/export.csv" + cat "$TMP/export.csv" + grep -q ',p1,2,200,100,300,0,0,0,0.03,30000000$' "$TMP/export.csv" || fail "export row for p1" + grep -q ',p2,1,100,50,150,0,0,0,0.5,500000000$' "$TMP/export.csv" || fail "export row for p2" + + echo "--- usage top and by-meta" + tokenmeter usage top --last 1d --by cost --dimension project | sed -n 2p | grep -q 'p2' || fail "top is not p2" + tokenmeter usage by-meta --key job_id --value j1 -o json | grep -q '"request_count": "2"' || fail "by-meta for m1" + + code=0 + TOKENMETER_TOKEN=wrong tokenmeter usage summary 2>/dev/null || code=$? + [ "$code" = 4 ] || fail "a wrong token exited $code, want 4" +} + +check pg +check clickhouse +echo "cli e2e: ok" From bec44058dd1d8c670a5a20215c2f0870b57597fe Mon Sep 17 00:00:00 2001 From: Vasil Date: Sat, 24 Oct 2026 02:26:35 +0200 Subject: [PATCH 2/2] tokenmeter CLI: review fixes Private files are written through a fresh temporary file; a context named on the command line beats the environment; keychain writes are read back, as security -i exits 0 on failure; capped windows are fetched with up to eight calls in flight; long durations are refused, not wrapped; exports to a file are renamed into place only when complete; TOKENMETER_TLS must be a boolean. --- README.md | 3 +- cmd/tokenmeter/conn.go | 25 ++++++++++---- cmd/tokenmeter/contexts.go | 25 +++++++++++--- cmd/tokenmeter/errors.go | 25 -------------- cmd/tokenmeter/event_cmd.go | 13 ++----- cmd/tokenmeter/main_test.go | 14 ++++++++ cmd/tokenmeter/output.go | 4 ++- cmd/tokenmeter/usage_cmd.go | 68 +++++++++++++++++++++---------------- pkg/client/client.go | 47 +++++++++++++++++-------- pkg/client/client_test.go | 11 +++--- 10 files changed, 138 insertions(+), 97 deletions(-) diff --git a/README.md b/README.md index c8bb9bc..0c42a15 100644 --- a/README.md +++ b/README.md @@ -248,7 +248,8 @@ tokenmeter completion zsh > "${fpath[1]}/_tokenmeter" # also bash an --token keychain`, then `tokenmeter context store-token prod` reads the token without echo. They live in `~/.config/tokenmeter/config.yaml` (`TOKENMETER_CONFIG`); a token source is `env:VAR`, `file:PATH` or `keychain`. Flags win over `TOKENMETER_ADDR`, `TOKENMETER_TLS`, - `TOKENMETER_HTTP` and `TOKENMETER_TOKEN`, which win over the context. + `TOKENMETER_HTTP` and `TOKENMETER_TOKEN`, which win over the current context; a context named + with `--context` or `TOKENMETER_CONTEXT` wins over them, so its token never goes elsewhere. - Exit codes: 0 done, 1 failed, 2 bad usage, 4 not authenticated or not allowed, 5 not found, 6 rejected as invalid, 7 unavailable (retry later). - [`pkg/client`](pkg/client) is the Go client underneath: `Record` in batches of 500, `Summary` diff --git a/cmd/tokenmeter/conn.go b/cmd/tokenmeter/conn.go index 4bdcd7a..df357b2 100644 --- a/cmd/tokenmeter/conn.go +++ b/cmd/tokenmeter/conn.go @@ -8,6 +8,7 @@ import ( "net" "net/http" "os" + "strconv" "strings" "github.com/ai-process/light-tokenmeter/pkg/client" @@ -25,6 +26,7 @@ type target struct { ctx *Context // the selected context, nil without one ctxName string + envToken string // TOKENMETER_TOKEN, unless a context was asked for by name warn io.Writer resolved bool } @@ -58,15 +60,26 @@ func (t *target) resolve() error { if c == nil { c = &Context{} } - t.addr = firstNonEmpty(t.addr, os.Getenv("TOKENMETER_ADDR"), c.Addr, "localhost:9090") + // A context asked for by name wins over the environment, so its token never + // goes to an address left in the shell, nor the shell's token to its address. + env := os.Getenv + if asked != "" { + env = func(string) string { return "" } + } + t.envToken = env("TOKENMETER_TOKEN") + t.addr = firstNonEmpty(t.addr, env("TOKENMETER_ADDR"), c.Addr, "localhost:9090") if !t.tlsSet { - if v := os.Getenv("TOKENMETER_TLS"); v != "" { - t.tls = v == "true" || v == "1" + if v := env("TOKENMETER_TLS"); v != "" { + b, err := strconv.ParseBool(v) + if err != nil { + return usageError{fmt.Errorf("TOKENMETER_TLS=%q: want true or false", v)} + } + t.tls = b } else { t.tls = c.TLS } } - t.http = firstNonEmpty(t.http, os.Getenv("TOKENMETER_HTTP"), c.HTTP) + t.http = firstNonEmpty(t.http, env("TOKENMETER_HTTP"), c.HTTP) if t.http == "" { host, _, err := net.SplitHostPort(t.addr) if err != nil { @@ -92,8 +105,8 @@ func (t *target) token() (string, error) { if t.tokenFile != "" { return readKeyFile(t.tokenFile) } - if k := os.Getenv("TOKENMETER_TOKEN"); k != "" { - return k, nil + if t.envToken != "" { + return t.envToken, nil } if t.ctx != nil && !t.ctx.Token.empty() { k, err := t.ctx.Token.read(t.ctxName) diff --git a/cmd/tokenmeter/contexts.go b/cmd/tokenmeter/contexts.go index 4873e8c..02fa86f 100644 --- a/cmd/tokenmeter/contexts.go +++ b/cmd/tokenmeter/contexts.go @@ -163,13 +163,26 @@ func (c *Contexts) names() []string { return out } -// writePrivate writes a file only its owner can read, replacing it whole. +// writePrivate writes a file only its owner can read, replacing it whole. A +// fresh temporary file keeps a leftover or planted one from setting the mode. func writePrivate(path string, b []byte) error { - tmp := path + ".tmp" - if err := os.WriteFile(tmp, b, 0o600); err != nil { + f, err := os.CreateTemp(filepath.Dir(path), "."+filepath.Base(path)+".*") + if err != nil { + return err + } + defer os.Remove(f.Name()) + if err := f.Chmod(0o600); err != nil { + f.Close() + return err + } + if _, err := f.Write(b); err != nil { + f.Close() + return err + } + if err := f.Close(); err != nil { return err } - return os.Rename(tmp, path) + return os.Rename(f.Name(), path) } func expandHome(p string) string { @@ -226,6 +239,10 @@ func keychainSet(account, key string) error { if err := cmd.Run(); err != nil { return fmt.Errorf("keychain: %v %s", err, strings.TrimSpace(errb.String())) } + // security -i exits 0 even when the command in it failed, so read it back. + if got, err := keychainGet(account); err != nil || got != key { + return fmt.Errorf("keychain did not store the entry %s/%s %s", keychainService, account, strings.TrimSpace(errb.String())) + } return nil } diff --git a/cmd/tokenmeter/errors.go b/cmd/tokenmeter/errors.go index 120d34a..92dbd4a 100644 --- a/cmd/tokenmeter/errors.go +++ b/cmd/tokenmeter/errors.go @@ -5,7 +5,6 @@ import ( "errors" "fmt" "net" - "net/http" "google.golang.org/grpc/codes" "google.golang.org/grpc/status" @@ -28,16 +27,6 @@ const exitCodesHelp = `Exit codes: 0 done, 1 failed, 2 bad usage, type usageError struct{ error } -// httpError is a non-2xx answer from the HTTP API. -type httpError struct { - status int - body string -} - -func (e httpError) Error() string { - return fmt.Sprintf("HTTP %d: %s", e.status, e.body) -} - func exitCode(err error) int { if err == nil { return exitOK @@ -46,20 +35,6 @@ func exitCode(err error) int { if errors.As(err, &u) { return exitUsage } - var h httpError - if errors.As(err, &h) { - switch { - case h.status == http.StatusUnauthorized || h.status == http.StatusForbidden: - return exitAuth - case h.status == http.StatusNotFound: - return exitNotFound - case h.status == http.StatusTooManyRequests || h.status >= 500: - return exitRetry - case h.status >= 400: - return exitRejected - } - return exitFailed - } var netErr net.Error if errors.As(err, &netErr) || errors.Is(err, context.DeadlineExceeded) { return exitRetry diff --git a/cmd/tokenmeter/event_cmd.go b/cmd/tokenmeter/event_cmd.go index 97a9886..1d22ff1 100644 --- a/cmd/tokenmeter/event_cmd.go +++ b/cmd/tokenmeter/event_cmd.go @@ -1,12 +1,12 @@ package main import ( - "crypto/rand" "fmt" "math" "strconv" "time" + "github.com/google/uuid" "github.com/spf13/cobra" "google.golang.org/protobuf/types/known/timestamppb" @@ -29,7 +29,7 @@ func (a *app) eventCmd() *cobra.Command { } ev.Status = s if ev.EventId == "" { - ev.EventId = newUUID() + ev.EventId = uuid.NewString() } if cost != "" { usd, err := strconv.ParseFloat(cost, 64) @@ -86,12 +86,3 @@ func (a *app) eventCmd() *cobra.Command { cmd.AddCommand(send) return cmd } - -// newUUID is a random (version 4) UUID. -func newUUID() string { - b := make([]byte, 16) - _, _ = rand.Read(b) - b[6] = b[6]&0x0f | 0x40 - b[8] = b[8]&0x3f | 0x80 - return fmt.Sprintf("%x-%x-%x-%x-%x", b[0:4], b[4:6], b[6:8], b[8:10], b[10:]) -} diff --git a/cmd/tokenmeter/main_test.go b/cmd/tokenmeter/main_test.go index afde1d5..64f6b68 100644 --- a/cmd/tokenmeter/main_test.go +++ b/cmd/tokenmeter/main_test.go @@ -123,6 +123,20 @@ func TestContexts(t *testing.T) { if _, out, _ := call("context", "list"); !strings.Contains(out, "* prod") { t.Fatalf("context list: %s", out) } + + // A context asked for by name beats the shell's address and token. + t.Setenv("TOKENMETER_ADDR", "localhost:9090") + tg = target{context: "prod"} + if tok, _ := tg.token(); tok != "tok-file" || tg.addr != "tm.example.com:443" { + t.Fatalf("named context: addr %s token %q", tg.addr, tok) + } + t.Setenv("TOKENMETER_TLS", "yes please") + if code, _, _ := call("usage", "summary"); code != exitUsage { + t.Fatalf("bad TOKENMETER_TLS = %d, want %d", code, exitUsage) + } + if info, err := os.Stat(filepath.Join(filepath.Dir(p), "tokenmeter", "prod.token")); err != nil || info.Mode().Perm() != 0o600 { + t.Fatalf("token file: %v %v", info.Mode(), err) + } } // serve runs the real gRPC and HTTP servers over a fresh database with these diff --git a/cmd/tokenmeter/output.go b/cmd/tokenmeter/output.go index 8bd03c1..0603b83 100644 --- a/cmd/tokenmeter/output.go +++ b/cmd/tokenmeter/output.go @@ -12,6 +12,8 @@ import ( "google.golang.org/protobuf/encoding/protojson" "google.golang.org/protobuf/proto" + + "github.com/ai-process/light-tokenmeter/pkg/client" ) // printer writes a result as a table for people, or as JSON or CSV for @@ -120,7 +122,7 @@ func (p *printer) usd(nano int64) string { if p.format == "csv" { return exactUSD(nano) } - return fmt.Sprintf("$%.4f", float64(nano)/1e9) + return fmt.Sprintf("$%.4f", client.USD(nano)) } // exactUSD writes nano-USD as dollars without float rounding: 36000000 is 0.036. diff --git a/cmd/tokenmeter/usage_cmd.go b/cmd/tokenmeter/usage_cmd.go index 97dbc17..a402c52 100644 --- a/cmd/tokenmeter/usage_cmd.go +++ b/cmd/tokenmeter/usage_cmd.go @@ -3,8 +3,8 @@ package main import ( "errors" "fmt" - "io" "os" + "path/filepath" "sort" "strconv" "strings" @@ -411,38 +411,23 @@ func (a *app) exportCmd() *cobra.Command { } rows = client.Merge(rows, true) sort.SliceStable(rows, func(i, j int) bool { return rows[i].GetBucket().GetSeconds() < rows[j].GetBucket().GetSeconds() }) - var w io.Writer = a.stdout - if out != "" { - f, err := os.Create(out) - if err != nil { - return err - } - defer f.Close() - w = f + if out == "" { + return writeExport(&printer{w: a.stdout, format: format}, req, groups, rows, truncated) } - p := &printer{w: w, format: format} - if format == "json" { - return p.data(map[string]any{"from": req.GetFrom(), "to": req.GetTo(), "rows": messages(rows), "truncated": truncated}) + // Into a temporary file renamed at the end, so a failed write leaves no partial export. + f, err := os.CreateTemp(filepath.Dir(out), "."+filepath.Base(out)+".*") + if err != nil { + return err } - dims := ordered(groups) - header := []string{"day"} - for _, d := range dims { - header = append(header, d) + defer os.Remove(f.Name()) + if err := writeExport(&printer{w: f, format: format}, req, groups, rows, truncated); err != nil { + f.Close() + return err } - header = append(header, "requests", "tokens_input", "tokens_output", "tokens_total", "cache_read_tokens", - "cache_write_tokens", "sum_duration_ms", "cost_usd", "cost_nanousd") - var lines [][]string - for _, r := range rows { - line := []string{r.GetBucket().AsTime().UTC().Format(time.DateOnly)} - for _, d := range dims { - line = append(line, dimValue(r, d)) - } - lines = append(lines, append(line, strconv.FormatInt(r.GetRequestCount(), 10), strconv.FormatInt(r.GetTokensInput(), 10), - strconv.FormatInt(r.GetTokensOutput(), 10), strconv.FormatInt(r.GetTokensTotal(), 10), - strconv.FormatInt(r.GetCacheReadTokens(), 10), strconv.FormatInt(r.GetCacheWriteTokens(), 10), - strconv.FormatInt(r.GetSumDurationMs(), 10), exactUSD(r.GetCostNanousd()), strconv.FormatInt(r.GetCostNanousd(), 10))) + if err := f.Close(); err != nil { + return err } - return p.rows(header, lines) + return os.Rename(f.Name(), out) }, } q.flags(cmd) @@ -452,3 +437,28 @@ func (a *app) exportCmd() *cobra.Command { _ = cmd.RegisterFlagCompletionFunc("format", cobra.FixedCompletions([]string{"csv", "json"}, cobra.ShellCompDirectiveNoFileComp)) return cmd } + +func writeExport(p *printer, req *usagev1.GetUsageSummaryRequest, groups []string, rows []*usagev1.UsageSummaryRow, truncated bool) error { + if p.format == "json" { + return p.data(map[string]any{"from": req.GetFrom(), "to": req.GetTo(), "rows": messages(rows), "truncated": truncated}) + } + dims := ordered(groups) + header := []string{"day"} + for _, d := range dims { + header = append(header, d) + } + header = append(header, "requests", "tokens_input", "tokens_output", "tokens_total", "cache_read_tokens", + "cache_write_tokens", "sum_duration_ms", "cost_usd", "cost_nanousd") + var lines [][]string + for _, r := range rows { + line := []string{r.GetBucket().AsTime().UTC().Format(time.DateOnly)} + for _, d := range dims { + line = append(line, dimValue(r, d)) + } + lines = append(lines, append(line, strconv.FormatInt(r.GetRequestCount(), 10), strconv.FormatInt(r.GetTokensInput(), 10), + strconv.FormatInt(r.GetTokensOutput(), 10), strconv.FormatInt(r.GetTokensTotal(), 10), + strconv.FormatInt(r.GetCacheReadTokens(), 10), strconv.FormatInt(r.GetCacheWriteTokens(), 10), + strconv.FormatInt(r.GetSumDurationMs(), 10), exactUSD(r.GetCostNanousd()), strconv.FormatInt(r.GetCostNanousd(), 10))) + } + return p.rows(header, lines) +} diff --git a/pkg/client/client.go b/pkg/client/client.go index 8418c41..c3f4341 100644 --- a/pkg/client/client.go +++ b/pkg/client/client.go @@ -8,6 +8,7 @@ import ( "crypto/tls" "errors" "fmt" + "math" "time" "google.golang.org/grpc" @@ -25,6 +26,9 @@ const MaxBatch = 500 // MinWindow is the shortest window Summary splits a truncated one into. const MinWindow = time.Hour +// parallel caps the summary calls one Client has in flight while splitting. +const parallel = 8 + // Options say where light-tokenmeter is and how to reach it. type Options struct { Addr string // gRPC host:port @@ -37,8 +41,9 @@ type Options struct { // Client talks to one light-tokenmeter. It is safe for concurrent use. type Client struct { - conn *grpc.ClientConn - svc usagev1.UsageServiceClient + conn *grpc.ClientConn + svc usagev1.UsageServiceClient + slots chan struct{} } // New connects lazily: the first call dials. @@ -62,7 +67,7 @@ func New(o Options) (*Client, error) { if err != nil { return nil, err } - return &Client{conn: conn, svc: usagev1.NewUsageServiceClient(conn)}, nil + return &Client{conn: conn, svc: usagev1.NewUsageServiceClient(conn), slots: make(chan struct{}, parallel)}, nil } // Close releases the connection. @@ -101,7 +106,15 @@ func (c *Client) Record(ctx context.Context, events ...*usagev1.UsageEvent) (cre // (truncated), Summary fetches its halves and merges them, down to MinWindow; // truncated is true only if a window that short is still capped. func (c *Client) Summary(ctx context.Context, req *usagev1.GetUsageSummaryRequest) (rows []*usagev1.UsageSummaryRow, truncated bool, err error) { + // A slot is held for the call only, never across the recursion, so halves + // waiting for their own halves cannot starve each other. + select { + case c.slots <- struct{}{}: + case <-ctx.Done(): + return nil, false, ctx.Err() + } res, err := c.svc.GetUsageSummary(ctx, req) + <-c.slots if err != nil { return nil, false, err } @@ -117,12 +130,19 @@ func (c *Client) Summary(ctx context.Context, req *usagev1.GetUsageSummaryReques left.To = timestamppb.New(mid) right := proto.Clone(req).(*usagev1.GetUsageSummaryRequest) right.From = timestamppb.New(mid) - a, ta, err := c.Summary(ctx, left) - if err != nil { - return nil, false, err - } - b, tb, err := c.Summary(ctx, right) - if err != nil { + // The halves are independent; fetching them together keeps a deep split + // within one deadline. + var b []*usagev1.UsageSummaryRow + var tb bool + var errB error + done := make(chan struct{}) + go func() { + defer close(done) + b, tb, errB = c.Summary(ctx, right) + }() + a, ta, errA := c.Summary(ctx, left) + <-done + if err := errors.Join(errA, errB); err != nil { return nil, false, err } // A day split at mid shows up in both halves; its parts add up. @@ -176,12 +196,6 @@ func Merge(rows []*usagev1.UsageSummaryRow, byDay bool) []*usagev1.UsageSummaryR // USD converts nano-USD, the API's unit of cost, to dollars. func USD(nanousd int64) float64 { return float64(nanousd) / 1e9 } -// Window is [from, to) for a duration ending now, in UTC, e.g. Last(7*24*time.Hour). -func Last(d time.Duration) (from, to time.Time) { - to = time.Now().UTC() - return to.Add(-d), to -} - // ParseDuration reads Go durations plus d (days) and w (weeks): 7d, 2w, 36h. func ParseDuration(s string) (time.Duration, error) { if n := len(s); n > 1 && (s[n-1] == 'd' || s[n-1] == 'w') { @@ -193,6 +207,9 @@ func ParseDuration(s string) (time.Duration, error) { if s[n-1] == 'w' { unit *= 7 } + if v > int(math.MaxInt64/int64(unit)) { + return 0, fmt.Errorf("duration %q is too long", s) + } return time.Duration(v) * unit, nil } d, err := time.ParseDuration(s) diff --git a/pkg/client/client_test.go b/pkg/client/client_test.go index b58c2cf..54eed62 100644 --- a/pkg/client/client_test.go +++ b/pkg/client/client_test.go @@ -3,6 +3,7 @@ package client import ( "context" "net" + "sync/atomic" "testing" "time" @@ -18,11 +19,11 @@ import ( // anything longer, as the server does past its row limit. type capped struct { usagev1.UnimplementedUsageServiceServer - calls int + calls atomic.Int32 } func (c *capped) GetUsageSummary(_ context.Context, req *usagev1.GetUsageSummaryRequest) (*usagev1.GetUsageSummaryResponse, error) { - c.calls++ + c.calls.Add(1) from, to := req.GetFrom().AsTime(), req.GetTo().AsTime() day := timestamppb.New(from.Truncate(24 * time.Hour)) if to.Sub(from) > 2*time.Hour { @@ -62,8 +63,8 @@ func TestSummarySplitsCappedWindows(t *testing.T) { if truncated || len(rows) != 1 || rows[0].GetRequestCount() != 10 || rows[0].GetCostNanousd() != 100 { t.Fatalf("rows = %v, truncated = %v", rows, truncated) } - if fake.calls < 3 { - t.Fatalf("calls = %d: the window was not split", fake.calls) + if n := fake.calls.Load(); n < 3 { + t.Fatalf("calls = %d: the window was not split", n) } // A window that stays capped at MinWindow is reported, not looped on. @@ -110,7 +111,7 @@ func TestParseDuration(t *testing.T) { t.Errorf("ParseDuration(%q) = %v, %v", in, got, err) } } - for _, in := range []string{"", "d", "0d", "-1d", "1.5d", "7x", "+3d", "-5h"} { + for _, in := range []string{"", "d", "0d", "-1d", "1.5d", "7x", "+3d", "-5h", "200000d", "99999999999w"} { if _, err := ParseDuration(in); err == nil { t.Errorf("ParseDuration(%q) accepted", in) }