A couple of the challenges running AI applications at scale are managing cost, and ensuring content is compliant. LLM calls and token costs can rapidly spin out of control, and prompts are sent to the LLMs without filtering, a potential risk with external providers. Go’s standard libraries provide sophisticated networking capability, including an elegant reverse proxy implementation. This seemed like a good fit for these problems.

In this post, we will use reverse proxy to create an endpoint for LLM calls, that intercepts and inspects incoming prompts, before forwarding these largely unchanged to the LLM provider if they pass inspection, or block them if not. For those that pass through, the LLM reply is again intercepted by the same reverse proxy, and token counts are extracted from the LLM completions. These counts can be aggregated, and high-usage accounts can be quickly flagged and even blocked.

In Go, this capability is relatively easy to implement, and is very fast. In this blog post, we will use Go’s standard reverse proxy implementation to inspect LLM calls, and count and report on token values.

The basic reverse proxy in Go

There are two kinds of proxy, a forward proxy and a reverse proxy.

• A forward proxy sits in front of clients to manage outbound requests to the internet, and thus can disguise the client’s IP address, to get around regional or IP specific browsing restrictions or filters, provide caching, or serve other purposes.

• A reverse proxy, on the other hand, sits in front of backend server(s) and receives requests on their behalf; clients talk to it as though it were the server itself, but the reverse proxy can intercept and inspect their traffic, or the responses received, and act upon these.

I found this a concise and helpful example of a reverse proxy in Go. It creates an internal endpoint that does nothing with the traffic, just sends it through to another destination (in this case Google) unchanged:

package main

import(
    "log"
    "net/url"
    "net/http"
    "net/http/httputil"
)

func main() {

    remote, err := url.Parse("http://google.com")
    if err != nil {
        panic(err)
    }

    handler := func(p *httputil.ReverseProxy) func(http.ResponseWriter, *http.Request) {
        return func(w http.ResponseWriter, r *http.Request) {
            log.Println(r.URL)
            r.Host = remote.Host
            p.ServeHTTP(w, r)
        }
    }
        
    proxy := httputil.NewSingleHostReverseProxy(remote)
    http.HandleFunc("/", handler(proxy))
    err = http.ListenAndServe(":8080", nil)
    if err != nil {
        panic(err)
    }
}

The standard library documentation also includes a slightly longer example.

The core line of this short program is this:

proxy := httputil.NewSingleHostReverseProxy(remote)

This builds a standard library ReverseProxy, which rewrites each incoming request so its host points at remote, sends it upstream, and copies the response (status, headers, body) back to the original client. Single host means it always forwards to the single same backend, with no load balancing.

You could use this ReverseProxy object directly, but for our purposes it makes sense to wrap it inside a function, which lets you inspect and block requests.

handler := func(p *httputil.ReverseProxy) func(http.ResponseWriter, *http.Request) {
    return func(w http.ResponseWriter, r *http.Request) {
        // wrapper contents here
        r.Header.Del("Accept-Encoding")
        r.Host = remote.Host
        p.ServeHTTP(w, r)
    }
}

