Skip to content

Latest commit

 

History

1 Commit

Folders and files

NameName
Last commit message
Last commit date
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 

Repository files navigation

tsgate

An identity-aware LLM gateway that runs on a tailnet instead of on a public port.

Callers reach it over Tailscale and never hold the upstream API key. The gateway resolves who they are from the connection, decides whether that identity may call that model, forwards the request, streams the answer back, and writes one audit line per request — including the ones that fail and the ones the caller abandons.

    caller on the tailnet                      tsgate                        model API
    ─────────────────────      ┌──────────────────────────────────┐      ───────────────
    POST /v1/chat/completions  │ WhoIs(conn) ─▶ policy ─▶ forward │      OpenAI-compatible
    (no API key, no public IP) │      │           │          │    │      streaming endpoint
                               │      └───────────┴──────────┴──▶ audit + metrics
                               └──────────────────────────────────┘

Why it exists

Two problems show up the moment a team puts a shared model key behind an internal tool:

  1. The key leaks by being useful. Every service that needs the model needs the key, and a key that lives in five places lives in a screenshot eventually. Here the key exists in one process, injected from the environment, and callers authenticate by being on the tailnet — there is nothing to hand out and nothing to rotate across consumers.

  2. Nobody can say who spent what. Usage arrives as one bill. This gateway attributes every request to a tailnet identity and enforces a per-identity hourly budget.

The third problem is the one that motivated the design, and it is the interesting one.

The bug this is built around

A client hanging up on a streaming LLM response does not, by itself, stop anything upstream.

Write the obvious proxy — read the body, http.NewRequest, copy the response — and when the caller closes the tab, the model keeps generating. The tokens are produced, the account is billed, and nothing anywhere logs an error, because from the proxy's point of view nothing went wrong. The only trace is a bill that does not match observed traffic.

Two things fix it, and both are in internal/proxy/proxy.go:

// 1. The inbound request's context is threaded into the upstream call.
//    This single argument is what makes a hang-up stop the generation.
upReq, err := http.NewRequestWithContext(ctx, http.MethodPost, url, bytes.NewReader(body))
// 2. The copy loop checks cancellation before every read/write pair, rather
//    than waiting to discover a write error.
select {
case <-ctx.Done():
    return bytesOut, frames, ctx.Err()
default:
}

The second one matters more than it looks. Writes to a departed client keep succeeding for a while — into the http.ResponseWriter buffer, into the kernel socket buffer — so "the write failed" is a late and unreliable way to learn the reader is gone. Waiting for it means paying for tokens after the fact.

There is an accounting half too: usage is recorded on abort. The tokens generated before the caller left were still generated. policy.Check and policy.Record are separate calls for exactly this reason — you cannot know a request's cost in advance, and the aborted requests are the ones most likely to be dropped from a naive accounting path.

Both behaviours have regression tests: TestClientHangupCancelsUpstream and TestUsageIsRecordedOnAbort.

Identity comes from the connection

tsgate never reads identity out of the request. No header, no bearer token, no query parameter — those are all things a caller controls. Identity comes from tsnet's WhoIs on the connection's remote address, and the caller's own Authorization header is dropped rather than forwarded, so a caller cannot smuggle their own key past the gateway's policy.

The rule is enforced structurally: internal/identity defines a Resolver interface that only receives a remote address. There is no request object in scope to misuse.

Layout

cmd/tsgate/main.go        tsnet wiring, config, graceful shutdown — the only file importing tsnet
internal/identity/        Resolver interface; identity from the connection
internal/policy/          per-identity model allow-list and rolling hourly token budget
internal/proxy/           the request path: resolve, check, forward, stream, account
internal/audit/           one JSON line per request, including failures and aborts
internal/metrics/         request / latency / error-rate / abort counters behind /metrics
docs/RUNBOOK.md           what to do when it misbehaves

tsnet is imported in exactly one file. Everything under internal/ depends only on the standard library, which is why the entire request path — cancellation, policy, streaming, accounting — is testable with net/http/httptest and no tailnet at all.

Running it

go mod tidy
go test ./...

cp config.example.json config.json    # edit rules and upstream_base_url
export TSGATE_UPSTREAM_API_KEY=sk-...
export TS_AUTHKEY=tskey-auth-...      # first run only; tsnet stores state in state_dir
go run ./cmd/tsgate -config config.json

Then, from any device on the tailnet:

curl -N http://tsgate/v1/chat/completions \
  -H 'Content-Type: application/json' \
  -d '{"model":"gpt-4o-mini","stream":true,"messages":[{"role":"user","content":"hi"}]}'

No API key on the client. -N disables curl's buffering so you can see the stream arrive incrementally — which is also the quickest way to confirm the gateway is flushing per frame.

Operational surface

Endpoint Purpose
/healthz liveness; returns ok
/metrics JSON: requests, error rate, abort rate, mean latency, tokens, status breakdown
anything else proxied to upstream_base_url with the path preserved

Abort rate is on the dashboard deliberately. It is the number that goes up when clients are timing out or users are giving up mid-answer, and it is invisible in an error rate because nothing errored.

Audit lines look like this:

{"time":"2026-09-04T12:00:03Z","login":"alice@example.com","node":"alice-laptop",
 "model":"gpt-4o-mini","outcome":"client_gone_mid_stream","status":499,
 "tokens":37,"bytes":2214,"duration_ms":1840.2,"aborted":true}

outcome is the gateway's own classification, not the HTTP status, because several distinct situations share a status. A 403 from deny_model and a 403 from deny_unknown_identity mean very different things at 3am.

Known limits

  • Token counts are approximate. The gateway counts SSE data frames; an OpenAI-style chat stream emits roughly one frame per token. Exact accounting means parsing every frame on the hot path to find a usage block that arrives only at the end — and never arrives at all on the aborted requests that most need charging. An approximation available on every path beat an exact number available only on the happy one.
  • Budgets are per-process. Running two replicas gives each its own counter. Shared state would mean a datastore, and this is a gateway for a tailnet, not a billing system.
  • Non-streaming responses are proxied but not parsed. Only the model field is inspected.

License

MIT — see LICENSE.

About

Identity-aware LLM gateway that runs on a tailnet instead of a public port. Callers are identified by the connection (tsnet WhoIs), never by a header, and never hold the upstream API key.

Topics

Resources

Stars

0 stars

Watchers

0 watching

Forks

Releases

Packages

Contributors

Languages