Extending TokenHub

This guide covers common extension points for adding functionality to TokenHub.

Adding a New Provider

Create the adapter package:

internal/providers/newprovider/
├── adapter.go      # Sender implementation
└── adapter_test.go # Tests

Implement the interfaces:

package newprovider

type Adapter struct {
    id      string
    apiKey  string
    baseURL string
    client  *http.Client
}

// Required: router.Sender
func (a *Adapter) ID() string { return a.id }
func (a *Adapter) Send(ctx context.Context, model string, req router.Request) (router.ProviderResponse, error) { ... }
func (a *Adapter) ClassifyError(err error) *router.ClassifiedError { ... }

// Optional: router.StreamSender
func (a *Adapter) SendStream(ctx context.Context, model string, req router.Request) (io.ReadCloser, error) { ... }

// Optional: health.Probeable
func (a *Adapter) HealthEndpoint() string { return a.baseURL + "/health" }

Register via the admin API (providers and models are registered at runtime, not compiled in):

curl -X POST http://localhost:8080/admin/v1/providers \
  -H "Authorization: Bearer $ADMIN_TOKEN" \
  -H "Content-Type: application/json" \
  -d '{"id":"newprovider","type":"openai","base_url":"https://api.newprovider.com","api_key":"..."}'

Add adapter construction in registerProviderAdapter() in handlers_admin.go:

case "newprovider":
    d.Engine.RegisterAdapter(newprovider.New(p.ID, apiKey, p.BaseURL, newprovider.WithTimeout(timeout)))

Adding a New Routing Mode

Define the weight profile in internal/router/engine.go:

var modeWeights = map[string]weights{
    // ...existing modes...
    "mymode": {Cost: 0.3, Latency: 0.2, Failure: 0.2, Weight: 0.3},
}

Add validation in internal/httpapi/handlers_chat.go and handlers_plan.go:

case "mymode":
    // valid

Add to routing config validation in handlers_routing.go.

Adding a New Orchestration Mode

Add the case in engine.Orchestrate():

case "mymode":
    // Implement multi-call pattern
    result, err := json.Marshal(map[string]any{...})
    return totalDecision, result, err

Add validation in handlers_plan.go.
Update Temporal if using workflows:

// In OrchestrationWorkflow
case "mymode":
    // Implement as Temporal activities

Adding New Admin Endpoints

Create handler in internal/httpapi/handlers_newfeature.go:

func NewFeatureHandler(d Dependencies) http.HandlerFunc {
    return func(w http.ResponseWriter, r *http.Request) {
        // Handler logic
    }
}

Mount route in internal/httpapi/routes.go:

r.Get("/admin/v1/newfeature", NewFeatureHandler(d))

Add to Dependencies if new services are needed.

Adding New Metrics

In internal/metrics/metrics.go:

type Registry struct {
    // ...existing metrics...
    NewMetric *prometheus.CounterVec
}

func New() *Registry {
    r := &Registry{
        NewMetric: prometheus.NewCounterVec(prometheus.CounterOpts{
            Namespace: "tokenhub",
            Name:      "new_metric_total",
            Help:      "Description of the new metric",
        }, []string{"label1", "label2"}),
    }
    // Register with Prometheus
    return r
}

Adding New Store Operations

Add to the interface in internal/store/store.go
Implement in SQLite in internal/store/sqlite.go
Add migration in Migrate() if new tables are needed
Write tests in internal/store/sqlite_test.go

Testing

TokenHub uses Go's standard testing package. Key test patterns:

Unit tests: Each package has *_test.go files
Integration tests: internal/httpapi/handlers_test.go tests the full HTTP stack
Mock adapters: mockSender in handler tests simulates provider responses
In-memory SQLite: Tests use :memory: DSN for isolated databases

Run all tests:

make test        # Standard tests
make test-race   # With race detector

Build

make build       # Build to bin/tokenhub
make package     # Build Docker image
make lint        # Run linter (requires golangci-lint)
make vet         # Go vet

TokenHub Documentation