Inside the wrapper, you need to change the header from the original request (e.g., http://localhost:8080), and change it to the destination URL. Otherwise, even if it gets to the destination, that server would see the header Host: localhost:8080 and reject the request.

The wrapper can then execute p.ServeHTTP(w, r), which sends the modified request to the reverse proxy, which performs the upstream call and writes the response to w.

Note the wrapper also removes the “Accept-Encoding” header, since requests from the OpenAI SDK send gzip content, which are automatically decompressed by Go, so this is no longer needed.

Using this to inspect incoming prompts

We can adapt this structure to intercept non-streaming API calls to OpenAI, and inspect the prompts and token counts.

To inspect the prompt, add a check inside the wrapper. This should grab the bytes in the request and look at them, e.g., look for regular expressions. In this example, we just do a substring search for “restricted” and reject the request if found:

// Read the whole request body, to check for restricted content
body, err := io.ReadAll(r.Body)
r.Body.Close()
if err != nil {
    http.Error(w, "could not read request body", http.StatusBadRequest)
    return
}

// Block the request if the body contains restricted content
if bytes.Contains(body, []byte("restricted")) {
    log.Printf("blocked request to %s: restricted content", r.URL)
    http.Error(w, "request blocked: restricted content", http.StatusForbidden)
    return
}

// Put a copy of the body back, so it can be forwarded
r.Body = io.NopCloser(bytes.NewReader(body))
r.ContentLength = int64(len(body))

Of course, you can do much more extensive checks, such as those in my recent post on in-browser PII redaction. And you can also handle the transgression more gracefully, by returning a status message, for example, as well as logging the error and possibly blocking the user.

Using this to monitor token costs

You need a different approach to inspect responses from the server being called, e.g., to inspect the token counts returned from OpenAI API calls.

To do this, first define a function that takes a response body, reads it, and performs the necessary checks (note that this does not work for streaming responses, but that can be changed):

// Handler that gets called for every response coming back from OpenAI
func trackUsage(resp *http.Response) error {

	// For now, streaming responses (SSE) are passed through
	// TODO: add support for streaming
	if strings.HasPrefix(resp.Header.Get("Content-Type"), "text/event-stream") {
		return nil
	}

	// Read the whole body
	body, err := io.ReadAll(resp.Body)
	resp.Body.Close()
	if err != nil {
		return err
	}

	// Put the body back into the response. Note that NopCloser
    // just adds a Close() to the NewReader, so it satisfies io.ReadCloser
	resp.Body = io.NopCloser(bytes.NewReader(body))
	resp.ContentLength = int64(len(body))
	resp.Header.Set("Content-Length", strconv.Itoa(len(body)))

	// Extract the usage block, if there is one
	var parsed struct {
		Model string `json:"model"`
		Usage *Usage `json:"usage"`
	}
	if json.Unmarshal(body, &parsed) != nil || parsed.Usage == nil {
		return nil
	}
	u := parsed.Usage
	in := u.PromptTokens + u.InputTokens
	out := u.CompletionTokens + u.OutputTokens

	// Get the API key from the request, used as key for accumulating counts
	key := maskKey(resp.Request.Header.Get("Authorization"))

	// Get or create the counter for this API key
	mu.Lock()
	t, ok := totals[key]
	if !ok {
		t = &Totals{}
		totals[key] = t
	}

	// Update and log the counts
	t.Requests++
	t.InputTokens += in
	t.OutputTokens += out
	t.TotalTokens += u.TotalTokens
	log.Printf("key=%s model=%s in=%d out=%d total=%d | running: requests=%d total=%d",
		key, parsed.Model, in, out, u.TotalTokens, t.Requests, t.TotalTokens)
	mu.Unlock()

	// No error
	return nil
}

Then, when you create the proxy, register this function by setting ModifyResponse:

proxy := httputil.NewSingleHostReverseProxy(remote)
proxy.ModifyResponse = trackUsage
http.HandleFunc("/", handler(proxy))

The trackUsage function gets called on every response, and can inspect and log values.

Testing it

To test this, build and run:

go build
./tokengate

There is a simple test program in the respository that just makes a completion call to OpenAI, but calls the local address of the reverse proxy instead of the remote LLM:

# Minimal OpenAI chat completion script, using OpenAI library.
# Set BASE_URL to point at the real OpenAI API or to a local reverse proxy.

import sys
import openai
from openai import OpenAI

# Point this at the local reverse proxy (e.g. "http://localhost:8000/v1")
# or at OpenAI directly ("https://api.openai.com/v1").
#BASE_URL = "https://api.openai.com/v1"
BASE_URL = "http://localhost:8000/v1"

# Cheapest model served by OpenAI's own API.
MODEL = "gpt-5-nano"

def chat_completion(prompt: str):

    # Create the client; assumes API key is defined in env OPENAI_API_KEY
    client = OpenAI(base_url=BASE_URL, timeout=60)

    # Attempt the request
    response = client.chat.completions.create(
        model=MODEL,
        messages=[{"role": "user", "content": prompt}],
    )

    # Return the full response, not just the answer
    return response

# Get prompt from command line, or use default
prompt = sys.argv[1] if len(sys.argv) == 2 else "What is a reverse proxy?"
print('Prompt:', prompt)

# Call the completion and split out answer and tokens
resp = chat_completion(prompt)
ans = resp.choices[0].message.content

# Show answer
#print(resp.model_dump())
print('Answer:', ans)
print('Usage:', resp.usage)

Open another console window, install openai in a new Python environment, set your OpenAI key as an environment variable and run the script with a prompt:

uv sync
source .venv/bin/activate
source my_openai_key.txt
python chat_completion.py "Recommend a solid breakfast"

Note the token counts that appear in the tokengate terminal.

Try with another prompt that triggers the prompt rejection:

python chat_completion.py "What do people mean by a restricted diet?"

Conclusion

The code snippets above are abbreviated from a slightly longer file, but the whole thing is less than 200 lines including comments and blank lines. Complete code is here.