Compare commits
23
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
aac502ea1f | ||
|
|
28753d11ed | ||
|
|
3068c68730 | ||
|
|
da361f2a21 | ||
|
|
1e17706367 | ||
|
|
49c9071a59 | ||
|
|
16ddd6d4ad | ||
|
|
fadc658ac0 | ||
|
|
bb7a490f5c | ||
|
|
4a3768ac35 | ||
|
|
f8882e86ae | ||
|
|
f9093fe713 | ||
|
|
ecaaaa8792 | ||
|
|
6d304a3e5a | ||
|
|
3ca654513f | ||
|
|
23fff1006a | ||
|
|
50bcd19634 | ||
|
|
a24e4f7174 | ||
|
|
0bde5dd8fe | ||
|
|
ef5e736673 | ||
|
|
aaa9a5d1cb | ||
|
|
ea7bd8336b | ||
|
|
ad5ba583ac |
+11
-16
@@ -21,9 +21,6 @@ jobs:
|
|||||||
run: |
|
run: |
|
||||||
apt-get update
|
apt-get update
|
||||||
apt-get install -y docker.io curl nodejs
|
apt-get install -y docker.io curl nodejs
|
||||||
curl -sLO "https://dl.k8s.io/release/$(curl -sL https://dl.k8s.io/release/stable.txt)/bin/linux/amd64/kubectl"
|
|
||||||
chmod +x kubectl && mv kubectl /usr/local/bin/
|
|
||||||
kubectl version --client
|
|
||||||
|
|
||||||
- name: Checkout code
|
- name: Checkout code
|
||||||
uses: actions/checkout@v4
|
uses: actions/checkout@v4
|
||||||
@@ -63,7 +60,7 @@ jobs:
|
|||||||
# (Temporal, Gotify). Tests run against its pod IP directly.
|
# (Temporal, Gotify). Tests run against its pod IP directly.
|
||||||
# Zero downtime: production pods untouched until tests pass.
|
# Zero downtime: production pods untouched until tests pass.
|
||||||
- name: Setup kubeconfig
|
- name: Setup kubeconfig
|
||||||
if: github.event_name == 'push' && github.ref == 'refs/heads/main'
|
if: false # SKIP: only needed for canary deploy, RBAC not ready
|
||||||
run: |
|
run: |
|
||||||
mkdir -p ~/.kube
|
mkdir -p ~/.kube
|
||||||
echo "${KUBECONFIG_B64}" | base64 -d > ~/.kube/config
|
echo "${KUBECONFIG_B64}" | base64 -d > ~/.kube/config
|
||||||
@@ -73,7 +70,7 @@ jobs:
|
|||||||
KUBECONFIG_B64: ${{ secrets.KUBECONFIG_B64 }}
|
KUBECONFIG_B64: ${{ secrets.KUBECONFIG_B64 }}
|
||||||
|
|
||||||
- name: Deploy canary pod
|
- name: Deploy canary pod
|
||||||
if: github.event_name == 'push' && github.ref == 'refs/heads/main'
|
if: false # SKIP: SA lacks pod create RBAC. Re-enable when fixed.
|
||||||
id: canary
|
id: canary
|
||||||
run: |
|
run: |
|
||||||
SHA="${{ steps.sha.outputs.short_sha }}"
|
SHA="${{ steps.sha.outputs.short_sha }}"
|
||||||
@@ -136,7 +133,7 @@ jobs:
|
|||||||
echo "✓ Canary ready at ${POD_IP}:8080"
|
echo "✓ Canary ready at ${POD_IP}:8080"
|
||||||
|
|
||||||
- name: Run integration tests against canary
|
- name: Run integration tests against canary
|
||||||
if: github.event_name == 'push' && github.ref == 'refs/heads/main'
|
if: false # SKIP: depends on canary pod
|
||||||
id: tests
|
id: tests
|
||||||
run: |
|
run: |
|
||||||
POD_IP="${{ steps.canary.outputs.pod_ip }}"
|
POD_IP="${{ steps.canary.outputs.pod_ip }}"
|
||||||
@@ -154,7 +151,7 @@ jobs:
|
|||||||
echo "result=${RESULT}" >> $GITHUB_OUTPUT
|
echo "result=${RESULT}" >> $GITHUB_OUTPUT
|
||||||
|
|
||||||
- name: Cleanup canary pod
|
- name: Cleanup canary pod
|
||||||
if: always() && github.event_name == 'push' && github.ref == 'refs/heads/main'
|
if: false # SKIP: depends on canary pod
|
||||||
run: |
|
run: |
|
||||||
POD_NAME="${{ steps.canary.outputs.pod_name }}"
|
POD_NAME="${{ steps.canary.outputs.pod_name }}"
|
||||||
if [ -n "${POD_NAME}" ]; then
|
if [ -n "${POD_NAME}" ]; then
|
||||||
@@ -165,7 +162,7 @@ jobs:
|
|||||||
fi
|
fi
|
||||||
|
|
||||||
- name: Gate on test result
|
- name: Gate on test result
|
||||||
if: github.event_name == 'push' && github.ref == 'refs/heads/main' && steps.tests.outputs.result != 'pass'
|
if: false # SKIP: depends on canary pod
|
||||||
run: |
|
run: |
|
||||||
echo "✗ Integration tests FAILED — production deployment unchanged"
|
echo "✗ Integration tests FAILED — production deployment unchanged"
|
||||||
exit 1
|
exit 1
|
||||||
@@ -173,14 +170,12 @@ jobs:
|
|||||||
# ── Zero-downtime promotion ──────────────────────────────
|
# ── Zero-downtime promotion ──────────────────────────────
|
||||||
# kubectl set image triggers a rolling update: new pods come up
|
# kubectl set image triggers a rolling update: new pods come up
|
||||||
# before old pods are terminated. Production traffic uninterrupted.
|
# before old pods are terminated. Production traffic uninterrupted.
|
||||||
- name: Promote — rolling update
|
# TODO: Re-enable rolling update when kubeconfig/RBAC is ready
|
||||||
if: github.event_name == 'push' && github.ref == 'refs/heads/main'
|
# - name: Promote — rolling update
|
||||||
run: |
|
# if: github.event_name == 'push' && github.ref == 'refs/heads/main'
|
||||||
SHA="${{ steps.sha.outputs.short_sha }}"
|
# run: |
|
||||||
kubectl set image deployment/api-gateway \
|
# kubectl set image deployment/api-gateway api-gateway="${IMAGE}:${SHA}" -n api
|
||||||
api-gateway="${IMAGE}:${SHA}" -n api
|
# kubectl rollout status deployment/api-gateway -n api --timeout=120s
|
||||||
kubectl rollout status deployment/api-gateway -n api --timeout=120s
|
|
||||||
echo "✓ Rolling update complete: ${IMAGE}:${SHA}"
|
|
||||||
|
|
||||||
- name: Tag :latest (informational)
|
- name: Tag :latest (informational)
|
||||||
if: github.event_name == 'push' && github.ref == 'refs/heads/main'
|
if: github.event_name == 'push' && github.ref == 'refs/heads/main'
|
||||||
|
|||||||
@@ -12,6 +12,7 @@ import (
|
|||||||
"forgejo.riotpiao.com/rock/homelab-frontend/internal/auth"
|
"forgejo.riotpiao.com/rock/homelab-frontend/internal/auth"
|
||||||
"forgejo.riotpiao.com/rock/homelab-frontend/internal/config"
|
"forgejo.riotpiao.com/rock/homelab-frontend/internal/config"
|
||||||
"forgejo.riotpiao.com/rock/homelab-frontend/internal/notification"
|
"forgejo.riotpiao.com/rock/homelab-frontend/internal/notification"
|
||||||
|
"forgejo.riotpiao.com/rock/homelab-frontend/internal/observability"
|
||||||
"forgejo.riotpiao.com/rock/homelab-frontend/internal/proxy"
|
"forgejo.riotpiao.com/rock/homelab-frontend/internal/proxy"
|
||||||
"forgejo.riotpiao.com/rock/homelab-frontend/internal/server"
|
"forgejo.riotpiao.com/rock/homelab-frontend/internal/server"
|
||||||
"forgejo.riotpiao.com/rock/homelab-frontend/internal/serviceadapter"
|
"forgejo.riotpiao.com/rock/homelab-frontend/internal/serviceadapter"
|
||||||
@@ -112,6 +113,14 @@ func main() {
|
|||||||
poimenAdapter := serviceadapter.NewPoimenAdapter(temporalHandler)
|
poimenAdapter := serviceadapter.NewPoimenAdapter(temporalHandler)
|
||||||
_ = registry.Add(poimenAdapter)
|
_ = registry.Add(poimenAdapter)
|
||||||
|
|
||||||
|
// Add memory service adapter (code-registered; config memory entry is skipped).
|
||||||
|
// Includes standing-context resource for CAG support (issue #44).
|
||||||
|
_ = registry.Add(serviceadapter.NewMemoryAdapter())
|
||||||
|
|
||||||
|
// Add sessions service adapter — REST CRUD bridged to poimen-memory ingest
|
||||||
|
// and poimen-graphql reads. Issue #50.
|
||||||
|
_ = registry.Add(serviceadapter.NewSessionAdapter())
|
||||||
|
|
||||||
// Add other adapters from config (skip if already registered in code)
|
// Add other adapters from config (skip if already registered in code)
|
||||||
for _, a := range cfg.Adapters {
|
for _, a := range cfg.Adapters {
|
||||||
if existing := registry.Get(a.ServiceName); existing != nil {
|
if existing := registry.Get(a.ServiceName); existing != nil {
|
||||||
@@ -135,6 +144,7 @@ func main() {
|
|||||||
// Create router that handles health endpoints, X-Service (ServiceAdapter) routing,
|
// Create router that handles health endpoints, X-Service (ServiceAdapter) routing,
|
||||||
// temporal endpoints, and passes others to upstream handler
|
// temporal endpoints, and passes others to upstream handler
|
||||||
router := server.NewRouter(healthChecker, dispatcher, temporalHandler, upstreamHandler)
|
router := server.NewRouter(healthChecker, dispatcher, temporalHandler, upstreamHandler)
|
||||||
|
router.SetMetricsHandler(observability.NewMetricsHandler(observability.Global()))
|
||||||
|
|
||||||
// Wrap router with tracing middleware
|
// Wrap router with tracing middleware
|
||||||
tracedRouter := tracing.Middleware(router)
|
tracedRouter := tracing.Middleware(router)
|
||||||
|
|||||||
@@ -33,7 +33,11 @@ type ModelUpstream struct {
|
|||||||
// Name is the model name clients send (e.g., "reasoning", "ornith:35b").
|
// Name is the model name clients send (e.g., "reasoning", "ornith:35b").
|
||||||
Name string
|
Name string
|
||||||
// Address is the upstream server address (host:port).
|
// Address is the upstream server address (host:port).
|
||||||
|
// When Addresses is set, this holds Addresses[0] for backward compatibility.
|
||||||
Address string
|
Address string
|
||||||
|
// Addresses holds all backend addresses for multi-backend cache-aware routing.
|
||||||
|
// Populated from the 'addresses' YAML field; backward-compat with single 'address'.
|
||||||
|
Addresses []string
|
||||||
// Path is the upstream path for this model (e.g., "/v1/chat/completions").
|
// Path is the upstream path for this model (e.g., "/v1/chat/completions").
|
||||||
Path string
|
Path string
|
||||||
// UpstreamModel is the model name to send to the upstream server.
|
// UpstreamModel is the model name to send to the upstream server.
|
||||||
@@ -116,6 +120,18 @@ func (c *Config) ModelNames() []string {
|
|||||||
return names
|
return names
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// AllAddresses returns all configured backend addresses for the model.
|
||||||
|
// Returns Addresses if non-empty; otherwise wraps Address in a slice.
|
||||||
|
func (m *ModelUpstream) AllAddresses() []string {
|
||||||
|
if len(m.Addresses) > 0 {
|
||||||
|
return m.Addresses
|
||||||
|
}
|
||||||
|
if m.Address != "" {
|
||||||
|
return []string{m.Address}
|
||||||
|
}
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
// Load reads configuration from environment variables with defaults.
|
// Load reads configuration from environment variables with defaults.
|
||||||
func Load() (*Config, error) {
|
func Load() (*Config, error) {
|
||||||
listenAddr := "127.0.0.1:8080"
|
listenAddr := "127.0.0.1:8080"
|
||||||
|
|||||||
@@ -37,9 +37,10 @@ type rawRoute struct {
|
|||||||
|
|
||||||
// rawModel represents a single model entry in the YAML configuration.
|
// rawModel represents a single model entry in the YAML configuration.
|
||||||
type rawModel struct {
|
type rawModel struct {
|
||||||
Name string `yaml:"name"`
|
Name string `yaml:"name"`
|
||||||
Address string `yaml:"address"`
|
Address string `yaml:"address"`
|
||||||
Path string `yaml:"path"`
|
Addresses []string `yaml:"addresses"`
|
||||||
|
Path string `yaml:"path"`
|
||||||
UpstreamModel string `yaml:"upstreamModel"`
|
UpstreamModel string `yaml:"upstreamModel"`
|
||||||
AuthRequired *bool `yaml:"authRequired"`
|
AuthRequired *bool `yaml:"authRequired"`
|
||||||
}
|
}
|
||||||
@@ -121,11 +122,18 @@ func LoadRoutesAndModelsFromFile(path string) (map[string]*Route, map[string]*Mo
|
|||||||
if _, exists := models[rawModel.Name]; exists {
|
if _, exists := models[rawModel.Name]; exists {
|
||||||
return nil, nil, nil, AuthConfig{}, fmt.Errorf("duplicate model: \"%s\"", rawModel.Name)
|
return nil, nil, nil, AuthConfig{}, fmt.Errorf("duplicate model: \"%s\"", rawModel.Name)
|
||||||
}
|
}
|
||||||
if rawModel.Address == "" {
|
// Support both 'address' (single) and 'addresses' (multi-backend) YAML fields.
|
||||||
return nil, nil, nil, AuthConfig{}, fmt.Errorf("model \"%s\": field 'address' is required", rawModel.Name)
|
addresses := rawModel.Addresses
|
||||||
|
if len(addresses) == 0 && rawModel.Address != "" {
|
||||||
|
addresses = []string{rawModel.Address}
|
||||||
}
|
}
|
||||||
if _, _, err := net.SplitHostPort(rawModel.Address); err != nil {
|
if len(addresses) == 0 {
|
||||||
return nil, nil, nil, AuthConfig{}, fmt.Errorf("model \"%s\": invalid address \"%s\": %w", rawModel.Name, rawModel.Address, err)
|
return nil, nil, nil, AuthConfig{}, fmt.Errorf("model \"%s\": 'address' or 'addresses' is required", rawModel.Name)
|
||||||
|
}
|
||||||
|
for _, addr := range addresses {
|
||||||
|
if _, _, err := net.SplitHostPort(addr); err != nil {
|
||||||
|
return nil, nil, nil, AuthConfig{}, fmt.Errorf("model \"%s\": invalid address \"%s\": %w", rawModel.Name, addr, err)
|
||||||
|
}
|
||||||
}
|
}
|
||||||
// Default authRequired to global auth.enabled if not specified per-model
|
// Default authRequired to global auth.enabled if not specified per-model
|
||||||
authRequired := false
|
authRequired := false
|
||||||
@@ -134,7 +142,8 @@ func LoadRoutesAndModelsFromFile(path string) (map[string]*Route, map[string]*Mo
|
|||||||
}
|
}
|
||||||
models[rawModel.Name] = &ModelUpstream{
|
models[rawModel.Name] = &ModelUpstream{
|
||||||
Name: rawModel.Name,
|
Name: rawModel.Name,
|
||||||
Address: rawModel.Address,
|
Address: addresses[0],
|
||||||
|
Addresses: addresses,
|
||||||
Path: rawModel.Path,
|
Path: rawModel.Path,
|
||||||
UpstreamModel: rawModel.UpstreamModel,
|
UpstreamModel: rawModel.UpstreamModel,
|
||||||
AuthRequired: authRequired,
|
AuthRequired: authRequired,
|
||||||
|
|||||||
@@ -161,7 +161,7 @@ models:
|
|||||||
if err == nil {
|
if err == nil {
|
||||||
t.Errorf("expected error for missing address, got nil")
|
t.Errorf("expected error for missing address, got nil")
|
||||||
}
|
}
|
||||||
if err.Error() != "model \"reasoning\": field 'address' is required" {
|
if err.Error() != "model \"reasoning\": 'address' or 'addresses' is required" {
|
||||||
t.Errorf("expected missing address error, got: %v", err)
|
t.Errorf("expected missing address error, got: %v", err)
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
@@ -283,3 +283,139 @@ models:
|
|||||||
t.Errorf("expected different ModelUpstream objects even for same address")
|
t.Errorf("expected different ModelUpstream objects even for same address")
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// minimalRouteYAML is a valid routes block required by LoadRoutesAndModelsFromFile.
|
||||||
|
const minimalRouteYAML = `
|
||||||
|
routes:
|
||||||
|
- name: test-route
|
||||||
|
upstream:
|
||||||
|
address: "localhost:8000"
|
||||||
|
pathRewrite: ""
|
||||||
|
connectTimeout: "10s"
|
||||||
|
readTimeout: "30s"
|
||||||
|
writeTimeout: "30s"
|
||||||
|
maxBodySize: 1048576
|
||||||
|
authRequired: false
|
||||||
|
`
|
||||||
|
|
||||||
|
// TestLoadModelsAddressesField verifies address / addresses YAML field handling.
|
||||||
|
func TestLoadModelsAddressesField(t *testing.T) {
|
||||||
|
tests := []struct {
|
||||||
|
name string
|
||||||
|
modelsYAML string
|
||||||
|
wantErr bool
|
||||||
|
errContains string
|
||||||
|
wantAddresses []string
|
||||||
|
wantAddress string
|
||||||
|
}{
|
||||||
|
{
|
||||||
|
name: "single address field",
|
||||||
|
modelsYAML: `
|
||||||
|
models:
|
||||||
|
- name: "mymodel"
|
||||||
|
address: "host:8080"
|
||||||
|
path: "/v1/chat"
|
||||||
|
`,
|
||||||
|
wantErr: false,
|
||||||
|
wantAddresses: []string{"host:8080"},
|
||||||
|
wantAddress: "host:8080",
|
||||||
|
},
|
||||||
|
{
|
||||||
|
name: "multi addresses field",
|
||||||
|
modelsYAML: `
|
||||||
|
models:
|
||||||
|
- name: "mymodel"
|
||||||
|
addresses:
|
||||||
|
- "a:80"
|
||||||
|
- "b:80"
|
||||||
|
path: "/v1/chat"
|
||||||
|
`,
|
||||||
|
wantErr: false,
|
||||||
|
wantAddresses: []string{"a:80", "b:80"},
|
||||||
|
wantAddress: "a:80",
|
||||||
|
},
|
||||||
|
{
|
||||||
|
name: "neither field",
|
||||||
|
modelsYAML: `
|
||||||
|
models:
|
||||||
|
- name: "mymodel"
|
||||||
|
path: "/v1/chat"
|
||||||
|
`,
|
||||||
|
wantErr: true,
|
||||||
|
errContains: "mymodel",
|
||||||
|
},
|
||||||
|
{
|
||||||
|
name: "invalid address format",
|
||||||
|
modelsYAML: `
|
||||||
|
models:
|
||||||
|
- name: "mymodel"
|
||||||
|
address: "notavalidhostport"
|
||||||
|
path: "/v1/chat"
|
||||||
|
`,
|
||||||
|
wantErr: true,
|
||||||
|
errContains: "mymodel",
|
||||||
|
},
|
||||||
|
}
|
||||||
|
|
||||||
|
for _, tc := range tests {
|
||||||
|
t.Run(tc.name, func(t *testing.T) {
|
||||||
|
data := minimalRouteYAML + tc.modelsYAML
|
||||||
|
|
||||||
|
tmpFile, err := os.CreateTemp("", "config-*.yaml")
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("CreateTemp: %v", err)
|
||||||
|
}
|
||||||
|
defer os.Remove(tmpFile.Name())
|
||||||
|
tmpFile.WriteString(data)
|
||||||
|
tmpFile.Close()
|
||||||
|
|
||||||
|
_, models, _, _, loadErr := LoadRoutesAndModelsFromFile(tmpFile.Name())
|
||||||
|
|
||||||
|
if tc.wantErr {
|
||||||
|
if loadErr == nil {
|
||||||
|
t.Fatalf("expected error, got nil")
|
||||||
|
}
|
||||||
|
if tc.errContains != "" && !containsStr(loadErr.Error(), tc.errContains) {
|
||||||
|
t.Errorf("error %q does not contain %q", loadErr.Error(), tc.errContains)
|
||||||
|
}
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
if loadErr != nil {
|
||||||
|
t.Fatalf("unexpected error: %v", loadErr)
|
||||||
|
}
|
||||||
|
|
||||||
|
m, ok := models["mymodel"]
|
||||||
|
if !ok {
|
||||||
|
t.Fatalf("model 'mymodel' not found in result")
|
||||||
|
}
|
||||||
|
|
||||||
|
if m.Address != tc.wantAddress {
|
||||||
|
t.Errorf("Address: want %q, got %q", tc.wantAddress, m.Address)
|
||||||
|
}
|
||||||
|
|
||||||
|
if len(m.Addresses) != len(tc.wantAddresses) {
|
||||||
|
t.Errorf("Addresses len: want %d, got %d (%v)", len(tc.wantAddresses), len(m.Addresses), m.Addresses)
|
||||||
|
} else {
|
||||||
|
for i, a := range tc.wantAddresses {
|
||||||
|
if m.Addresses[i] != a {
|
||||||
|
t.Errorf("Addresses[%d]: want %q, got %q", i, a, m.Addresses[i])
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
})
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// containsStr is a helper to avoid importing strings just for Contains.
|
||||||
|
func containsStr(s, sub string) bool {
|
||||||
|
return len(s) >= len(sub) && (s == sub || len(sub) == 0 ||
|
||||||
|
func() bool {
|
||||||
|
for i := 0; i <= len(s)-len(sub); i++ {
|
||||||
|
if s[i:i+len(sub)] == sub {
|
||||||
|
return true
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return false
|
||||||
|
}())
|
||||||
|
}
|
||||||
|
|||||||
@@ -198,3 +198,11 @@ func FromContext(ctx context.Context) *Logger {
|
|||||||
func WithContext(ctx context.Context, l *Logger) context.Context {
|
func WithContext(ctx context.Context, l *Logger) context.Context {
|
||||||
return context.WithValue(ctx, "logger", l)
|
return context.WithValue(ctx, "logger", l)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// Event emits arbitrary fields as a structured JSON line.
|
||||||
|
// Use when field values include non-string types (numerics, floats).
|
||||||
|
func Event(fields map[string]interface{}) {
|
||||||
|
ensure()
|
||||||
|
data, _ := json.Marshal(fields)
|
||||||
|
_, _ = global.out.Write(append(data, '\n'))
|
||||||
|
}
|
||||||
|
|||||||
@@ -36,6 +36,11 @@ type Metrics struct {
|
|||||||
itlMs map[string][]int64 // samples for histogram
|
itlMs map[string][]int64 // samples for histogram
|
||||||
// Token counts
|
// Token counts
|
||||||
tokenCount map[string]int64
|
tokenCount map[string]int64
|
||||||
|
|
||||||
|
// CAG metrics
|
||||||
|
warmupRequestsDelayed int64
|
||||||
|
prefixCacheHitTokens int64
|
||||||
|
queueFullRequests int64
|
||||||
}
|
}
|
||||||
|
|
||||||
// NewMetrics creates a new Metrics instance.
|
// NewMetrics creates a new Metrics instance.
|
||||||
@@ -280,4 +285,61 @@ func (m *Metrics) Reset() {
|
|||||||
m.ttftMs = make(map[string][]int64)
|
m.ttftMs = make(map[string][]int64)
|
||||||
m.itlMs = make(map[string][]int64)
|
m.itlMs = make(map[string][]int64)
|
||||||
m.tokenCount = make(map[string]int64)
|
m.tokenCount = make(map[string]int64)
|
||||||
|
m.warmupRequestsDelayed = 0
|
||||||
|
m.prefixCacheHitTokens = 0
|
||||||
|
m.queueFullRequests = 0
|
||||||
|
}
|
||||||
|
|
||||||
|
// IncrQueueFull increments queue_full_requests_total.
|
||||||
|
func (m *Metrics) IncrQueueFull() {
|
||||||
|
m.mu.Lock()
|
||||||
|
defer m.mu.Unlock()
|
||||||
|
m.queueFullRequests++
|
||||||
|
}
|
||||||
|
|
||||||
|
// GetQueueFull returns queue_full_requests_total.
|
||||||
|
func (m *Metrics) GetQueueFull() int64 {
|
||||||
|
m.mu.RLock()
|
||||||
|
defer m.mu.RUnlock()
|
||||||
|
return m.queueFullRequests
|
||||||
|
}
|
||||||
|
|
||||||
|
// IncrWarmupDelayed increments warmup_requests_delayed_total.
|
||||||
|
func (m *Metrics) IncrWarmupDelayed() {
|
||||||
|
m.mu.Lock()
|
||||||
|
defer m.mu.Unlock()
|
||||||
|
m.warmupRequestsDelayed++
|
||||||
|
}
|
||||||
|
|
||||||
|
// GetWarmupDelayed returns warmup_requests_delayed_total.
|
||||||
|
func (m *Metrics) GetWarmupDelayed() int64 {
|
||||||
|
m.mu.RLock()
|
||||||
|
defer m.mu.RUnlock()
|
||||||
|
return m.warmupRequestsDelayed
|
||||||
|
}
|
||||||
|
|
||||||
|
// IncrPrefixCacheHitTokens increments vllm_prefix_cache_hit_tokens_total by n.
|
||||||
|
func (m *Metrics) IncrPrefixCacheHitTokens(n int64) {
|
||||||
|
m.mu.Lock()
|
||||||
|
defer m.mu.Unlock()
|
||||||
|
m.prefixCacheHitTokens += n
|
||||||
|
}
|
||||||
|
|
||||||
|
// GetPrefixCacheHitTokens returns vllm_prefix_cache_hit_tokens_total.
|
||||||
|
func (m *Metrics) GetPrefixCacheHitTokens() int64 {
|
||||||
|
m.mu.RLock()
|
||||||
|
defer m.mu.RUnlock()
|
||||||
|
return m.prefixCacheHitTokens
|
||||||
|
}
|
||||||
|
|
||||||
|
// globalMetrics is the singleton Metrics instance for gateway-wide telemetry.
|
||||||
|
var (
|
||||||
|
globalMetrics *Metrics
|
||||||
|
globalMetricsOnce sync.Once
|
||||||
|
)
|
||||||
|
|
||||||
|
// Global returns the singleton Metrics instance, initialising on first call.
|
||||||
|
func Global() *Metrics {
|
||||||
|
globalMetricsOnce.Do(func() { globalMetrics = NewMetrics() })
|
||||||
|
return globalMetrics
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -36,6 +36,18 @@ func (p *PrometheusExporter) Export() string {
|
|||||||
lines = append(lines, "# TYPE request_duration_seconds histogram")
|
lines = append(lines, "# TYPE request_duration_seconds histogram")
|
||||||
p.exportRequestDuration(&lines)
|
p.exportRequestDuration(&lines)
|
||||||
|
|
||||||
|
lines = append(lines, "# HELP queue_full_requests_total Requests rejected due to backend queue saturation")
|
||||||
|
lines = append(lines, "# TYPE queue_full_requests_total counter")
|
||||||
|
lines = append(lines, fmt.Sprintf("queue_full_requests_total %d", p.metrics.GetQueueFull()))
|
||||||
|
|
||||||
|
lines = append(lines, "# HELP warmup_requests_delayed_total LLM warmup requests delayed for vLLM slot management")
|
||||||
|
lines = append(lines, "# TYPE warmup_requests_delayed_total counter")
|
||||||
|
lines = append(lines, fmt.Sprintf("warmup_requests_delayed_total %d", p.metrics.GetWarmupDelayed()))
|
||||||
|
|
||||||
|
lines = append(lines, "# HELP vllm_prefix_cache_hit_tokens_total vLLM KV prefix cache hit tokens")
|
||||||
|
lines = append(lines, "# TYPE vllm_prefix_cache_hit_tokens_total counter")
|
||||||
|
lines = append(lines, fmt.Sprintf("vllm_prefix_cache_hit_tokens_total %d", p.metrics.GetPrefixCacheHitTokens()))
|
||||||
|
|
||||||
return strings.Join(lines, "\n") + "\n"
|
return strings.Join(lines, "\n") + "\n"
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
@@ -3,7 +3,6 @@ package proxy
|
|||||||
import (
|
import (
|
||||||
"bufio"
|
"bufio"
|
||||||
"fmt"
|
"fmt"
|
||||||
"io"
|
|
||||||
"net"
|
"net"
|
||||||
"net/http"
|
"net/http"
|
||||||
"strings"
|
"strings"
|
||||||
@@ -12,87 +11,37 @@ import (
|
|||||||
"forgejo.riotpiao.com/rock/homelab-frontend/internal/observability"
|
"forgejo.riotpiao.com/rock/homelab-frontend/internal/observability"
|
||||||
)
|
)
|
||||||
|
|
||||||
// LLMMetricsCapture wraps a response writer to capture TTFT and ITL metrics
|
// ResponseWriterWrapper wraps an http.ResponseWriter to capture per-request LLM
|
||||||
type LLMMetricsCapture struct {
|
// latency metrics (TTFT, ITL, generated tokens) while a streaming response is
|
||||||
writer io.WriteCloser
|
// copied from the upstream by httputil.ReverseProxy.
|
||||||
model string
|
//
|
||||||
metrics *observability.Metrics
|
// Only server-sent-event responses are measured. A non-streaming completion is
|
||||||
firstTokenTime time.Time
|
// written in a single Write, so recording it would emit one "TTFT" sample equal
|
||||||
lastTokenTime time.Time
|
// to the whole request latency and no ITL samples at all — numbers that look
|
||||||
requestStartTime time.Time
|
// plausible but mean nothing.
|
||||||
ttftRecorded bool
|
|
||||||
tokenCount int64
|
|
||||||
responseStartTime time.Time
|
|
||||||
}
|
|
||||||
|
|
||||||
// NewLLMMetricsCapture creates a new metrics capture wrapper
|
|
||||||
func NewLLMMetricsCapture(writer io.WriteCloser, model string, metrics *observability.Metrics, startTime time.Time) *LLMMetricsCapture {
|
|
||||||
return &LLMMetricsCapture{
|
|
||||||
writer: writer,
|
|
||||||
model: model,
|
|
||||||
metrics: metrics,
|
|
||||||
requestStartTime: startTime,
|
|
||||||
responseStartTime: time.Now(),
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
// Write intercepts writes to detect tokens and record metrics
|
|
||||||
func (c *LLMMetricsCapture) Write(p []byte) (int, error) {
|
|
||||||
// Record first token time
|
|
||||||
if !c.ttftRecorded && len(p) > 0 {
|
|
||||||
now := time.Now()
|
|
||||||
ttft := now.Sub(c.requestStartTime).Milliseconds()
|
|
||||||
c.metrics.RecordTTFT(c.model, ttft)
|
|
||||||
c.ttftRecorded = true
|
|
||||||
c.firstTokenTime = now
|
|
||||||
c.lastTokenTime = now
|
|
||||||
}
|
|
||||||
|
|
||||||
// Count tokens in SSE stream (simple: count "data: " lines)
|
|
||||||
if c.ttftRecorded {
|
|
||||||
tokenCount := strings.Count(string(p), "data: ")
|
|
||||||
if tokenCount > 0 {
|
|
||||||
now := time.Now()
|
|
||||||
if !c.firstTokenTime.IsZero() && c.lastTokenTime != now {
|
|
||||||
itl := now.Sub(c.lastTokenTime).Milliseconds()
|
|
||||||
c.metrics.RecordITL(c.model, itl)
|
|
||||||
}
|
|
||||||
c.lastTokenTime = now
|
|
||||||
c.tokenCount += int64(tokenCount)
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
return c.writer.Write(p)
|
|
||||||
}
|
|
||||||
|
|
||||||
// Close records final metrics and closes writer
|
|
||||||
func (c *LLMMetricsCapture) Close() error {
|
|
||||||
if c.tokenCount > 0 {
|
|
||||||
c.metrics.RecordTokenCount(c.model, c.tokenCount)
|
|
||||||
}
|
|
||||||
return c.writer.Close()
|
|
||||||
}
|
|
||||||
|
|
||||||
// ResponseWriterWrapper wraps http.ResponseWriter to capture metrics
|
|
||||||
type ResponseWriterWrapper struct {
|
type ResponseWriterWrapper struct {
|
||||||
writer http.ResponseWriter
|
writer http.ResponseWriter
|
||||||
statusCode int
|
statusCode int
|
||||||
metrics *observability.Metrics
|
metrics *observability.Metrics
|
||||||
model string
|
model string
|
||||||
startTime time.Time
|
startTime time.Time
|
||||||
firstByteTime time.Time
|
firstByteTime time.Time
|
||||||
lastWriteTime time.Time
|
lastWriteTime time.Time
|
||||||
ttftRecorded bool
|
ttftRecorded bool
|
||||||
|
streamChecked bool
|
||||||
|
streaming bool
|
||||||
}
|
}
|
||||||
|
|
||||||
// NewResponseWriterWrapper creates a wrapper for response writer
|
// NewResponseWriterWrapper creates a wrapper for response writer.
|
||||||
|
// startTime must be the time the request entered the gateway, not the time this
|
||||||
|
// wrapper was constructed, otherwise TTFT measures the wrong interval.
|
||||||
func NewResponseWriterWrapper(w http.ResponseWriter, model string, metrics *observability.Metrics, startTime time.Time) *ResponseWriterWrapper {
|
func NewResponseWriterWrapper(w http.ResponseWriter, model string, metrics *observability.Metrics, startTime time.Time) *ResponseWriterWrapper {
|
||||||
return &ResponseWriterWrapper{
|
return &ResponseWriterWrapper{
|
||||||
writer: w,
|
writer: w,
|
||||||
model: model,
|
model: model,
|
||||||
metrics: metrics,
|
metrics: metrics,
|
||||||
startTime: startTime,
|
startTime: startTime,
|
||||||
statusCode: 200,
|
statusCode: 200,
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -101,29 +50,48 @@ func (w *ResponseWriterWrapper) Header() http.Header {
|
|||||||
return w.writer.Header()
|
return w.writer.Header()
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// isStreaming reports whether this is an SSE response. Resolved once from the
|
||||||
|
// Content-Type header, which ReverseProxy has already set before the first write.
|
||||||
|
func (w *ResponseWriterWrapper) isStreaming() bool {
|
||||||
|
if !w.streamChecked {
|
||||||
|
w.streamChecked = true
|
||||||
|
w.streaming = strings.Contains(w.writer.Header().Get("Content-Type"), "text/event-stream")
|
||||||
|
}
|
||||||
|
return w.streaming
|
||||||
|
}
|
||||||
|
|
||||||
|
// countSSETokens counts generated tokens in a chunk of an SSE stream. These
|
||||||
|
// upstreams emit one `data:` event per token, so the event count is the token
|
||||||
|
// count; the `[DONE]` terminator is not a token.
|
||||||
|
func countSSETokens(b []byte) int64 {
|
||||||
|
var n int64
|
||||||
|
for _, line := range strings.Split(string(b), "\n") {
|
||||||
|
line = strings.TrimSpace(line)
|
||||||
|
if strings.HasPrefix(line, "data:") && !strings.HasPrefix(line, "data: [DONE]") {
|
||||||
|
n++
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return n
|
||||||
|
}
|
||||||
|
|
||||||
// Write implements http.ResponseWriter
|
// Write implements http.ResponseWriter
|
||||||
func (w *ResponseWriterWrapper) Write(b []byte) (int, error) {
|
func (w *ResponseWriterWrapper) Write(b []byte) (int, error) {
|
||||||
// Record TTFT on first write
|
if len(b) > 0 && w.isStreaming() {
|
||||||
if !w.ttftRecorded && len(b) > 0 {
|
|
||||||
now := time.Now()
|
now := time.Now()
|
||||||
ttft := now.Sub(w.startTime).Milliseconds()
|
if !w.ttftRecorded {
|
||||||
w.metrics.RecordTTFT(w.model, ttft)
|
w.metrics.RecordTTFT(w.model, now.Sub(w.startTime).Milliseconds())
|
||||||
w.ttftRecorded = true
|
w.ttftRecorded = true
|
||||||
w.firstByteTime = now
|
w.firstByteTime = now
|
||||||
w.lastWriteTime = now
|
} else if itl := now.Sub(w.lastWriteTime).Milliseconds(); itl > 0 {
|
||||||
}
|
// Only record positive ITL: a zero means two chunks landed in the
|
||||||
|
// same millisecond, which carries no timing information.
|
||||||
// Record ITL for subsequent writes (for streaming)
|
w.metrics.RecordITL(w.model, itl)
|
||||||
if w.ttftRecorded && len(b) > 0 {
|
|
||||||
now := time.Now()
|
|
||||||
if !w.firstByteTime.IsZero() && w.lastWriteTime != now {
|
|
||||||
itl := now.Sub(w.lastWriteTime).Milliseconds()
|
|
||||||
// Only record if ITL > 0 (avoid recording same millisecond twice)
|
|
||||||
if itl > 0 {
|
|
||||||
w.metrics.RecordITL(w.model, itl)
|
|
||||||
}
|
|
||||||
}
|
}
|
||||||
w.lastWriteTime = now
|
w.lastWriteTime = now
|
||||||
|
|
||||||
|
if n := countSSETokens(b); n > 0 {
|
||||||
|
w.metrics.RecordTokenCount(w.model, n)
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
return w.writer.Write(b)
|
return w.writer.Write(b)
|
||||||
@@ -150,4 +118,8 @@ func (w *ResponseWriterWrapper) Hijack() (net.Conn, *bufio.ReadWriter, error) {
|
|||||||
return nil, nil, fmt.Errorf("response writer does not implement Hijacker")
|
return nil, nil, fmt.Errorf("response writer does not implement Hijacker")
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// Unwrap exposes the underlying writer to http.ResponseController so optional
|
||||||
|
// capabilities (Flush, Hijack, deadlines) keep working through the wrapper.
|
||||||
|
func (w *ResponseWriterWrapper) Unwrap() http.ResponseWriter {
|
||||||
|
return w.writer
|
||||||
|
}
|
||||||
|
|||||||
@@ -0,0 +1,171 @@
|
|||||||
|
package proxy
|
||||||
|
|
||||||
|
import (
|
||||||
|
"fmt"
|
||||||
|
"io"
|
||||||
|
"net/http"
|
||||||
|
"net/http/httptest"
|
||||||
|
"regexp"
|
||||||
|
"strings"
|
||||||
|
"testing"
|
||||||
|
"time"
|
||||||
|
|
||||||
|
"forgejo.riotpiao.com/rock/homelab-frontend/internal/config"
|
||||||
|
"forgejo.riotpiao.com/rock/homelab-frontend/internal/observability"
|
||||||
|
)
|
||||||
|
|
||||||
|
// exportMetrics renders the Prometheus text exposition for m.
|
||||||
|
func exportMetrics(m *observability.Metrics) string {
|
||||||
|
return observability.NewPrometheusExporter(m).Export()
|
||||||
|
}
|
||||||
|
|
||||||
|
// requireSeries asserts the exporter emitted a sample line with the exact value.
|
||||||
|
func requireSeries(t *testing.T, out, want string) {
|
||||||
|
t.Helper()
|
||||||
|
for _, line := range strings.Split(out, "\n") {
|
||||||
|
if strings.TrimSpace(line) == want {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
}
|
||||||
|
t.Errorf("missing metrics series %q\n--- got ---\n%s", want, out)
|
||||||
|
}
|
||||||
|
|
||||||
|
// TestResponseWriterWrapperRecordsSSEMetrics guards the regression where the
|
||||||
|
// LLM latency wrapper existed but was never constructed, leaving
|
||||||
|
// llm_ttft_seconds / llm_itl_seconds / llm_tokens_total as HELP-only stubs.
|
||||||
|
func TestResponseWriterWrapperRecordsSSEMetrics(t *testing.T) {
|
||||||
|
m := observability.NewMetrics()
|
||||||
|
|
||||||
|
rec := httptest.NewRecorder()
|
||||||
|
rec.Header().Set("Content-Type", "text/event-stream")
|
||||||
|
|
||||||
|
w := NewResponseWriterWrapper(rec, "reasoning", m, time.Now().Add(-50*time.Millisecond))
|
||||||
|
|
||||||
|
if _, err := w.Write([]byte("data: {\"token\":\"a\"}\n\n")); err != nil {
|
||||||
|
t.Fatalf("write: %v", err)
|
||||||
|
}
|
||||||
|
time.Sleep(5 * time.Millisecond)
|
||||||
|
if _, err := w.Write([]byte("data: {\"token\":\"b\"}\n\n")); err != nil {
|
||||||
|
t.Fatalf("write: %v", err)
|
||||||
|
}
|
||||||
|
// The terminator is not a generated token.
|
||||||
|
if _, err := w.Write([]byte("data: [DONE]\n\n")); err != nil {
|
||||||
|
t.Fatalf("write: %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
out := exportMetrics(m)
|
||||||
|
requireSeries(t, out, `llm_ttft_seconds_count{model="reasoning"} 1`)
|
||||||
|
requireSeries(t, out, `llm_itl_seconds_count{model="reasoning"} 1`)
|
||||||
|
requireSeries(t, out, `llm_tokens_total{model="reasoning"} 2`)
|
||||||
|
}
|
||||||
|
|
||||||
|
// TestResponseWriterWrapperIgnoresNonStreaming verifies a single-shot JSON
|
||||||
|
// response does not produce a bogus TTFT sample equal to full request latency.
|
||||||
|
func TestResponseWriterWrapperIgnoresNonStreaming(t *testing.T) {
|
||||||
|
m := observability.NewMetrics()
|
||||||
|
|
||||||
|
rec := httptest.NewRecorder()
|
||||||
|
rec.Header().Set("Content-Type", "application/json")
|
||||||
|
|
||||||
|
w := NewResponseWriterWrapper(rec, "reasoning", m, time.Now().Add(-50*time.Millisecond))
|
||||||
|
if _, err := w.Write([]byte(`{"choices":[]}`)); err != nil {
|
||||||
|
t.Fatalf("write: %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
out := exportMetrics(m)
|
||||||
|
if strings.Contains(out, `llm_ttft_seconds_count{model="reasoning"}`) {
|
||||||
|
t.Errorf("non-streaming response recorded TTFT\n--- got ---\n%s", out)
|
||||||
|
}
|
||||||
|
if strings.Contains(out, `llm_tokens_total{model="reasoning"}`) {
|
||||||
|
t.Errorf("non-streaming response recorded tokens\n--- got ---\n%s", out)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// TestResponseWriterWrapperPreservesStreaming verifies the wrapper stays
|
||||||
|
// transparent: bytes reach the client and Flush still works, so wrapping the
|
||||||
|
// writer cannot silently break SSE (issue #33).
|
||||||
|
func TestResponseWriterWrapperPreservesStreaming(t *testing.T) {
|
||||||
|
m := observability.NewMetrics()
|
||||||
|
|
||||||
|
rec := httptest.NewRecorder()
|
||||||
|
rec.Header().Set("Content-Type", "text/event-stream")
|
||||||
|
|
||||||
|
w := NewResponseWriterWrapper(rec, "reasoning", m, time.Now())
|
||||||
|
payload := "data: {\"token\":\"a\"}\n\n"
|
||||||
|
if _, err := w.Write([]byte(payload)); err != nil {
|
||||||
|
t.Fatalf("write: %v", err)
|
||||||
|
}
|
||||||
|
w.Flush()
|
||||||
|
|
||||||
|
if got := rec.Body.String(); got != payload {
|
||||||
|
t.Errorf("body = %q, want %q", got, payload)
|
||||||
|
}
|
||||||
|
if !rec.Flushed {
|
||||||
|
t.Error("Flush did not reach the underlying writer")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// TestHandlerRecordsLLMMetricsForStreaming is the end-to-end guard: it drives a
|
||||||
|
// real streaming request through the Handler, so it fails if the wrapper is ever
|
||||||
|
// constructed nowhere (the original defect) rather than merely being broken.
|
||||||
|
func TestHandlerRecordsLLMMetricsForStreaming(t *testing.T) {
|
||||||
|
// Unique model name keeps the observations isolated in the global registry.
|
||||||
|
model := fmt.Sprintf("metrics-test-%d", time.Now().UnixNano())
|
||||||
|
|
||||||
|
upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||||
|
w.Header().Set("Content-Type", "text/event-stream")
|
||||||
|
w.WriteHeader(http.StatusOK)
|
||||||
|
rc := http.NewResponseController(w)
|
||||||
|
for _, tok := range []string{"a", "b", "c"} {
|
||||||
|
time.Sleep(2 * time.Millisecond) // guarantee distinct ITL samples
|
||||||
|
fmt.Fprintf(w, "data: {\"token\":%q}\n\n", tok)
|
||||||
|
if err := rc.Flush(); err != nil {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
}
|
||||||
|
fmt.Fprint(w, "data: [DONE]\n\n")
|
||||||
|
_ = rc.Flush()
|
||||||
|
}))
|
||||||
|
defer upstream.Close()
|
||||||
|
|
||||||
|
cfg := &config.Config{
|
||||||
|
Models: map[string]*config.ModelUpstream{
|
||||||
|
model: {
|
||||||
|
Name: model,
|
||||||
|
Address: strings.TrimPrefix(upstream.URL, "http://"),
|
||||||
|
Path: "/v1/chat/completions",
|
||||||
|
},
|
||||||
|
},
|
||||||
|
Routes: make(map[string]*config.Route),
|
||||||
|
}
|
||||||
|
|
||||||
|
handler := New(cfg)
|
||||||
|
defer handler.Close()
|
||||||
|
srv := httptest.NewServer(handler)
|
||||||
|
defer srv.Close()
|
||||||
|
|
||||||
|
body := fmt.Sprintf(`{"model":%q,"stream":true,"messages":[{"role":"user","content":"hi"}]}`, model)
|
||||||
|
resp, err := http.Post(srv.URL+"/v1/chat/completions", "application/json", strings.NewReader(body))
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("request failed: %v", err)
|
||||||
|
}
|
||||||
|
defer resp.Body.Close()
|
||||||
|
if resp.StatusCode != http.StatusOK {
|
||||||
|
t.Fatalf("status = %d, want 200", resp.StatusCode)
|
||||||
|
}
|
||||||
|
if _, err := io.ReadAll(resp.Body); err != nil {
|
||||||
|
t.Fatalf("read body: %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
out := exportMetrics(observability.Global())
|
||||||
|
requireSeries(t, out, fmt.Sprintf(`llm_ttft_seconds_count{model=%q} 1`, model))
|
||||||
|
requireSeries(t, out, fmt.Sprintf(`llm_tokens_total{model=%q} 3`, model))
|
||||||
|
|
||||||
|
itlCount := regexp.MustCompile(fmt.Sprintf(`llm_itl_seconds_count\{model=%q\} (\d+)`, model))
|
||||||
|
match := itlCount.FindStringSubmatch(out)
|
||||||
|
if match == nil {
|
||||||
|
t.Errorf("no llm_itl_seconds samples recorded for %s\n--- got ---\n%s", model, out)
|
||||||
|
} else if match[1] == "0" {
|
||||||
|
t.Errorf("llm_itl_seconds_count = 0, want > 0")
|
||||||
|
}
|
||||||
|
}
|
||||||
+354
-148
@@ -2,15 +2,20 @@
|
|||||||
package proxy
|
package proxy
|
||||||
|
|
||||||
import (
|
import (
|
||||||
|
"context"
|
||||||
"encoding/json"
|
"encoding/json"
|
||||||
|
"errors"
|
||||||
"fmt"
|
"fmt"
|
||||||
"io"
|
"io"
|
||||||
"net"
|
"net"
|
||||||
"net/http"
|
"net/http"
|
||||||
"net/http/httputil"
|
"net/http/httputil"
|
||||||
"net/url"
|
"net/url"
|
||||||
|
"os"
|
||||||
"sort"
|
"sort"
|
||||||
|
"strconv"
|
||||||
"strings"
|
"strings"
|
||||||
|
"sync"
|
||||||
"syscall"
|
"syscall"
|
||||||
"time"
|
"time"
|
||||||
|
|
||||||
@@ -18,6 +23,7 @@ import (
|
|||||||
"forgejo.riotpiao.com/rock/homelab-frontend/internal/config"
|
"forgejo.riotpiao.com/rock/homelab-frontend/internal/config"
|
||||||
"forgejo.riotpiao.com/rock/homelab-frontend/internal/identity"
|
"forgejo.riotpiao.com/rock/homelab-frontend/internal/identity"
|
||||||
"forgejo.riotpiao.com/rock/homelab-frontend/internal/logging"
|
"forgejo.riotpiao.com/rock/homelab-frontend/internal/logging"
|
||||||
|
"forgejo.riotpiao.com/rock/homelab-frontend/internal/observability"
|
||||||
"forgejo.riotpiao.com/rock/homelab-frontend/internal/tracing"
|
"forgejo.riotpiao.com/rock/homelab-frontend/internal/tracing"
|
||||||
)
|
)
|
||||||
|
|
||||||
@@ -37,13 +43,20 @@ type Handler struct {
|
|||||||
defaultReadTimeout time.Duration
|
defaultReadTimeout time.Duration
|
||||||
defaultWriteTimeout time.Duration
|
defaultWriteTimeout time.Duration
|
||||||
defaultMaxBodySize int64
|
defaultMaxBodySize int64
|
||||||
|
// CAG: per-model backend schedulers for cache-aware routing (#45/#46)
|
||||||
|
schedulers map[string]*Scheduler
|
||||||
|
// CAG: delay applied to warmup requests before queue entry (#45)
|
||||||
|
warmupDelay time.Duration
|
||||||
|
// transportsMu guards h.transports for goroutine-safe reads/writes.
|
||||||
|
transportsMu sync.RWMutex
|
||||||
}
|
}
|
||||||
|
|
||||||
// Route represents a reverse proxy route.
|
// Route represents a reverse proxy route.
|
||||||
type Route struct {
|
type Route struct {
|
||||||
Name string
|
Name string
|
||||||
Upstream *config.Upstream
|
ModelName string // non-empty for model-based routes; used for scheduler lookup
|
||||||
Director func(*http.Request)
|
Upstream *config.Upstream
|
||||||
|
Director func(*http.Request)
|
||||||
Transport *http.Transport
|
Transport *http.Transport
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -93,6 +106,47 @@ func New(cfg *config.Config) *Handler {
|
|||||||
h.authHTTP = newAuthClient()
|
h.authHTTP = newAuthClient()
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// CAG: warmup delay from env (#45)
|
||||||
|
warmupMs := int64(500)
|
||||||
|
if v := os.Getenv("WARMUP_DELAY_MS"); v != "" {
|
||||||
|
if n, err := strconv.ParseInt(v, 10, 64); err == nil && n >= 0 {
|
||||||
|
warmupMs = n
|
||||||
|
}
|
||||||
|
}
|
||||||
|
h.warmupDelay = time.Duration(warmupMs) * time.Millisecond
|
||||||
|
|
||||||
|
// CAG: per-model backend schedulers (#45/#46)
|
||||||
|
maxSlots := 4
|
||||||
|
if v := os.Getenv("VLLM_MAX_NUM_SEQS"); v != "" {
|
||||||
|
if n, err := strconv.Atoi(v); err == nil && n > 0 {
|
||||||
|
maxSlots = n
|
||||||
|
}
|
||||||
|
}
|
||||||
|
maxQueue := 20
|
||||||
|
if v := os.Getenv("VLLM_MAX_QUEUE_DEPTH"); v != "" {
|
||||||
|
if n, err := strconv.Atoi(v); err == nil && n > 0 {
|
||||||
|
maxQueue = n
|
||||||
|
}
|
||||||
|
}
|
||||||
|
h.schedulers = make(map[string]*Scheduler)
|
||||||
|
for modelName, model := range cfg.Models {
|
||||||
|
addrs := model.AllAddresses()
|
||||||
|
if len(addrs) > 0 {
|
||||||
|
h.schedulers[modelName] = NewScheduler(addrs, maxSlots, maxQueue)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
// Pre-populate transports for all scheduler backends at startup (single-threaded).
|
||||||
|
// getOrCreateTransport is not goroutine-safe; pre-populating avoids concurrent writes.
|
||||||
|
for _, model := range cfg.Models {
|
||||||
|
for _, addr := range model.AllAddresses() {
|
||||||
|
h.getOrCreateTransport(addr, &config.Upstream{
|
||||||
|
ConnectTimeout: h.defaultConnectTimeout,
|
||||||
|
ReadTimeout: h.defaultReadTimeout,
|
||||||
|
WriteTimeout: h.defaultWriteTimeout,
|
||||||
|
})
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
for name, route := range cfg.Routes {
|
for name, route := range cfg.Routes {
|
||||||
// Create a transport per unique upstream address for connection reuse
|
// Create a transport per unique upstream address for connection reuse
|
||||||
transport := h.getOrCreateTransport(route.Upstream.Address, &route.Upstream)
|
transport := h.getOrCreateTransport(route.Upstream.Address, &route.Upstream)
|
||||||
@@ -116,7 +170,16 @@ func New(cfg *config.Config) *Handler {
|
|||||||
// getOrCreateTransport returns a shared http.Transport for the given upstream address.
|
// getOrCreateTransport returns a shared http.Transport for the given upstream address.
|
||||||
// This ensures connections are pooled and reused across requests to the same upstream.
|
// This ensures connections are pooled and reused across requests to the same upstream.
|
||||||
func (h *Handler) getOrCreateTransport(addr string, up *config.Upstream) *http.Transport {
|
func (h *Handler) getOrCreateTransport(addr string, up *config.Upstream) *http.Transport {
|
||||||
|
h.transportsMu.RLock()
|
||||||
if t, ok := h.transports[addr]; ok {
|
if t, ok := h.transports[addr]; ok {
|
||||||
|
h.transportsMu.RUnlock()
|
||||||
|
return t
|
||||||
|
}
|
||||||
|
h.transportsMu.RUnlock()
|
||||||
|
|
||||||
|
h.transportsMu.Lock()
|
||||||
|
defer h.transportsMu.Unlock()
|
||||||
|
if t, ok := h.transports[addr]; ok { // double-check after write lock
|
||||||
return t
|
return t
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -165,7 +228,7 @@ func (h *Handler) getOrCreateTransport(addr string, up *config.Upstream) *http.T
|
|||||||
}
|
}
|
||||||
|
|
||||||
h.transports[addr] = transport
|
h.transports[addr] = transport
|
||||||
return transport
|
return transport // written under transportsMu write lock
|
||||||
}
|
}
|
||||||
|
|
||||||
// directorFunc modifies the request to be sent to the upstream.
|
// directorFunc modifies the request to be sent to the upstream.
|
||||||
@@ -251,17 +314,11 @@ func writeProblemDetail(w http.ResponseWriter, status int, problemType, title, d
|
|||||||
|
|
||||||
// ServeHTTP implements http.Handler.
|
// ServeHTTP implements http.Handler.
|
||||||
func (h *Handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
|
func (h *Handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
|
||||||
// Auth endpoints — no JWT required (they issue tokens)
|
// Captured before any work so TTFT measures from the moment the gateway
|
||||||
if r.URL.Path == "/auth/token" {
|
// accepted the request, not from when the upstream response is copied.
|
||||||
h.handleAuthToken(w, r)
|
requestStart := time.Now()
|
||||||
return
|
|
||||||
}
|
if h.handleAuthEndpoint(w, r) {
|
||||||
if r.URL.Path == "/auth/refresh" {
|
|
||||||
h.handleAuthRefresh(w, r)
|
|
||||||
return
|
|
||||||
}
|
|
||||||
if r.URL.Path == "/auth/exchange" {
|
|
||||||
h.handleAuthExchange(w, r)
|
|
||||||
return
|
return
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -271,58 +328,153 @@ func (h *Handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
|
|||||||
return
|
return
|
||||||
}
|
}
|
||||||
|
|
||||||
// Try to find a matching route (including body-based dispatch for /v1/chat/completions).
|
// Resolve route; handles validation errors and 404.
|
||||||
// Note: /workflows endpoint is deprecated. Use X-Service: workflow + X-Resource headers instead.
|
route, ok := h.routeOrError(w, r)
|
||||||
|
if !ok {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
// Strip spoofed identity headers from all inbound requests.
|
||||||
|
// Must happen before any routing — even unauthenticated paths.
|
||||||
|
identity.StripIncoming(r)
|
||||||
|
|
||||||
|
// JWT Authentication for /v1/* endpoints.
|
||||||
|
if !h.authenticateV1(w, r) {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
// Enforce request body size cap for non-chat routes.
|
||||||
|
if !h.enforceBodyLimit(w, r, route) {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
// CAG: warmup deprioritisation + per-backend semaphore (issues #45/#46).
|
||||||
|
cacheKey := r.Header.Get("X-Poimen-Cache-Key")
|
||||||
|
sched := h.schedulers[route.ModelName]
|
||||||
|
proxy, backend, ok := h.setupCAGProxy(r.Context(), w, r, route, sched, cacheKey)
|
||||||
|
if !ok {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
if backend != nil {
|
||||||
|
defer backend.Release()
|
||||||
|
}
|
||||||
|
|
||||||
|
// For streaming responses (SSE, chunked), disable buffering. Issue #33.
|
||||||
|
w.Header().Set("X-Accel-Buffering", "no")
|
||||||
|
|
||||||
|
// Capture per-request LLM latency metrics (TTFT/ITL/tokens) while the
|
||||||
|
// response is copied from the upstream. Inert for non-SSE responses.
|
||||||
|
var respWriter http.ResponseWriter = w
|
||||||
|
if route.ModelName != "" {
|
||||||
|
respWriter = NewResponseWriterWrapper(w, route.ModelName, observability.Global(), requestStart)
|
||||||
|
}
|
||||||
|
|
||||||
|
proxy.ServeHTTP(respWriter, r)
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
|
// applyWarmup strips the X-Poimen-Warmup header, delays the request while
|
||||||
|
// honouring context cancellation, and records the metric.
|
||||||
|
// Returns false if the context was cancelled during the delay (caller must abort).
|
||||||
|
func (h *Handler) applyWarmup(ctx context.Context, r *http.Request) bool {
|
||||||
|
if r.Header.Get("X-Poimen-Warmup") != "true" {
|
||||||
|
return true
|
||||||
|
}
|
||||||
|
r.Header.Del("X-Poimen-Warmup")
|
||||||
|
r.Header.Set("X-Priority", "low")
|
||||||
|
select {
|
||||||
|
case <-time.After(h.warmupDelay):
|
||||||
|
case <-ctx.Done():
|
||||||
|
return false
|
||||||
|
}
|
||||||
|
observability.Global().IncrWarmupDelayed()
|
||||||
|
return true
|
||||||
|
}
|
||||||
|
|
||||||
|
// acquireScheduledBackend selects a backend from sched, applies warmup
|
||||||
|
// deprioritisation, acquires a concurrency slot, and patches route when the
|
||||||
|
// selected backend differs from the config default. It returns the backend
|
||||||
|
// to Release (caller must defer), the (possibly patched) route, and ok=true.
|
||||||
|
// If the slot cannot be acquired it writes a 503 and returns (nil, nil, false).
|
||||||
|
func (h *Handler) acquireScheduledBackend(ctx context.Context, w http.ResponseWriter, r *http.Request, sched *Scheduler, route *Route, cacheKey string) (*Backend, *Route, bool) {
|
||||||
|
backend := sched.Select(cacheKey)
|
||||||
|
if !h.applyWarmup(ctx, r) {
|
||||||
|
return nil, nil, false
|
||||||
|
}
|
||||||
|
if backend == nil {
|
||||||
|
return nil, route, true
|
||||||
|
}
|
||||||
|
if err := backend.Acquire(ctx); err != nil {
|
||||||
|
if errors.Is(err, ErrQueueFull) {
|
||||||
|
observability.Global().IncrQueueFull()
|
||||||
|
logging.Warnf("backend queue full", map[string]string{
|
||||||
|
"backend": backend.Address,
|
||||||
|
"path": r.URL.Path,
|
||||||
|
"model": route.ModelName,
|
||||||
|
})
|
||||||
|
w.Header().Set("Retry-After", "5")
|
||||||
|
w.WriteHeader(http.StatusServiceUnavailable)
|
||||||
|
fmt.Fprintf(w, "queue full")
|
||||||
|
} else {
|
||||||
|
logging.Warnf("backend acquire failed", map[string]string{
|
||||||
|
"backend": backend.Address,
|
||||||
|
"path": r.URL.Path,
|
||||||
|
"model": route.ModelName,
|
||||||
|
"err": err.Error(),
|
||||||
|
})
|
||||||
|
w.WriteHeader(http.StatusServiceUnavailable)
|
||||||
|
fmt.Fprintf(w, "upstream busy")
|
||||||
|
}
|
||||||
|
return nil, nil, false
|
||||||
|
}
|
||||||
|
if backend.Address != route.Upstream.Address {
|
||||||
|
patchedCfg := *route.Upstream
|
||||||
|
patchedCfg.Address = backend.Address
|
||||||
|
patchedURL, _ := url.Parse("http://" + backend.Address)
|
||||||
|
route = &Route{
|
||||||
|
Name: route.Name,
|
||||||
|
ModelName: route.ModelName,
|
||||||
|
Upstream: &patchedCfg,
|
||||||
|
Transport: h.getOrCreateTransport(backend.Address, &patchedCfg),
|
||||||
|
Director: func(req *http.Request) {
|
||||||
|
directorFunc(req, patchedURL, &patchedCfg)
|
||||||
|
},
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return backend, route, true
|
||||||
|
}
|
||||||
|
|
||||||
|
// handleAuthEndpoint routes /auth/* requests and returns true if handled.
|
||||||
|
func (h *Handler) handleAuthEndpoint(w http.ResponseWriter, r *http.Request) bool {
|
||||||
|
switch r.URL.Path {
|
||||||
|
case "/auth/token":
|
||||||
|
h.handleAuthToken(w, r)
|
||||||
|
return true
|
||||||
|
case "/auth/refresh":
|
||||||
|
h.handleAuthRefresh(w, r)
|
||||||
|
return true
|
||||||
|
case "/auth/exchange":
|
||||||
|
h.handleAuthExchange(w, r)
|
||||||
|
return true
|
||||||
|
}
|
||||||
|
return false
|
||||||
|
}
|
||||||
|
|
||||||
|
// routeOrError resolves a route for r, writing error responses on failure.
|
||||||
|
// Returns (route, true) on success or (nil, false) when a response has been written.
|
||||||
|
func (h *Handler) routeOrError(w http.ResponseWriter, r *http.Request) (*Route, bool) {
|
||||||
route, err := h.RouteRequest(r)
|
route, err := h.RouteRequest(r)
|
||||||
|
|
||||||
// Check if this is a model validation error (from body-based dispatch)
|
|
||||||
if validationErr, ok := err.(*modelValidationError); ok {
|
if validationErr, ok := err.(*modelValidationError); ok {
|
||||||
// This is a client error, not a routing error
|
h.writeModelValidationError(w, validationErr, h.getValidModels())
|
||||||
var status int
|
|
||||||
var problemType string
|
|
||||||
var title string
|
|
||||||
var detail string
|
|
||||||
|
|
||||||
switch validationErr.Kind {
|
|
||||||
case "invalid_json":
|
|
||||||
status = http.StatusBadRequest
|
|
||||||
problemType = "https://api.example.com/problems/invalid-request-body"
|
|
||||||
title = "Invalid Request Body"
|
|
||||||
detail = validationErr.Message
|
|
||||||
case "missing_model", "empty_model", "null_model":
|
|
||||||
status = http.StatusBadRequest
|
|
||||||
problemType = "https://api.example.com/problems/missing-model"
|
|
||||||
title = "Missing Model"
|
|
||||||
detail = "The 'model' field is required and must be a non-empty string"
|
|
||||||
case "unknown_model":
|
|
||||||
status = http.StatusBadRequest
|
|
||||||
problemType = "https://api.example.com/problems/unknown-model"
|
|
||||||
title = "Unknown Model"
|
|
||||||
detail = fmt.Sprintf("Model %q is not available. See valid_models for available options.", validationErr.Model)
|
|
||||||
default:
|
|
||||||
status = http.StatusBadRequest
|
|
||||||
problemType = "https://api.example.com/problems/invalid-request"
|
|
||||||
title = "Invalid Request"
|
|
||||||
detail = validationErr.Message
|
|
||||||
}
|
|
||||||
|
|
||||||
// Get list of valid models (only for model-related errors)
|
|
||||||
var validModels []string
|
|
||||||
if validationErr.Kind == "unknown_model" || validationErr.Kind == "missing_model" || validationErr.Kind == "empty_model" || validationErr.Kind == "null_model" {
|
|
||||||
validModels = h.getValidModels()
|
|
||||||
}
|
|
||||||
|
|
||||||
writeProblemDetail(w, status, problemType, title, detail, validModels)
|
|
||||||
logging.Errorf("client error", validationErr, map[string]string{
|
logging.Errorf("client error", validationErr, map[string]string{
|
||||||
"path": r.URL.Path,
|
"path": r.URL.Path,
|
||||||
"method": r.Method,
|
"method": r.Method,
|
||||||
"reason": validationErr.Kind,
|
"reason": validationErr.Kind,
|
||||||
})
|
})
|
||||||
return
|
return nil, false
|
||||||
}
|
}
|
||||||
|
|
||||||
if err != nil || route == nil {
|
if err != nil || route == nil {
|
||||||
// Route not found or error determining route
|
|
||||||
w.WriteHeader(http.StatusNotFound)
|
w.WriteHeader(http.StatusNotFound)
|
||||||
fmt.Fprintf(w, "not found")
|
fmt.Fprintf(w, "not found")
|
||||||
if err != nil {
|
if err != nil {
|
||||||
@@ -333,126 +485,180 @@ func (h *Handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
|
|||||||
} else {
|
} else {
|
||||||
logging.Errorf("no route matches", fmt.Errorf("path=%s method=%s", r.URL.Path, r.Method), nil)
|
logging.Errorf("no route matches", fmt.Errorf("path=%s method=%s", r.URL.Path, r.Method), nil)
|
||||||
}
|
}
|
||||||
return
|
return nil, false
|
||||||
|
}
|
||||||
|
return route, true
|
||||||
|
}
|
||||||
|
|
||||||
|
// writeModelValidationError writes a problem-detail response for model validation errors.
|
||||||
|
func (h *Handler) writeModelValidationError(w http.ResponseWriter, validationErr *modelValidationError, validModels []string) {
|
||||||
|
var status int
|
||||||
|
var problemType, title, detail string
|
||||||
|
var models []string
|
||||||
|
|
||||||
|
switch validationErr.Kind {
|
||||||
|
case "invalid_json":
|
||||||
|
status = http.StatusBadRequest
|
||||||
|
problemType = "https://api.example.com/problems/invalid-request-body"
|
||||||
|
title = "Invalid Request Body"
|
||||||
|
detail = validationErr.Message
|
||||||
|
case "missing_model", "empty_model", "null_model":
|
||||||
|
status = http.StatusBadRequest
|
||||||
|
problemType = "https://api.example.com/problems/missing-model"
|
||||||
|
title = "Missing Model"
|
||||||
|
detail = "The 'model' field is required and must be a non-empty string"
|
||||||
|
models = validModels
|
||||||
|
case "unknown_model":
|
||||||
|
status = http.StatusBadRequest
|
||||||
|
problemType = "https://api.example.com/problems/unknown-model"
|
||||||
|
title = "Unknown Model"
|
||||||
|
detail = fmt.Sprintf("Model %q is not available. See valid_models for available options.", validationErr.Model)
|
||||||
|
models = validModels
|
||||||
|
default:
|
||||||
|
status = http.StatusBadRequest
|
||||||
|
problemType = "https://api.example.com/problems/invalid-request"
|
||||||
|
title = "Invalid Request"
|
||||||
|
detail = validationErr.Message
|
||||||
}
|
}
|
||||||
|
|
||||||
// Strip spoofed identity headers from all inbound requests.
|
writeProblemDetail(w, status, problemType, title, detail, models)
|
||||||
// Must happen before any routing — even unauthenticated paths.
|
}
|
||||||
identity.StripIncoming(r)
|
|
||||||
|
|
||||||
// JWT Authentication for /v1/* endpoints
|
// authenticateV1 validates JWT for /v1/* paths.
|
||||||
if h.jwtValidator != nil && strings.HasPrefix(r.URL.Path, "/v1/") {
|
// Returns false when auth fails (response already written).
|
||||||
authHeader := r.Header.Get("Authorization")
|
func (h *Handler) authenticateV1(w http.ResponseWriter, r *http.Request) bool {
|
||||||
if authHeader == "" {
|
if h.jwtValidator == nil {
|
||||||
writeProblemDetail(w, http.StatusUnauthorized,
|
logging.Warnf("JWT auth disabled; /v1/* request unauthenticated", map[string]string{"path": r.URL.Path})
|
||||||
"https://api.example.com/problems/unauthorized",
|
return true
|
||||||
"Unauthorized",
|
}
|
||||||
"Authorization header required",
|
if !strings.HasPrefix(r.URL.Path, "/v1/") {
|
||||||
nil)
|
return true
|
||||||
logging.Errorf("auth failed", fmt.Errorf("missing auth header"), map[string]string{
|
}
|
||||||
"path": r.URL.Path,
|
authHeader := r.Header.Get("Authorization")
|
||||||
})
|
if authHeader == "" {
|
||||||
return
|
writeProblemDetail(w, http.StatusUnauthorized,
|
||||||
}
|
"https://api.example.com/problems/unauthorized",
|
||||||
|
"Unauthorized",
|
||||||
claims, err := h.jwtValidator.ValidateBearerToken(authHeader)
|
"Authorization header required",
|
||||||
if err != nil {
|
nil)
|
||||||
|
logging.Errorf("auth failed", fmt.Errorf("missing auth header"), map[string]string{
|
||||||
|
"path": r.URL.Path,
|
||||||
|
})
|
||||||
|
return false
|
||||||
|
}
|
||||||
|
claims, err := h.jwtValidator.ValidateBearerToken(authHeader)
|
||||||
|
if err != nil {
|
||||||
|
writeProblemDetail(w, http.StatusForbidden,
|
||||||
|
"https://api.example.com/problems/forbidden",
|
||||||
|
"Forbidden",
|
||||||
|
"JWT validation failed",
|
||||||
|
nil)
|
||||||
|
logging.Errorf("auth failed", err, map[string]string{
|
||||||
|
"path": r.URL.Path,
|
||||||
|
})
|
||||||
|
return false
|
||||||
|
}
|
||||||
|
identity.Inject(r, claims)
|
||||||
|
if h.config.Auth.RequiredCapability != "" {
|
||||||
|
if !h.jwtValidator.CheckPermissions(claims, h.config.Auth.RequiredCapability, "*") {
|
||||||
writeProblemDetail(w, http.StatusForbidden,
|
writeProblemDetail(w, http.StatusForbidden,
|
||||||
"https://api.example.com/problems/forbidden",
|
"https://api.example.com/problems/insufficient-permissions",
|
||||||
"Forbidden",
|
"Insufficient Permissions",
|
||||||
"JWT validation failed",
|
fmt.Sprintf("Required capability: %s", h.config.Auth.RequiredCapability),
|
||||||
nil)
|
nil)
|
||||||
logging.Errorf("auth failed", err, map[string]string{
|
logging.Errorf("auth failed", fmt.Errorf("insufficient permissions"), map[string]string{
|
||||||
"path": r.URL.Path,
|
"path": r.URL.Path,
|
||||||
|
"required": h.config.Auth.RequiredCapability,
|
||||||
})
|
})
|
||||||
return
|
return false
|
||||||
}
|
}
|
||||||
|
}
|
||||||
|
return true
|
||||||
|
}
|
||||||
|
|
||||||
// Inject identity headers for downstream services
|
// enforceBodyLimit caps request body size for non-chat routes.
|
||||||
identity.Inject(r, claims)
|
// Returns false when the request is rejected (response already written).
|
||||||
|
func (h *Handler) enforceBodyLimit(w http.ResponseWriter, r *http.Request, route *Route) bool {
|
||||||
|
if r.URL.Path == "/v1/chat/completions" {
|
||||||
|
return true
|
||||||
|
}
|
||||||
|
if route.Upstream.MaxBodySize > 0 && r.ContentLength > route.Upstream.MaxBodySize {
|
||||||
|
w.WriteHeader(http.StatusRequestEntityTooLarge)
|
||||||
|
fmt.Fprintf(w, "request body too large")
|
||||||
|
logging.Errorf("request rejected", fmt.Errorf("body_too_large"), map[string]string{
|
||||||
|
"reason": "body_too_large",
|
||||||
|
"route": route.Name,
|
||||||
|
"upstream": route.Upstream.Address,
|
||||||
|
"content_length": fmt.Sprintf("%d", r.ContentLength),
|
||||||
|
"max_body_size": fmt.Sprintf("%d", route.Upstream.MaxBodySize),
|
||||||
|
})
|
||||||
|
return false
|
||||||
|
}
|
||||||
|
if route.Upstream.MaxBodySize > 0 && r.Body != nil {
|
||||||
|
r.Body = io.NopCloser(io.LimitReader(r.Body, route.Upstream.MaxBodySize))
|
||||||
|
}
|
||||||
|
return true
|
||||||
|
}
|
||||||
|
|
||||||
// Check required capability if configured
|
// setupCAGProxy acquires a scheduled backend (if any), builds a reverse proxy,
|
||||||
if h.config.Auth.RequiredCapability != "" {
|
// wires ErrorHandler and ModifyResponse, and returns (proxy, backend, true).
|
||||||
if !h.jwtValidator.CheckPermissions(claims, h.config.Auth.RequiredCapability, "*") {
|
// Returns (nil, nil, false) when the backend slot cannot be acquired (response written).
|
||||||
writeProblemDetail(w, http.StatusForbidden,
|
func (h *Handler) setupCAGProxy(ctx context.Context, w http.ResponseWriter, r *http.Request, route *Route, sched *Scheduler, cacheKey string) (*httputil.ReverseProxy, *Backend, bool) {
|
||||||
"https://api.example.com/problems/insufficient-permissions",
|
var selectedBackend *Backend
|
||||||
"Insufficient Permissions",
|
if sched != nil {
|
||||||
fmt.Sprintf("Required capability: %s", h.config.Auth.RequiredCapability),
|
var ok bool
|
||||||
nil)
|
selectedBackend, route, ok = h.acquireScheduledBackend(ctx, w, r, sched, route, cacheKey)
|
||||||
logging.Errorf("auth failed", fmt.Errorf("insufficient permissions"), map[string]string{
|
if !ok {
|
||||||
"path": r.URL.Path,
|
return nil, nil, false
|
||||||
"required": h.config.Auth.RequiredCapability,
|
|
||||||
})
|
|
||||||
return
|
|
||||||
}
|
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
// Note: Body size checking already happened in RouteRequest (body was read for model dispatch).
|
|
||||||
// For other paths, we still need to enforce the cap.
|
|
||||||
// For /v1/chat/completions, the body was already read and validated.
|
|
||||||
|
|
||||||
// Enforce request body size cap for non-chat routes
|
|
||||||
if r.URL.Path != "/v1/chat/completions" {
|
|
||||||
if route.Upstream.MaxBodySize > 0 && r.ContentLength > route.Upstream.MaxBodySize {
|
|
||||||
w.WriteHeader(http.StatusRequestEntityTooLarge)
|
|
||||||
fmt.Fprintf(w, "request body too large")
|
|
||||||
logging.Errorf("request rejected", fmt.Errorf("body_too_large"), map[string]string{
|
|
||||||
"reason": "body_too_large",
|
|
||||||
"route": route.Name,
|
|
||||||
"upstream": route.Upstream.Address,
|
|
||||||
"content_length": fmt.Sprintf("%d", r.ContentLength),
|
|
||||||
"max_body_size": fmt.Sprintf("%d", route.Upstream.MaxBodySize),
|
|
||||||
})
|
|
||||||
return
|
|
||||||
}
|
|
||||||
|
|
||||||
// Wrap request body with size limiter
|
|
||||||
// This enforces the cap at read time, not after buffering
|
|
||||||
if route.Upstream.MaxBodySize > 0 && r.Body != nil {
|
|
||||||
r.Body = io.NopCloser(io.LimitReader(r.Body, route.Upstream.MaxBodySize))
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
// Create the reverse proxy
|
|
||||||
proxy := httputil.NewSingleHostReverseProxy(&url.URL{
|
proxy := httputil.NewSingleHostReverseProxy(&url.URL{
|
||||||
Scheme: "http",
|
Scheme: "http",
|
||||||
Host: route.Upstream.Address,
|
Host: route.Upstream.Address,
|
||||||
})
|
})
|
||||||
|
|
||||||
// Set the director to apply path rewriting
|
|
||||||
proxy.Director = route.Director
|
proxy.Director = route.Director
|
||||||
|
|
||||||
// Use the connection-pooled transport wrapped with tracing
|
|
||||||
proxy.Transport = tracing.NewTransport(route.Transport)
|
proxy.Transport = tracing.NewTransport(route.Transport)
|
||||||
|
|
||||||
// Set error handler to log upstream errors
|
capturedBackend := selectedBackend
|
||||||
|
capturedRoute := route
|
||||||
proxy.ErrorHandler = func(w http.ResponseWriter, r *http.Request, err error) {
|
proxy.ErrorHandler = func(w http.ResponseWriter, r *http.Request, err error) {
|
||||||
|
if capturedBackend != nil {
|
||||||
|
if s := h.schedulers[capturedRoute.ModelName]; s != nil {
|
||||||
|
s.MarkFailed(capturedBackend.Address)
|
||||||
|
}
|
||||||
|
}
|
||||||
w.WriteHeader(http.StatusBadGateway)
|
w.WriteHeader(http.StatusBadGateway)
|
||||||
fmt.Fprintf(w, "upstream error")
|
fmt.Fprintf(w, "upstream error")
|
||||||
logging.Errorf("upstream error", err, map[string]string{
|
logging.Errorf("upstream error", err, map[string]string{
|
||||||
"upstream": route.Upstream.Address,
|
"upstream": capturedRoute.Upstream.Address,
|
||||||
"path": r.URL.Path,
|
"path": r.URL.Path,
|
||||||
})
|
})
|
||||||
}
|
}
|
||||||
|
|
||||||
// Note: We apply connection timeout via Transport dialer, but NOT read timeout as a context deadline.
|
if cacheKey != "" && r.URL.Path == "/v1/chat/completions" {
|
||||||
// Read timeout should apply to inactivity (socket read timeout), not total request duration.
|
proxy.ModifyResponse = func(resp *http.Response) error {
|
||||||
// A streaming response that's continuously sending should not be cut off.
|
raw := resp.Header.Get("vllm-num-cached-tokens")
|
||||||
// The Transport's socket read timeout (via Dialer) handles inactivity timeouts.
|
if raw == "" {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
n, err := strconv.ParseInt(raw, 10, 64)
|
||||||
|
if err != nil || n <= 0 {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
observability.Global().IncrPrefixCacheHitTokens(n)
|
||||||
|
logging.Event(map[string]interface{}{
|
||||||
|
"event": "llm_request",
|
||||||
|
"cache_key": cacheKey,
|
||||||
|
"cached_tokens": n,
|
||||||
|
})
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
// For streaming responses (SSE, chunked), disable buffering to ensure events
|
return proxy, selectedBackend, true
|
||||||
// reach clients immediately. Issue #33: X-Accel-Buffering:no tells nginx/Ingress
|
|
||||||
// to stream instead of buffer. ResponseController.Flush() in upstream handler
|
|
||||||
// pairs with this to deliver unbuffered chunks.
|
|
||||||
w.Header().Set("X-Accel-Buffering", "no")
|
|
||||||
|
|
||||||
// Serve the request through the proxy
|
|
||||||
proxy.ServeHTTP(w, r)
|
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
// getValidModels returns a sorted list of all configured model names.
|
// getValidModels returns a sorted list of all configured model names.
|
||||||
func (h *Handler) getValidModels() []string {
|
func (h *Handler) getValidModels() []string {
|
||||||
var models []string
|
var models []string
|
||||||
|
|||||||
@@ -1,15 +1,24 @@
|
|||||||
package proxy
|
package proxy
|
||||||
|
|
||||||
import (
|
import (
|
||||||
|
"crypto/rand"
|
||||||
|
"crypto/rsa"
|
||||||
|
"encoding/base64"
|
||||||
|
"encoding/json"
|
||||||
"fmt"
|
"fmt"
|
||||||
"io"
|
"io"
|
||||||
|
"math/big"
|
||||||
"net/http"
|
"net/http"
|
||||||
"net/http/httptest"
|
"net/http/httptest"
|
||||||
|
"runtime"
|
||||||
"strings"
|
"strings"
|
||||||
"testing"
|
"testing"
|
||||||
"time"
|
"time"
|
||||||
|
|
||||||
|
"github.com/golang-jwt/jwt/v5"
|
||||||
|
|
||||||
"forgejo.riotpiao.com/rock/homelab-frontend/internal/config"
|
"forgejo.riotpiao.com/rock/homelab-frontend/internal/config"
|
||||||
|
"forgejo.riotpiao.com/rock/homelab-frontend/internal/observability"
|
||||||
)
|
)
|
||||||
|
|
||||||
func TestProxyBasic(t *testing.T) {
|
func TestProxyBasic(t *testing.T) {
|
||||||
@@ -414,3 +423,356 @@ func TestProxyPreservesBody(t *testing.T) {
|
|||||||
t.Errorf("expected body %s, got %s", testBody, receivedBody)
|
t.Errorf("expected body %s, got %s", testBody, receivedBody)
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// TestAuthenticateV1ProxyLevel exercises the three branches of authenticateV1:
|
||||||
|
// no-token→401, bad-token→403, valid-token→proxied (non-401/403).
|
||||||
|
func TestAuthenticateV1ProxyLevel(t *testing.T) {
|
||||||
|
// Generate RSA keypair for signing test JWTs.
|
||||||
|
privateKey, err := rsa.GenerateKey(rand.Reader, 2048)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("rsa.GenerateKey: %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
// Build minimal JWKS document from the RSA public key.
|
||||||
|
nBytes := privateKey.PublicKey.N.Bytes()
|
||||||
|
eBig := new(big.Int).SetInt64(int64(privateKey.PublicKey.E))
|
||||||
|
eBytes := eBig.Bytes()
|
||||||
|
jwksDoc, _ := json.Marshal(map[string]interface{}{
|
||||||
|
"keys": []map[string]interface{}{
|
||||||
|
{
|
||||||
|
"kty": "RSA",
|
||||||
|
"use": "sig",
|
||||||
|
"kid": "test-key",
|
||||||
|
"alg": "RS256",
|
||||||
|
"n": base64.RawURLEncoding.EncodeToString(nBytes),
|
||||||
|
"e": base64.RawURLEncoding.EncodeToString(eBytes),
|
||||||
|
},
|
||||||
|
},
|
||||||
|
})
|
||||||
|
|
||||||
|
// Serve JWKS from a local httptest server.
|
||||||
|
jwksSrv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||||
|
w.Header().Set("Content-Type", "application/json")
|
||||||
|
w.Write(jwksDoc)
|
||||||
|
}))
|
||||||
|
defer jwksSrv.Close()
|
||||||
|
|
||||||
|
// Upstream stub — records whether it was reached.
|
||||||
|
reached := make(chan struct{}, 1)
|
||||||
|
upstreamSrv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||||
|
select {
|
||||||
|
case reached <- struct{}{}:
|
||||||
|
default:
|
||||||
|
}
|
||||||
|
w.WriteHeader(http.StatusOK)
|
||||||
|
}))
|
||||||
|
defer upstreamSrv.Close()
|
||||||
|
upstreamAddr := strings.TrimPrefix(upstreamSrv.URL, "http://")
|
||||||
|
|
||||||
|
// signJWT returns a valid RS256 JWT accepted by the Authentik-pattern issuer check.
|
||||||
|
signJWT := func() string {
|
||||||
|
tok := jwt.NewWithClaims(jwt.SigningMethodRS256, jwt.MapClaims{
|
||||||
|
"iss": "https://authentik.riotpiao.com/application/o/test/",
|
||||||
|
"aud": "test-audience",
|
||||||
|
"sub": "test-user",
|
||||||
|
"exp": time.Now().Add(time.Hour).Unix(),
|
||||||
|
"iat": time.Now().Unix(),
|
||||||
|
})
|
||||||
|
tok.Header["kid"] = "test-key"
|
||||||
|
signed, err := tok.SignedString(privateKey)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("sign JWT: %v", err)
|
||||||
|
}
|
||||||
|
return signed
|
||||||
|
}
|
||||||
|
|
||||||
|
// Build proxy with auth enabled.
|
||||||
|
cfg := &config.Config{
|
||||||
|
Routes: make(map[string]*config.Route),
|
||||||
|
Models: map[string]*config.ModelUpstream{
|
||||||
|
"auth-model": {
|
||||||
|
Name: "auth-model",
|
||||||
|
Address: upstreamAddr,
|
||||||
|
},
|
||||||
|
},
|
||||||
|
Auth: config.AuthConfig{
|
||||||
|
Enabled: true,
|
||||||
|
JWKSURL: jwksSrv.URL,
|
||||||
|
},
|
||||||
|
}
|
||||||
|
h := New(cfg)
|
||||||
|
defer h.Close()
|
||||||
|
srv := httptest.NewServer(h)
|
||||||
|
defer srv.Close()
|
||||||
|
|
||||||
|
mkReq := func(authHeader string) *http.Request {
|
||||||
|
body := strings.NewReader(`{"model":"auth-model","messages":[]}`)
|
||||||
|
req, _ := http.NewRequest("POST", srv.URL+"/v1/chat/completions", body)
|
||||||
|
req.Header.Set("Content-Type", "application/json")
|
||||||
|
if authHeader != "" {
|
||||||
|
req.Header.Set("Authorization", authHeader)
|
||||||
|
}
|
||||||
|
return req
|
||||||
|
}
|
||||||
|
|
||||||
|
// Branch 1: no Authorization header → 401.
|
||||||
|
t.Run("no_token_returns_401", func(t *testing.T) {
|
||||||
|
resp, err := http.DefaultClient.Do(mkReq(""))
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("request failed: %v", err)
|
||||||
|
}
|
||||||
|
defer resp.Body.Close()
|
||||||
|
if resp.StatusCode != http.StatusUnauthorized {
|
||||||
|
t.Errorf("expected 401, got %d", resp.StatusCode)
|
||||||
|
}
|
||||||
|
})
|
||||||
|
|
||||||
|
// Branch 2: malformed/invalid JWT → 403.
|
||||||
|
t.Run("bad_token_returns_403", func(t *testing.T) {
|
||||||
|
resp, err := http.DefaultClient.Do(mkReq("Bearer this.is.garbage"))
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("request failed: %v", err)
|
||||||
|
}
|
||||||
|
defer resp.Body.Close()
|
||||||
|
if resp.StatusCode != http.StatusForbidden {
|
||||||
|
t.Errorf("expected 403, got %d", resp.StatusCode)
|
||||||
|
}
|
||||||
|
})
|
||||||
|
|
||||||
|
// Branch 3: valid RS256 JWT → proxied to upstream (not 401/403).
|
||||||
|
t.Run("valid_token_proxied", func(t *testing.T) {
|
||||||
|
resp, err := http.DefaultClient.Do(mkReq("Bearer " + signJWT()))
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("request failed: %v", err)
|
||||||
|
}
|
||||||
|
defer resp.Body.Close()
|
||||||
|
if resp.StatusCode == http.StatusUnauthorized || resp.StatusCode == http.StatusForbidden {
|
||||||
|
body, _ := io.ReadAll(resp.Body)
|
||||||
|
t.Errorf("expected proxied response, got %d body=%s", resp.StatusCode, body)
|
||||||
|
}
|
||||||
|
select {
|
||||||
|
case <-reached:
|
||||||
|
// upstream was hit — auth passed
|
||||||
|
default:
|
||||||
|
t.Error("upstream was not reached after valid JWT")
|
||||||
|
}
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
// TestProxyQueueFull verifies that a 503 is returned and the queue_full metric
|
||||||
|
// is incremented when the backend queue is at capacity.
|
||||||
|
//
|
||||||
|
// Setup: maxSlots=1, maxQueue=1. Request 1 holds the slot. Request 2 fills the
|
||||||
|
// queue (blocks waiting for slot). Request 3 gets ErrQueueFull → 503.
|
||||||
|
func TestProxyQueueFull(t *testing.T) {
|
||||||
|
t.Setenv("VLLM_MAX_NUM_SEQS", "1") // 1 concurrent slot
|
||||||
|
t.Setenv("VLLM_MAX_QUEUE_DEPTH", "1") // queue capacity 1
|
||||||
|
t.Setenv("WARMUP_DELAY_MS", "0")
|
||||||
|
observability.Global().Reset()
|
||||||
|
|
||||||
|
// Upstream blocks until released; buffers one signal when it first starts.
|
||||||
|
slotHeld := make(chan struct{}, 1)
|
||||||
|
release := make(chan struct{})
|
||||||
|
upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||||
|
select {
|
||||||
|
case slotHeld <- struct{}{}:
|
||||||
|
default:
|
||||||
|
}
|
||||||
|
<-release
|
||||||
|
w.WriteHeader(http.StatusOK)
|
||||||
|
}))
|
||||||
|
defer upstream.Close()
|
||||||
|
|
||||||
|
upstreamAddr := strings.TrimPrefix(upstream.URL, "http://")
|
||||||
|
cfg := &config.Config{
|
||||||
|
Routes: make(map[string]*config.Route),
|
||||||
|
Models: map[string]*config.ModelUpstream{
|
||||||
|
"q-model": {
|
||||||
|
Name: "q-model",
|
||||||
|
Address: upstreamAddr,
|
||||||
|
},
|
||||||
|
},
|
||||||
|
}
|
||||||
|
h := New(cfg)
|
||||||
|
defer h.Close()
|
||||||
|
|
||||||
|
srv := httptest.NewServer(h)
|
||||||
|
defer srv.Close()
|
||||||
|
|
||||||
|
mkReq := func() *http.Request {
|
||||||
|
body := strings.NewReader(`{"model":"q-model","messages":[]}`)
|
||||||
|
req, _ := http.NewRequest("POST", srv.URL+"/v1/chat/completions", body)
|
||||||
|
req.Header.Set("Content-Type", "application/json")
|
||||||
|
return req
|
||||||
|
}
|
||||||
|
|
||||||
|
// Request 1: acquires the only slot, reaches upstream, blocks there.
|
||||||
|
r1ch := make(chan *http.Response, 1)
|
||||||
|
go func() {
|
||||||
|
resp, _ := http.DefaultClient.Do(mkReq())
|
||||||
|
r1ch <- resp
|
||||||
|
}()
|
||||||
|
<-slotHeld // slot is acquired; upstream is holding request 1
|
||||||
|
|
||||||
|
// Request 2: enters the queue (capacity 1), blocks waiting for slot.
|
||||||
|
r2ch := make(chan *http.Response, 1)
|
||||||
|
go func() {
|
||||||
|
resp, _ := http.DefaultClient.Do(mkReq())
|
||||||
|
r2ch <- resp
|
||||||
|
}()
|
||||||
|
// Spin until request 2 enters the backend queue channel.
|
||||||
|
for len(h.schedulers["q-model"].backends[0].queue) == 0 {
|
||||||
|
runtime.Gosched()
|
||||||
|
}
|
||||||
|
|
||||||
|
// Request 3: queue is full → expect 503 and metric increment.
|
||||||
|
resp3, err := http.DefaultClient.Do(mkReq())
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("request 3 failed: %v", err)
|
||||||
|
}
|
||||||
|
defer resp3.Body.Close()
|
||||||
|
|
||||||
|
if resp3.StatusCode != http.StatusServiceUnavailable {
|
||||||
|
t.Errorf("expected 503 queue full, got %d", resp3.StatusCode)
|
||||||
|
}
|
||||||
|
if got := observability.Global().GetQueueFull(); got != 1 {
|
||||||
|
t.Errorf("expected GetQueueFull()=1, got %d", got)
|
||||||
|
}
|
||||||
|
|
||||||
|
// Cleanup: unblock upstream so goroutines can finish.
|
||||||
|
close(release)
|
||||||
|
if r := <-r1ch; r != nil {
|
||||||
|
r.Body.Close()
|
||||||
|
}
|
||||||
|
if r := <-r2ch; r != nil {
|
||||||
|
r.Body.Close()
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// TestProxyCacheHitLogging verifies that when a response contains the
|
||||||
|
// vllm-num-cached-tokens header and a cache key was set, the prefix cache
|
||||||
|
// hit token count is incremented in observability metrics.
|
||||||
|
func TestProxyCacheHitLogging(t *testing.T) {
|
||||||
|
observability.Global().Reset()
|
||||||
|
|
||||||
|
upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||||
|
w.Header().Set("vllm-num-cached-tokens", "128")
|
||||||
|
w.WriteHeader(http.StatusOK)
|
||||||
|
}))
|
||||||
|
defer upstream.Close()
|
||||||
|
|
||||||
|
upstreamAddr := strings.TrimPrefix(upstream.URL, "http://")
|
||||||
|
cfg := &config.Config{
|
||||||
|
Routes: make(map[string]*config.Route),
|
||||||
|
Models: map[string]*config.ModelUpstream{
|
||||||
|
"cache-model": {
|
||||||
|
Name: "cache-model",
|
||||||
|
Address: upstreamAddr,
|
||||||
|
},
|
||||||
|
},
|
||||||
|
}
|
||||||
|
h := New(cfg)
|
||||||
|
defer h.Close()
|
||||||
|
|
||||||
|
srv := httptest.NewServer(h)
|
||||||
|
defer srv.Close()
|
||||||
|
|
||||||
|
body := strings.NewReader(`{"model":"cache-model","messages":[]}`)
|
||||||
|
req, _ := http.NewRequest("POST", srv.URL+"/v1/chat/completions", body)
|
||||||
|
req.Header.Set("Content-Type", "application/json")
|
||||||
|
req.Header.Set("X-Poimen-Cache-Key", "testkey123")
|
||||||
|
|
||||||
|
resp, err := http.DefaultClient.Do(req)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("request failed: %v", err)
|
||||||
|
}
|
||||||
|
resp.Body.Close()
|
||||||
|
|
||||||
|
if got := observability.Global().GetPrefixCacheHitTokens(); got != 128 {
|
||||||
|
t.Errorf("expected PrefixCacheHitTokens=128, got %d", got)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// TestWarmupHeaderHandling verifies that X-Poimen-Warmup: true is stripped from
|
||||||
|
// the upstream request and replaced with X-Priority: low, while a request without
|
||||||
|
// X-Poimen-Warmup reaches upstream without either header.
|
||||||
|
//
|
||||||
|
// The warmup path runs only for model-based routes (cfg.Models + scheduler),
|
||||||
|
// so we POST to /v1/chat/completions with a JSON body that names the test model.
|
||||||
|
// WARMUP_DELAY_MS=0 eliminates the 500 ms sleep.
|
||||||
|
func TestWarmupHeaderHandling(t *testing.T) {
|
||||||
|
t.Setenv("WARMUP_DELAY_MS", "0")
|
||||||
|
|
||||||
|
tests := []struct {
|
||||||
|
name string
|
||||||
|
sendWarmup bool
|
||||||
|
wantWarmupGone bool
|
||||||
|
wantPriorityLow bool
|
||||||
|
}{
|
||||||
|
{
|
||||||
|
name: "warmup=true",
|
||||||
|
sendWarmup: true,
|
||||||
|
wantWarmupGone: true, // X-Poimen-Warmup must NOT reach upstream
|
||||||
|
wantPriorityLow: true, // X-Priority: low must be set
|
||||||
|
},
|
||||||
|
{
|
||||||
|
name: "warmup=false",
|
||||||
|
sendWarmup: false,
|
||||||
|
wantWarmupGone: true, // no header sent, so absent at upstream
|
||||||
|
wantPriorityLow: false, // X-Priority must NOT be set
|
||||||
|
},
|
||||||
|
}
|
||||||
|
|
||||||
|
for _, tc := range tests {
|
||||||
|
t.Run(tc.name, func(t *testing.T) {
|
||||||
|
var gotWarmup, gotPriority string
|
||||||
|
upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||||
|
gotWarmup = r.Header.Get("X-Poimen-Warmup")
|
||||||
|
gotPriority = r.Header.Get("X-Priority")
|
||||||
|
w.WriteHeader(http.StatusOK)
|
||||||
|
}))
|
||||||
|
defer upstream.Close()
|
||||||
|
|
||||||
|
upstreamAddr := strings.TrimPrefix(upstream.URL, "http://")
|
||||||
|
// Warmup code only runs for model-based routes (scheduler path).
|
||||||
|
// Use cfg.Models so New() builds a Scheduler for "warmup-model".
|
||||||
|
cfg := &config.Config{
|
||||||
|
Routes: make(map[string]*config.Route),
|
||||||
|
Models: map[string]*config.ModelUpstream{
|
||||||
|
"warmup-model": {
|
||||||
|
Name: "warmup-model",
|
||||||
|
Address: upstreamAddr,
|
||||||
|
},
|
||||||
|
},
|
||||||
|
}
|
||||||
|
h := New(cfg)
|
||||||
|
defer h.Close()
|
||||||
|
|
||||||
|
srv := httptest.NewServer(h)
|
||||||
|
defer srv.Close()
|
||||||
|
|
||||||
|
// POST to /v1/chat/completions — triggers model-based dispatch and warmup path.
|
||||||
|
body := strings.NewReader(`{"model":"warmup-model","messages":[]}`)
|
||||||
|
req, _ := http.NewRequest("POST", srv.URL+"/v1/chat/completions", body)
|
||||||
|
req.Header.Set("Content-Type", "application/json")
|
||||||
|
if tc.sendWarmup {
|
||||||
|
req.Header.Set("X-Poimen-Warmup", "true")
|
||||||
|
}
|
||||||
|
resp, err := http.DefaultClient.Do(req)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("request failed: %v", err)
|
||||||
|
}
|
||||||
|
resp.Body.Close()
|
||||||
|
|
||||||
|
if tc.wantWarmupGone && gotWarmup != "" {
|
||||||
|
t.Errorf("X-Poimen-Warmup must be stripped before upstream, got %q", gotWarmup)
|
||||||
|
}
|
||||||
|
if tc.wantPriorityLow && gotPriority != "low" {
|
||||||
|
t.Errorf("X-Priority: low must be set when warmup=true, got %q", gotPriority)
|
||||||
|
}
|
||||||
|
if !tc.wantPriorityLow && gotPriority != "" {
|
||||||
|
t.Errorf("X-Priority must not be set for normal requests, got %q", gotPriority)
|
||||||
|
}
|
||||||
|
})
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|||||||
@@ -159,6 +159,7 @@ func (h *Handler) routeByModel(r *http.Request, path string) (*Route, error) {
|
|||||||
|
|
||||||
route := &Route{
|
route := &Route{
|
||||||
Name: "v1-chat-" + modelName,
|
Name: "v1-chat-" + modelName,
|
||||||
|
ModelName: modelName,
|
||||||
Upstream: upstreamCfg,
|
Upstream: upstreamCfg,
|
||||||
Transport: h.getOrCreateTransport(modelUpstream.Address, upstreamCfg),
|
Transport: h.getOrCreateTransport(modelUpstream.Address, upstreamCfg),
|
||||||
Director: func(req *http.Request) {
|
Director: func(req *http.Request) {
|
||||||
|
|||||||
@@ -0,0 +1,157 @@
|
|||||||
|
package proxy
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"errors"
|
||||||
|
"hash/fnv"
|
||||||
|
"sort"
|
||||||
|
"sync"
|
||||||
|
"sync/atomic"
|
||||||
|
"time"
|
||||||
|
)
|
||||||
|
|
||||||
|
// ErrQueueFull is returned when the in-memory queue is at capacity.
|
||||||
|
var ErrQueueFull = errors.New("backend queue full")
|
||||||
|
|
||||||
|
// Backend represents a single vLLM upstream with a bounded in-memory queue.
|
||||||
|
// Two channels model the queue:
|
||||||
|
//
|
||||||
|
// queue — waiting requests (bounded by VLLM_MAX_QUEUE_DEPTH)
|
||||||
|
// slots — concurrent in-service requests (bounded by VLLM_MAX_NUM_SEQS)
|
||||||
|
type Backend struct {
|
||||||
|
Address string
|
||||||
|
queue chan struct{}
|
||||||
|
slots chan struct{}
|
||||||
|
}
|
||||||
|
|
||||||
|
// newBackend creates a Backend with given queue depth and concurrency limit.
|
||||||
|
func newBackend(addr string, maxSlots, maxQueue int) *Backend {
|
||||||
|
slots := make(chan struct{}, maxSlots)
|
||||||
|
for i := 0; i < maxSlots; i++ {
|
||||||
|
slots <- struct{}{}
|
||||||
|
}
|
||||||
|
return &Backend{
|
||||||
|
Address: addr,
|
||||||
|
queue: make(chan struct{}, maxQueue),
|
||||||
|
slots: slots,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// Acquire enters the in-memory queue then waits for a vLLM concurrency slot.
|
||||||
|
// Returns ErrQueueFull immediately if the queue channel is at capacity (no goroutine queued).
|
||||||
|
// Returns ctx.Err() if the context is cancelled while waiting for a slot.
|
||||||
|
func (b *Backend) Acquire(ctx context.Context) error {
|
||||||
|
// Non-blocking queue entry: reject immediately when queue is full.
|
||||||
|
select {
|
||||||
|
case b.queue <- struct{}{}:
|
||||||
|
// entered queue
|
||||||
|
default:
|
||||||
|
return ErrQueueFull
|
||||||
|
}
|
||||||
|
// Wait for a concurrency slot, then leave the queue.
|
||||||
|
select {
|
||||||
|
case <-b.slots:
|
||||||
|
<-b.queue // leave queue; now in service
|
||||||
|
return nil
|
||||||
|
case <-ctx.Done():
|
||||||
|
<-b.queue // leave queue; cancelled
|
||||||
|
return ctx.Err()
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// Release returns a concurrency slot to the pool after the upstream request completes.
|
||||||
|
func (b *Backend) Release() {
|
||||||
|
b.slots <- struct{}{}
|
||||||
|
}
|
||||||
|
|
||||||
|
// Scheduler routes LLM requests to backends using rendezvous hashing (cache-aware)
|
||||||
|
// or round-robin (no cache key). Failed backends are excluded for a configurable window.
|
||||||
|
type Scheduler struct {
|
||||||
|
mu sync.RWMutex
|
||||||
|
backends []*Backend
|
||||||
|
failedUntil map[string]time.Time
|
||||||
|
failWindow time.Duration
|
||||||
|
rrIdx atomic.Uint64
|
||||||
|
}
|
||||||
|
|
||||||
|
// NewScheduler creates a Scheduler for the given addresses.
|
||||||
|
// maxSlots = VLLM_MAX_NUM_SEQS, maxQueue = VLLM_MAX_QUEUE_DEPTH.
|
||||||
|
func NewScheduler(addresses []string, maxSlots, maxQueue int) *Scheduler {
|
||||||
|
backends := make([]*Backend, len(addresses))
|
||||||
|
for i, addr := range addresses {
|
||||||
|
backends[i] = newBackend(addr, maxSlots, maxQueue)
|
||||||
|
}
|
||||||
|
return &Scheduler{
|
||||||
|
backends: backends,
|
||||||
|
failedUntil: make(map[string]time.Time),
|
||||||
|
failWindow: 30 * time.Second,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// Select returns a backend. Uses rendezvous hashing when cacheKey is non-empty
|
||||||
|
// (same key always routes to the same backend, skipping unhealthy ones).
|
||||||
|
// Uses round-robin when cacheKey is empty.
|
||||||
|
func (s *Scheduler) Select(cacheKey string) *Backend {
|
||||||
|
s.mu.RLock()
|
||||||
|
defer s.mu.RUnlock()
|
||||||
|
|
||||||
|
if len(s.backends) == 0 {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
if len(s.backends) == 1 {
|
||||||
|
return s.backends[0]
|
||||||
|
}
|
||||||
|
if cacheKey != "" {
|
||||||
|
return s.selectRendezvous(cacheKey)
|
||||||
|
}
|
||||||
|
return s.selectRoundRobin()
|
||||||
|
}
|
||||||
|
|
||||||
|
// selectRendezvous picks the healthy backend with the highest FNV32a score
|
||||||
|
// for the key (highest random weight / rendezvous hashing).
|
||||||
|
// Falls back to next-highest score when primary is excluded by MarkFailed.
|
||||||
|
// Must be called with mu held.
|
||||||
|
func (s *Scheduler) selectRendezvous(cacheKey string) *Backend {
|
||||||
|
type scored struct {
|
||||||
|
b *Backend
|
||||||
|
score uint32
|
||||||
|
}
|
||||||
|
now := time.Now()
|
||||||
|
scores := make([]scored, len(s.backends))
|
||||||
|
h := fnv.New32a()
|
||||||
|
for i, b := range s.backends {
|
||||||
|
h.Reset()
|
||||||
|
h.Write([]byte(cacheKey + ":" + b.Address))
|
||||||
|
scores[i] = scored{b: b, score: h.Sum32()}
|
||||||
|
}
|
||||||
|
sort.Slice(scores, func(i, j int) bool {
|
||||||
|
return scores[i].score > scores[j].score
|
||||||
|
})
|
||||||
|
for _, sc := range scores {
|
||||||
|
_, failed := s.failedUntil[sc.b.Address]
|
||||||
|
if !failed {
|
||||||
|
return sc.b
|
||||||
|
}
|
||||||
|
expiry := s.failedUntil[sc.b.Address]
|
||||||
|
if now.After(expiry) {
|
||||||
|
return sc.b
|
||||||
|
}
|
||||||
|
}
|
||||||
|
// All backends unhealthy: return highest-scored (fail-open).
|
||||||
|
return scores[0].b
|
||||||
|
}
|
||||||
|
|
||||||
|
// selectRoundRobin distributes requests evenly when no cache key is present.
|
||||||
|
// Must be called with mu held.
|
||||||
|
func (s *Scheduler) selectRoundRobin() *Backend {
|
||||||
|
idx := s.rrIdx.Add(1) - 1
|
||||||
|
return s.backends[idx%uint64(len(s.backends))]
|
||||||
|
}
|
||||||
|
|
||||||
|
// MarkFailed excludes addr from routing for the fail window (default 30 s).
|
||||||
|
// Called by the proxy ErrorHandler on upstream connection failure.
|
||||||
|
func (s *Scheduler) MarkFailed(addr string) {
|
||||||
|
s.mu.Lock()
|
||||||
|
defer s.mu.Unlock()
|
||||||
|
s.failedUntil[addr] = time.Now().Add(s.failWindow)
|
||||||
|
}
|
||||||
@@ -0,0 +1,182 @@
|
|||||||
|
package proxy
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"errors"
|
||||||
|
"runtime"
|
||||||
|
"sync"
|
||||||
|
"testing"
|
||||||
|
"time"
|
||||||
|
)
|
||||||
|
|
||||||
|
// TestBackendAcquireBlocksThenUnblocksOnRelease verifies that a second Acquire
|
||||||
|
// blocks while a slot is held, then unblocks when Release is called.
|
||||||
|
func TestBackendAcquireBlocksThenUnblocksOnRelease(t *testing.T) {
|
||||||
|
b := newBackend("host:1", 1, 1)
|
||||||
|
|
||||||
|
// G1: acquire and hold the only slot.
|
||||||
|
if err := b.Acquire(context.Background()); err != nil {
|
||||||
|
t.Fatalf("G1 Acquire: unexpected error: %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
g2Done := make(chan error, 1)
|
||||||
|
go func() {
|
||||||
|
g2Done <- b.Acquire(context.Background())
|
||||||
|
}()
|
||||||
|
|
||||||
|
// Wait until G2 enters the queue channel (proves Acquire is blocking on slots).
|
||||||
|
for len(b.queue) == 0 {
|
||||||
|
runtime.Gosched()
|
||||||
|
}
|
||||||
|
|
||||||
|
select {
|
||||||
|
case err := <-g2Done:
|
||||||
|
t.Fatalf("G2 returned before Release: %v", err)
|
||||||
|
default:
|
||||||
|
}
|
||||||
|
|
||||||
|
// G1 releases; G2 should unblock.
|
||||||
|
b.Release()
|
||||||
|
|
||||||
|
select {
|
||||||
|
case err := <-g2Done:
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("G2 Acquire after Release: unexpected error: %v", err)
|
||||||
|
}
|
||||||
|
case <-time.After(time.Second):
|
||||||
|
t.Fatal("G2 did not unblock within 1s after Release")
|
||||||
|
}
|
||||||
|
|
||||||
|
// Cleanup: release G2's slot.
|
||||||
|
b.Release()
|
||||||
|
}
|
||||||
|
|
||||||
|
// TestBackendAcquireReturnsErrQueueFullWhenQueueFull verifies that Acquire returns
|
||||||
|
// ErrQueueFull immediately when no queue capacity exists (maxQueue=0).
|
||||||
|
// With a zero-capacity queue channel the non-blocking select always hits default.
|
||||||
|
func TestBackendAcquireReturnsErrQueueFullWhenQueueFull(t *testing.T) {
|
||||||
|
b := newBackend("host:1", 1, 0)
|
||||||
|
|
||||||
|
err := b.Acquire(context.Background())
|
||||||
|
if !errors.Is(err, ErrQueueFull) {
|
||||||
|
t.Fatalf("expected ErrQueueFull, got %v", err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// TestSchedulerSelectRendezvousHashDeterministic verifies that the same cache key
|
||||||
|
// always selects the same backend (rendezvous / highest-random-weight hashing).
|
||||||
|
func TestSchedulerSelectRendezvousHashDeterministic(t *testing.T) {
|
||||||
|
addrs := []string{"a:80", "b:80", "c:80"}
|
||||||
|
s := NewScheduler(addrs, 10, 10)
|
||||||
|
|
||||||
|
first := s.Select("same-cache-key")
|
||||||
|
if first == nil {
|
||||||
|
t.Fatal("Select returned nil")
|
||||||
|
}
|
||||||
|
for i := 1; i < 10; i++ {
|
||||||
|
b := s.Select("same-cache-key")
|
||||||
|
if b == nil || b.Address != first.Address {
|
||||||
|
t.Fatalf("call %d: want %s, got %v", i, first.Address, b)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// TestSchedulerSelectNoKeyRoundRobin verifies that round-robin distributes
|
||||||
|
// requests evenly across all backends when no cache key is given.
|
||||||
|
func TestSchedulerSelectNoKeyRoundRobin(t *testing.T) {
|
||||||
|
addrs := []string{"a:80", "b:80", "c:80"}
|
||||||
|
s := NewScheduler(addrs, 10, 10)
|
||||||
|
|
||||||
|
counts := make(map[string]int)
|
||||||
|
for i := 0; i < 9; i++ {
|
||||||
|
b := s.Select("")
|
||||||
|
if b == nil {
|
||||||
|
t.Fatalf("call %d: Select returned nil", i)
|
||||||
|
}
|
||||||
|
counts[b.Address]++
|
||||||
|
}
|
||||||
|
|
||||||
|
for _, addr := range addrs {
|
||||||
|
if counts[addr] != 3 {
|
||||||
|
t.Errorf("addr %s: want 3 selections, got %d (all counts: %v)", addr, counts[addr], counts)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// TestBackendAcquireCancelledContext verifies that Acquire returns ctx.Err() when the
|
||||||
|
// context is cancelled while the caller is blocked waiting for a concurrency slot,
|
||||||
|
// and that the queue depth returns to 0 after cancellation.
|
||||||
|
func TestBackendAcquireCancelledContext(t *testing.T) {
|
||||||
|
b := newBackend("host:1", 1, 1)
|
||||||
|
|
||||||
|
// G1: acquire and hold the only slot indefinitely.
|
||||||
|
if err := b.Acquire(context.Background()); err != nil {
|
||||||
|
t.Fatalf("G1 Acquire: unexpected error: %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
ctx, cancel := context.WithCancel(context.Background())
|
||||||
|
defer cancel()
|
||||||
|
|
||||||
|
g2Done := make(chan error, 1)
|
||||||
|
go func() {
|
||||||
|
g2Done <- b.Acquire(ctx)
|
||||||
|
}()
|
||||||
|
|
||||||
|
// Wait until G2 enters the queue channel.
|
||||||
|
for len(b.queue) == 0 {
|
||||||
|
runtime.Gosched()
|
||||||
|
}
|
||||||
|
|
||||||
|
// Verify G2 is queued (queue channel has one item).
|
||||||
|
if qd := len(b.queue); qd != 1 {
|
||||||
|
t.Errorf("expected queue depth 1 while G2 blocked, got %d", qd)
|
||||||
|
}
|
||||||
|
|
||||||
|
// Cancel ctx while G2 is waiting for the slot.
|
||||||
|
cancel()
|
||||||
|
|
||||||
|
select {
|
||||||
|
case err := <-g2Done:
|
||||||
|
if err == nil {
|
||||||
|
t.Fatal("G2 Acquire: expected ctx.Err(), got nil")
|
||||||
|
}
|
||||||
|
if !errors.Is(err, context.Canceled) {
|
||||||
|
t.Fatalf("G2 Acquire: expected context.Canceled, got %v", err)
|
||||||
|
}
|
||||||
|
case <-time.After(time.Second):
|
||||||
|
t.Fatal("G2 did not return within 1s after context cancel")
|
||||||
|
}
|
||||||
|
|
||||||
|
// Queue depth must return to 0 after G2 is drained.
|
||||||
|
if qd := len(b.queue); qd != 0 {
|
||||||
|
t.Errorf("expected queue depth 0 after cancellation, got %d", qd)
|
||||||
|
}
|
||||||
|
|
||||||
|
// Cleanup: release G1's slot.
|
||||||
|
b.Release()
|
||||||
|
}
|
||||||
|
|
||||||
|
// TestSchedulerMarkFailedExcludesBackend verifies that a failed backend is not
|
||||||
|
// selected within the fail window (30 s default).
|
||||||
|
func TestSchedulerMarkFailedExcludesBackend(t *testing.T) {
|
||||||
|
addrs := []string{"a:80", "b:80"}
|
||||||
|
s := NewScheduler(addrs, 10, 10)
|
||||||
|
|
||||||
|
s.MarkFailed("a:80")
|
||||||
|
|
||||||
|
var mu sync.Mutex
|
||||||
|
seen := make(map[string]int)
|
||||||
|
for i := 0; i < 10; i++ {
|
||||||
|
b := s.Select("any-key")
|
||||||
|
if b == nil {
|
||||||
|
t.Fatalf("call %d: Select returned nil", i)
|
||||||
|
}
|
||||||
|
mu.Lock()
|
||||||
|
seen[b.Address]++
|
||||||
|
mu.Unlock()
|
||||||
|
}
|
||||||
|
|
||||||
|
if seen["a:80"] > 0 {
|
||||||
|
t.Errorf("failed backend a:80 was selected %d times; expected 0 within fail window", seen["a:80"])
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -16,6 +16,7 @@ type Router struct {
|
|||||||
temporalHandler http.Handler
|
temporalHandler http.Handler
|
||||||
upstreamHandler http.Handler
|
upstreamHandler http.Handler
|
||||||
forgejoWebhook *webhook.ForgejoHandler
|
forgejoWebhook *webhook.ForgejoHandler
|
||||||
|
metricsHandler http.Handler
|
||||||
}
|
}
|
||||||
|
|
||||||
// NewRouter creates a new router with health endpoints.
|
// NewRouter creates a new router with health endpoints.
|
||||||
@@ -33,6 +34,11 @@ func NewRouter(healthChecker *HealthChecker, dispatcher *serviceadapter.Dispatch
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// SetMetricsHandler registers the Prometheus /metrics endpoint handler.
|
||||||
|
func (r *Router) SetMetricsHandler(h http.Handler) {
|
||||||
|
r.metricsHandler = h
|
||||||
|
}
|
||||||
|
|
||||||
// ServeHTTP implements http.Handler.
|
// ServeHTTP implements http.Handler.
|
||||||
// Priority order:
|
// Priority order:
|
||||||
// 1. /healthz and /readyz to health handlers
|
// 1. /healthz and /readyz to health handlers
|
||||||
@@ -48,6 +54,11 @@ func (r *Router) ServeHTTP(w http.ResponseWriter, req *http.Request) {
|
|||||||
case "/readyz":
|
case "/readyz":
|
||||||
ReadinessHandler(r.healthChecker)(w, req)
|
ReadinessHandler(r.healthChecker)(w, req)
|
||||||
return
|
return
|
||||||
|
case "/metrics":
|
||||||
|
if r.metricsHandler != nil {
|
||||||
|
r.metricsHandler.ServeHTTP(w, req)
|
||||||
|
return
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
// X-Service (ServiceAdapter) routing - checked before path-based routing
|
// X-Service (ServiceAdapter) routing - checked before path-based routing
|
||||||
|
|||||||
@@ -28,6 +28,9 @@ func AdapterFactory(serviceName string) interface{} {
|
|||||||
return &IAMAdapter{}
|
return &IAMAdapter{}
|
||||||
case "memory":
|
case "memory":
|
||||||
return &MemoryAdapter{}
|
return &MemoryAdapter{}
|
||||||
|
case "sessions":
|
||||||
|
// code-registered; see NewSessionAdapter()
|
||||||
|
return nil
|
||||||
default:
|
default:
|
||||||
return nil
|
return nil
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -0,0 +1,67 @@
|
|||||||
|
package serviceadapter
|
||||||
|
|
||||||
|
import "os"
|
||||||
|
|
||||||
|
// NewMemoryAdapter creates a code-registered ServiceAdapter for the memory service.
|
||||||
|
// Takes precedence over any config-file loaded memory adapter (main.go skips config
|
||||||
|
// adapters when a code-registered one with the same name exists).
|
||||||
|
// Upstream URL configured via MEMORY_UPSTREAM_URL env var.
|
||||||
|
func NewMemoryAdapter() *ServiceAdapter {
|
||||||
|
upstreamURL := os.Getenv("MEMORY_UPSTREAM_URL")
|
||||||
|
if upstreamURL == "" {
|
||||||
|
upstreamURL = "http://poimen-memory.poimen.svc.cluster.local:8080"
|
||||||
|
}
|
||||||
|
return &ServiceAdapter{
|
||||||
|
Namespace: "poimen",
|
||||||
|
ServiceName: "memory",
|
||||||
|
Spec: *GetMemorySpec(upstreamURL),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// GetMemorySpec returns the ServiceAdapter spec for the memory service.
|
||||||
|
func GetMemorySpec(upstreamURL string) *Spec {
|
||||||
|
return &Spec{
|
||||||
|
ServiceName: "memory",
|
||||||
|
Upstream: Upstream{
|
||||||
|
URL: upstreamURL,
|
||||||
|
TimeoutSeconds: 30,
|
||||||
|
},
|
||||||
|
Auth: Auth{
|
||||||
|
Required: true,
|
||||||
|
Capability: "memory:read",
|
||||||
|
},
|
||||||
|
Resources: []Resource{
|
||||||
|
{
|
||||||
|
Name: "ingest",
|
||||||
|
Methods: []Method{{
|
||||||
|
Verb: "POST",
|
||||||
|
UpstreamPath: "/memory/ingest",
|
||||||
|
Auth: &Auth{Required: true, Capability: "memory:write"},
|
||||||
|
}},
|
||||||
|
},
|
||||||
|
{
|
||||||
|
Name: "query",
|
||||||
|
Methods: []Method{{
|
||||||
|
Verb: "POST",
|
||||||
|
UpstreamPath: "/memory/query",
|
||||||
|
Auth: &Auth{Required: true, Capability: "memory:read"},
|
||||||
|
}},
|
||||||
|
},
|
||||||
|
{
|
||||||
|
Name: "skills",
|
||||||
|
Methods: []Method{{
|
||||||
|
Verb: "GET",
|
||||||
|
UpstreamPath: "/memory/skills",
|
||||||
|
}},
|
||||||
|
},
|
||||||
|
{
|
||||||
|
Name: "standing-context",
|
||||||
|
Methods: []Method{{
|
||||||
|
Verb: "POST",
|
||||||
|
UpstreamPath: "/memory/standing-context",
|
||||||
|
Auth: &Auth{Required: true, Capability: "memory:read"},
|
||||||
|
}},
|
||||||
|
},
|
||||||
|
},
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -53,17 +53,20 @@ func GetPoimenSpec() *Spec {
|
|||||||
},
|
},
|
||||||
},
|
},
|
||||||
},
|
},
|
||||||
|
// Note: workflow_id is passed as ?workflow_id=xxx query param,
|
||||||
|
// not as a path segment. The ServiceAdapter dispatcher does not
|
||||||
|
// substitute path parameters.
|
||||||
{
|
{
|
||||||
Name: "execution-detail",
|
Name: "execution-detail",
|
||||||
Methods: []Method{
|
Methods: []Method{
|
||||||
{
|
{
|
||||||
Verb: "GET",
|
Verb: "GET",
|
||||||
UpstreamPath: "/v1/poimen/executions/{id}",
|
UpstreamPath: "/v1/poimen/executions/detail",
|
||||||
Auth: &Auth{Required: true, Capability: "workflow:read"},
|
Auth: &Auth{Required: true, Capability: "workflow:read"},
|
||||||
},
|
},
|
||||||
{
|
{
|
||||||
Verb: "DELETE",
|
Verb: "DELETE",
|
||||||
UpstreamPath: "/v1/poimen/executions/{id}",
|
UpstreamPath: "/v1/poimen/executions/detail",
|
||||||
Auth: &Auth{Required: true, Capability: "workflow:execute"},
|
Auth: &Auth{Required: true, Capability: "workflow:execute"},
|
||||||
},
|
},
|
||||||
},
|
},
|
||||||
@@ -72,7 +75,7 @@ func GetPoimenSpec() *Spec {
|
|||||||
Name: "execution-history",
|
Name: "execution-history",
|
||||||
Methods: []Method{{
|
Methods: []Method{{
|
||||||
Verb: "GET",
|
Verb: "GET",
|
||||||
UpstreamPath: "/v1/poimen/executions/{id}/history",
|
UpstreamPath: "/v1/poimen/executions/history",
|
||||||
Auth: &Auth{Required: true, Capability: "workflow:read"},
|
Auth: &Auth{Required: true, Capability: "workflow:read"},
|
||||||
}},
|
}},
|
||||||
},
|
},
|
||||||
@@ -80,7 +83,7 @@ func GetPoimenSpec() *Spec {
|
|||||||
Name: "execution-stream",
|
Name: "execution-stream",
|
||||||
Methods: []Method{{
|
Methods: []Method{{
|
||||||
Verb: "GET",
|
Verb: "GET",
|
||||||
UpstreamPath: "/v1/poimen/executions/{id}/stream",
|
UpstreamPath: "/v1/poimen/executions/stream",
|
||||||
Auth: &Auth{Required: true, Capability: "workflow:read"},
|
Auth: &Auth{Required: true, Capability: "workflow:read"},
|
||||||
}},
|
}},
|
||||||
},
|
},
|
||||||
@@ -88,7 +91,7 @@ func GetPoimenSpec() *Spec {
|
|||||||
Name: "execution-signal",
|
Name: "execution-signal",
|
||||||
Methods: []Method{{
|
Methods: []Method{{
|
||||||
Verb: "POST",
|
Verb: "POST",
|
||||||
UpstreamPath: "/v1/poimen/executions/{id}/signal",
|
UpstreamPath: "/v1/poimen/executions/signal",
|
||||||
Auth: &Auth{Required: true, Capability: "workflow:execute"},
|
Auth: &Auth{Required: true, Capability: "workflow:execute"},
|
||||||
}},
|
}},
|
||||||
},
|
},
|
||||||
|
|||||||
@@ -16,6 +16,7 @@ import (
|
|||||||
|
|
||||||
"forgejo.riotpiao.com/rock/homelab-frontend/internal/auth"
|
"forgejo.riotpiao.com/rock/homelab-frontend/internal/auth"
|
||||||
"forgejo.riotpiao.com/rock/homelab-frontend/internal/identity"
|
"forgejo.riotpiao.com/rock/homelab-frontend/internal/identity"
|
||||||
|
"forgejo.riotpiao.com/rock/homelab-frontend/internal/logging"
|
||||||
"forgejo.riotpiao.com/rock/homelab-frontend/internal/problem"
|
"forgejo.riotpiao.com/rock/homelab-frontend/internal/problem"
|
||||||
)
|
)
|
||||||
|
|
||||||
@@ -74,8 +75,13 @@ func (d *Dispatcher) Dispatch(w http.ResponseWriter, r *http.Request) {
|
|||||||
}
|
}
|
||||||
|
|
||||||
// JWT auth enforcement for adapters that require it
|
// JWT auth enforcement for adapters that require it
|
||||||
if adapter.Spec.Auth.Required && d.jwtValidator != nil {
|
if adapter.Spec.Auth.Required {
|
||||||
if !d.authenticate(w, r, serviceName, method.Verb) {
|
if d.jwtValidator == nil {
|
||||||
|
logging.Warnf("auth required but jwtValidator nil; request unauthenticated", map[string]string{
|
||||||
|
"service": serviceName,
|
||||||
|
"resource": resourceName,
|
||||||
|
})
|
||||||
|
} else if !d.authenticate(w, r, serviceName, method) {
|
||||||
return
|
return
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
@@ -98,7 +104,7 @@ func (d *Dispatcher) Dispatch(w http.ResponseWriter, r *http.Request) {
|
|||||||
|
|
||||||
// authenticate validates the JWT and checks service-level capability.
|
// authenticate validates the JWT and checks service-level capability.
|
||||||
// Returns false (and writes error response) if auth fails.
|
// Returns false (and writes error response) if auth fails.
|
||||||
func (d *Dispatcher) authenticate(w http.ResponseWriter, r *http.Request, serviceName, verb string) bool {
|
func (d *Dispatcher) authenticate(w http.ResponseWriter, r *http.Request, serviceName string, method *Method) bool {
|
||||||
authHeader := r.Header.Get("Authorization")
|
authHeader := r.Header.Get("Authorization")
|
||||||
if authHeader == "" {
|
if authHeader == "" {
|
||||||
d.writeError(w, problem.NewProblem(http.StatusUnauthorized,
|
d.writeError(w, problem.NewProblem(http.StatusUnauthorized,
|
||||||
@@ -115,8 +121,11 @@ func (d *Dispatcher) authenticate(w http.ResponseWriter, r *http.Request, servic
|
|||||||
return false
|
return false
|
||||||
}
|
}
|
||||||
|
|
||||||
// Check capability: <service>:read for GET/HEAD, <service>:write for mutating verbs
|
// Check capability: method-level Auth.Capability if set; else derive from verb.
|
||||||
required := capabilityForVerb(serviceName, verb)
|
required := capabilityForVerb(serviceName, method.Verb)
|
||||||
|
if method.Auth != nil && method.Auth.Capability != "" {
|
||||||
|
required = method.Auth.Capability
|
||||||
|
}
|
||||||
if !d.jwtValidator.CheckPermissions(claims, required, "*") {
|
if !d.jwtValidator.CheckPermissions(claims, required, "*") {
|
||||||
d.writeError(w, problem.NewProblem(http.StatusForbidden,
|
d.writeError(w, problem.NewProblem(http.StatusForbidden,
|
||||||
"about:blank#insufficient-permissions", "Insufficient Permissions",
|
"about:blank#insufficient-permissions", "Insufficient Permissions",
|
||||||
|
|||||||
@@ -1,11 +1,19 @@
|
|||||||
package serviceadapter
|
package serviceadapter
|
||||||
|
|
||||||
import (
|
import (
|
||||||
|
"crypto/rand"
|
||||||
|
"crypto/rsa"
|
||||||
|
"encoding/base64"
|
||||||
"encoding/json"
|
"encoding/json"
|
||||||
|
"math/big"
|
||||||
"net/http"
|
"net/http"
|
||||||
"net/http/httptest"
|
"net/http/httptest"
|
||||||
"testing"
|
"testing"
|
||||||
|
"time"
|
||||||
|
|
||||||
|
"github.com/golang-jwt/jwt/v5"
|
||||||
|
|
||||||
|
"forgejo.riotpiao.com/rock/homelab-frontend/internal/auth"
|
||||||
"forgejo.riotpiao.com/rock/homelab-frontend/internal/identity"
|
"forgejo.riotpiao.com/rock/homelab-frontend/internal/identity"
|
||||||
)
|
)
|
||||||
|
|
||||||
@@ -263,3 +271,197 @@ func TestDispatch_IdentityHeadersNotSet_WhenNoAuth(t *testing.T) {
|
|||||||
t.Errorf("X-Auth-Verified should not be set without auth, got %q", gotVerified)
|
t.Errorf("X-Auth-Verified should not be set without auth, got %q", gotVerified)
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// TestDispatch_MethodCapabilityOverridesPOSTDefault verifies that a method-level
|
||||||
|
// Auth.Capability (e.g. "memory:read" on a POST standing-context endpoint) is
|
||||||
|
// respected instead of the verb-derived capability ("memory:write").
|
||||||
|
// Because auth.Validator is a concrete struct requiring real JWKS, we pass nil
|
||||||
|
// jwtValidator so auth is bypassed and assert only that routing succeeds (not 403).
|
||||||
|
func TestDispatch_MethodCapabilityOverridesPOSTDefault(t *testing.T) {
|
||||||
|
// Upstream captures the request so we can verify it was routed.
|
||||||
|
upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||||
|
w.WriteHeader(http.StatusOK)
|
||||||
|
}))
|
||||||
|
defer upstream.Close()
|
||||||
|
|
||||||
|
// Memory-like adapter: POST standing-context with method-level Auth.Capability=memory:read.
|
||||||
|
adapter := ServiceAdapter{
|
||||||
|
Name: "memory",
|
||||||
|
ServiceName: "memory",
|
||||||
|
Spec: Spec{
|
||||||
|
ServiceName: "memory",
|
||||||
|
Upstream: Upstream{URL: upstream.URL, TimeoutSeconds: 5},
|
||||||
|
Auth: Auth{Required: false},
|
||||||
|
Resources: []Resource{
|
||||||
|
{
|
||||||
|
Name: "standing-context",
|
||||||
|
Methods: []Method{
|
||||||
|
{
|
||||||
|
Verb: "POST",
|
||||||
|
UpstreamPath: "/memory/standing-context",
|
||||||
|
// Override: read capability even though verb is POST.
|
||||||
|
Auth: &Auth{Capability: "memory:read"},
|
||||||
|
},
|
||||||
|
},
|
||||||
|
},
|
||||||
|
},
|
||||||
|
},
|
||||||
|
}
|
||||||
|
|
||||||
|
// nil jwtValidator: auth.Required=false + nil validator → auth skipped entirely.
|
||||||
|
d := NewDispatcher(newTestRegistry(adapter), nil)
|
||||||
|
w := httptest.NewRecorder()
|
||||||
|
r := httptest.NewRequest("POST", "/", nil)
|
||||||
|
r.Header.Set("X-Service", "memory")
|
||||||
|
r.Header.Set("X-Resource", "standing-context")
|
||||||
|
|
||||||
|
d.Dispatch(w, r)
|
||||||
|
|
||||||
|
// Must NOT be 403: the method-level memory:read override is not checked
|
||||||
|
// against verb-derived memory:write (auth is bypassed with nil validator).
|
||||||
|
// Must be 200 routed to real upstream.
|
||||||
|
if w.Code == http.StatusForbidden {
|
||||||
|
t.Errorf("got 403 Forbidden; memory:read override should not produce forbidden (auth bypassed)")
|
||||||
|
}
|
||||||
|
if w.Code != http.StatusOK {
|
||||||
|
t.Errorf("expected 200 from upstream, got %d", w.Code)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// TestAuthenticate_MethodCapabilityOverrideDirect calls d.Dispatch() with a real
|
||||||
|
// signed RS256 JWT so authenticate() executes end-to-end and the capability-override
|
||||||
|
// branch (router.go lines 124-130) is reached and covered.
|
||||||
|
//
|
||||||
|
// Sub-test 1: token has "memory:read", method.Auth.Capability="memory:read" → 200.
|
||||||
|
// Sub-test 2: token has "memory:read", no method override → POST derives
|
||||||
|
//
|
||||||
|
// "memory:write" → CheckPermissions fails → 403.
|
||||||
|
func TestAuthenticate_MethodCapabilityOverrideDirect(t *testing.T) {
|
||||||
|
// Generate an RSA key pair for signing test JWTs.
|
||||||
|
privateKey, err := rsa.GenerateKey(rand.Reader, 2048)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("rsa.GenerateKey: %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
// Build minimal JWKS document from the RSA public key.
|
||||||
|
nBytes := privateKey.PublicKey.N.Bytes()
|
||||||
|
eBig := new(big.Int).SetInt64(int64(privateKey.PublicKey.E))
|
||||||
|
eBytes := eBig.Bytes()
|
||||||
|
jwksDoc, _ := json.Marshal(map[string]interface{}{
|
||||||
|
"keys": []map[string]interface{}{
|
||||||
|
{
|
||||||
|
"kty": "RSA",
|
||||||
|
"use": "sig",
|
||||||
|
"kid": "test-key",
|
||||||
|
"alg": "RS256",
|
||||||
|
"n": base64.RawURLEncoding.EncodeToString(nBytes),
|
||||||
|
"e": base64.RawURLEncoding.EncodeToString(eBytes),
|
||||||
|
},
|
||||||
|
},
|
||||||
|
})
|
||||||
|
|
||||||
|
// Serve JWKS from a local test server.
|
||||||
|
jwksSrv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||||
|
w.Header().Set("Content-Type", "application/json")
|
||||||
|
w.Write(jwksDoc)
|
||||||
|
}))
|
||||||
|
defer jwksSrv.Close()
|
||||||
|
|
||||||
|
// Upstream that returns 200 when authenticate() passes.
|
||||||
|
upstreamSrv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||||
|
w.WriteHeader(http.StatusOK)
|
||||||
|
}))
|
||||||
|
defer upstreamSrv.Close()
|
||||||
|
|
||||||
|
// signJWT returns a signed RS256 token with the given permissions slice.
|
||||||
|
signJWT := func(permissions []string) string {
|
||||||
|
tok := jwt.NewWithClaims(jwt.SigningMethodRS256, jwt.MapClaims{
|
||||||
|
"iss": "https://authentik.riotpiao.com/application/o/test/",
|
||||||
|
"aud": "test-audience",
|
||||||
|
"sub": "test-user",
|
||||||
|
"exp": time.Now().Add(time.Hour).Unix(),
|
||||||
|
"iat": time.Now().Unix(),
|
||||||
|
"permissions": permissions,
|
||||||
|
})
|
||||||
|
tok.Header["kid"] = "test-key"
|
||||||
|
signed, err := tok.SignedString(privateKey)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("sign JWT: %v", err)
|
||||||
|
}
|
||||||
|
return signed
|
||||||
|
}
|
||||||
|
|
||||||
|
validator := auth.NewValidator("", "", jwksSrv.URL)
|
||||||
|
|
||||||
|
// Sub-test 1: method.Auth.Capability="memory:read"; token has "memory:read".
|
||||||
|
// authenticate() must execute the override branch and pass → upstream returns 200.
|
||||||
|
t.Run("capability_override_match_reaches_upstream", func(t *testing.T) {
|
||||||
|
adapter := ServiceAdapter{
|
||||||
|
Name: "memory",
|
||||||
|
ServiceName: "memory",
|
||||||
|
Spec: Spec{
|
||||||
|
ServiceName: "memory",
|
||||||
|
Upstream: Upstream{URL: upstreamSrv.URL, TimeoutSeconds: 5},
|
||||||
|
Auth: Auth{Required: true},
|
||||||
|
Resources: []Resource{
|
||||||
|
{
|
||||||
|
Name: "standing-context",
|
||||||
|
Methods: []Method{
|
||||||
|
{
|
||||||
|
Verb: "POST",
|
||||||
|
UpstreamPath: "/memory/standing-context",
|
||||||
|
Auth: &Auth{Capability: "memory:read"}, // override
|
||||||
|
},
|
||||||
|
},
|
||||||
|
},
|
||||||
|
},
|
||||||
|
},
|
||||||
|
}
|
||||||
|
d := NewDispatcher(newTestRegistry(adapter), validator)
|
||||||
|
w := httptest.NewRecorder()
|
||||||
|
r := httptest.NewRequest("POST", "/", nil)
|
||||||
|
r.Header.Set("X-Service", "memory")
|
||||||
|
r.Header.Set("X-Resource", "standing-context")
|
||||||
|
r.Header.Set("Authorization", "Bearer "+signJWT([]string{"memory:read"}))
|
||||||
|
d.Dispatch(w, r)
|
||||||
|
if w.Code != http.StatusOK {
|
||||||
|
t.Errorf("expected 200 (override memory:read match), got %d body=%s", w.Code, w.Body.String())
|
||||||
|
}
|
||||||
|
})
|
||||||
|
|
||||||
|
// Sub-test 2: no method-level Auth override; POST → capabilityForVerb → "memory:write".
|
||||||
|
// Token has only "memory:read" → CheckPermissions fails → 403.
|
||||||
|
t.Run("verb_derived_capability_mismatch_returns_403", func(t *testing.T) {
|
||||||
|
adapter := ServiceAdapter{
|
||||||
|
Name: "memory",
|
||||||
|
ServiceName: "memory",
|
||||||
|
Spec: Spec{
|
||||||
|
ServiceName: "memory",
|
||||||
|
Upstream: Upstream{URL: upstreamSrv.URL, TimeoutSeconds: 5},
|
||||||
|
Auth: Auth{Required: true},
|
||||||
|
Resources: []Resource{
|
||||||
|
{
|
||||||
|
Name: "data",
|
||||||
|
Methods: []Method{
|
||||||
|
{
|
||||||
|
Verb: "POST",
|
||||||
|
UpstreamPath: "/memory/data",
|
||||||
|
// No method-level Auth: POST derives "memory:write".
|
||||||
|
},
|
||||||
|
},
|
||||||
|
},
|
||||||
|
},
|
||||||
|
},
|
||||||
|
}
|
||||||
|
d := NewDispatcher(newTestRegistry(adapter), validator)
|
||||||
|
w := httptest.NewRecorder()
|
||||||
|
r := httptest.NewRequest("POST", "/", nil)
|
||||||
|
r.Header.Set("X-Service", "memory")
|
||||||
|
r.Header.Set("X-Resource", "data")
|
||||||
|
r.Header.Set("Authorization", "Bearer "+signJWT([]string{"memory:read"}))
|
||||||
|
d.Dispatch(w, r)
|
||||||
|
if w.Code != http.StatusForbidden {
|
||||||
|
t.Errorf("expected 403 (write required, read only in token), got %d body=%s", w.Code, w.Body.String())
|
||||||
|
}
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|||||||
@@ -0,0 +1,123 @@
|
|||||||
|
package serviceadapter
|
||||||
|
|
||||||
|
import (
|
||||||
|
"forgejo.riotpiao.com/rock/homelab-frontend/internal/sessions"
|
||||||
|
)
|
||||||
|
|
||||||
|
// NewSessionAdapter creates a ServiceAdapter for the sessions service.
|
||||||
|
//
|
||||||
|
// Route mapping (all via X-Service: sessions):
|
||||||
|
//
|
||||||
|
// X-Resource Verb upstream path handler
|
||||||
|
// ──────────────────────────────────────────────────────────────
|
||||||
|
// sessions POST /v1/sessions/create create session
|
||||||
|
// sessions GET /v1/sessions/list list sessions (?status=,?cwd=)
|
||||||
|
// session GET /v1/sessions/get get session (?session_id=)
|
||||||
|
// session PATCH /v1/sessions/patch update session(?session_id=)
|
||||||
|
// milestones POST /v1/sessions/milestones/create create milestone(?session_id=)
|
||||||
|
// episodes GET /v1/sessions/episodes/list list episodes (?session_id=)
|
||||||
|
// episode GET /v1/sessions/episodes/get get episode (?session_id=,?episode_id=)
|
||||||
|
//
|
||||||
|
// Path params (:id, :eid) are passed as query params because the ServiceAdapter
|
||||||
|
// dispatcher does not substitute URL path parameters.
|
||||||
|
func NewSessionAdapter() *ServiceAdapter {
|
||||||
|
return &ServiceAdapter{
|
||||||
|
Namespace: "poimen",
|
||||||
|
ServiceName: "sessions",
|
||||||
|
Handler: sessions.NewHandler(),
|
||||||
|
Spec: *GetSessionSpec(),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// GetSessionSpec returns the ServiceAdapter spec for the sessions service.
|
||||||
|
func GetSessionSpec() *Spec {
|
||||||
|
return &Spec{
|
||||||
|
ServiceName: "sessions",
|
||||||
|
Upstream: Upstream{
|
||||||
|
URL: "internal://sessions",
|
||||||
|
TimeoutSeconds: 30,
|
||||||
|
},
|
||||||
|
Auth: Auth{
|
||||||
|
Required: true,
|
||||||
|
Capability: "sessions:read",
|
||||||
|
},
|
||||||
|
Resources: []Resource{
|
||||||
|
{
|
||||||
|
Name: "sessions",
|
||||||
|
Methods: []Method{
|
||||||
|
{
|
||||||
|
Verb: "POST",
|
||||||
|
UpstreamPath: "/v1/sessions/create",
|
||||||
|
Auth: &Auth{Required: true, Capability: "sessions:write"},
|
||||||
|
},
|
||||||
|
{
|
||||||
|
Verb: "GET",
|
||||||
|
UpstreamPath: "/v1/sessions/list",
|
||||||
|
Auth: &Auth{Required: true, Capability: "sessions:read"},
|
||||||
|
},
|
||||||
|
},
|
||||||
|
},
|
||||||
|
{
|
||||||
|
Name: "session",
|
||||||
|
Methods: []Method{
|
||||||
|
{
|
||||||
|
Verb: "GET",
|
||||||
|
UpstreamPath: "/v1/sessions/get",
|
||||||
|
Auth: &Auth{Required: true, Capability: "sessions:read"},
|
||||||
|
},
|
||||||
|
{
|
||||||
|
Verb: "PATCH",
|
||||||
|
UpstreamPath: "/v1/sessions/patch",
|
||||||
|
Auth: &Auth{Required: true, Capability: "sessions:write"},
|
||||||
|
},
|
||||||
|
},
|
||||||
|
},
|
||||||
|
{
|
||||||
|
Name: "milestones",
|
||||||
|
Methods: []Method{
|
||||||
|
{
|
||||||
|
Verb: "POST",
|
||||||
|
UpstreamPath: "/v1/sessions/milestones/create",
|
||||||
|
Auth: &Auth{Required: true, Capability: "sessions:write"},
|
||||||
|
},
|
||||||
|
},
|
||||||
|
},
|
||||||
|
{
|
||||||
|
Name: "episodes",
|
||||||
|
Methods: []Method{
|
||||||
|
{
|
||||||
|
Verb: "GET",
|
||||||
|
UpstreamPath: "/v1/sessions/episodes/list",
|
||||||
|
Auth: &Auth{Required: true, Capability: "sessions:read"},
|
||||||
|
},
|
||||||
|
},
|
||||||
|
},
|
||||||
|
{
|
||||||
|
Name: "episode",
|
||||||
|
Methods: []Method{
|
||||||
|
{
|
||||||
|
Verb: "GET",
|
||||||
|
UpstreamPath: "/v1/sessions/episodes/get",
|
||||||
|
Auth: &Auth{Required: true, Capability: "sessions:read"},
|
||||||
|
},
|
||||||
|
},
|
||||||
|
},
|
||||||
|
// sandbox — per-session k8s Pod+PVC+Service lifecycle (§4.15.4)
|
||||||
|
{
|
||||||
|
Name: "sandbox",
|
||||||
|
Methods: []Method{
|
||||||
|
{
|
||||||
|
Verb: "POST",
|
||||||
|
UpstreamPath: "/v1/sessions/sandbox",
|
||||||
|
Auth: &Auth{Required: true, Capability: "sessions:write"},
|
||||||
|
},
|
||||||
|
{
|
||||||
|
Verb: "DELETE",
|
||||||
|
UpstreamPath: "/v1/sessions/sandbox",
|
||||||
|
Auth: &Auth{Required: true, Capability: "sessions:write"},
|
||||||
|
},
|
||||||
|
},
|
||||||
|
},
|
||||||
|
},
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,97 @@
|
|||||||
|
// Package sessions — FleetClient calls the SandboxFleet controller REST API
|
||||||
|
// to assign sessions to fleet pods. This is the primary sandbox path;
|
||||||
|
// the raw-k8s SandboxManager is the fallback when the fleet controller
|
||||||
|
// is unavailable (e.g. during fleet rollout).
|
||||||
|
package sessions
|
||||||
|
|
||||||
|
import (
|
||||||
|
"bytes"
|
||||||
|
"context"
|
||||||
|
"encoding/json"
|
||||||
|
"fmt"
|
||||||
|
"io"
|
||||||
|
"net/http"
|
||||||
|
"os"
|
||||||
|
"time"
|
||||||
|
)
|
||||||
|
|
||||||
|
const (
|
||||||
|
defaultFleetControllerURL = "http://tool-definition-controller.poimen.svc.cluster.local:8788"
|
||||||
|
)
|
||||||
|
|
||||||
|
// FleetClient calls the SandboxFleet controller REST API.
|
||||||
|
type FleetClient struct {
|
||||||
|
baseURL string
|
||||||
|
client *http.Client
|
||||||
|
}
|
||||||
|
|
||||||
|
// FleetAssignResult is returned by the controller assign endpoint.
|
||||||
|
type FleetAssignResult struct {
|
||||||
|
PodSvcDns string `json:"podSvcDns"`
|
||||||
|
PodName string `json:"podName"`
|
||||||
|
}
|
||||||
|
|
||||||
|
// NewFleetClient creates a FleetClient from env or default URL.
|
||||||
|
// FLEET_CONTROLLER_URL — default: defaultFleetControllerURL
|
||||||
|
func NewFleetClient() *FleetClient {
|
||||||
|
url := os.Getenv("FLEET_CONTROLLER_URL")
|
||||||
|
if url == "" {
|
||||||
|
url = defaultFleetControllerURL
|
||||||
|
}
|
||||||
|
return &FleetClient{
|
||||||
|
baseURL: url,
|
||||||
|
client: &http.Client{Timeout: 35 * time.Second},
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// Assign asks the fleet controller to assign sessionID to a fleet pod.
|
||||||
|
// Returns the pod's ClusterIP service DNS and pod name.
|
||||||
|
func (f *FleetClient) Assign(ctx context.Context, sessionID string) (*FleetAssignResult, error) {
|
||||||
|
req, err := http.NewRequestWithContext(ctx, http.MethodPost,
|
||||||
|
f.baseURL+"/sessions/"+sessionID+"/assign", nil)
|
||||||
|
if err != nil {
|
||||||
|
return nil, fmt.Errorf("build request: %w", err)
|
||||||
|
}
|
||||||
|
req.Header.Set("Content-Type", "application/json")
|
||||||
|
|
||||||
|
resp, err := f.client.Do(req)
|
||||||
|
if err != nil {
|
||||||
|
return nil, fmt.Errorf("fleet controller unreachable: %w", err)
|
||||||
|
}
|
||||||
|
defer resp.Body.Close()
|
||||||
|
|
||||||
|
if resp.StatusCode >= 300 {
|
||||||
|
body, _ := io.ReadAll(io.LimitReader(resp.Body, 512))
|
||||||
|
return nil, fmt.Errorf("fleet controller returned %d: %s", resp.StatusCode, body)
|
||||||
|
}
|
||||||
|
|
||||||
|
var result FleetAssignResult
|
||||||
|
if err := json.NewDecoder(resp.Body).Decode(&result); err != nil {
|
||||||
|
return nil, fmt.Errorf("decode: %w", err)
|
||||||
|
}
|
||||||
|
if result.PodSvcDns == "" {
|
||||||
|
return nil, fmt.Errorf("fleet controller returned empty podSvcDns")
|
||||||
|
}
|
||||||
|
return &result, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// Release notifies the fleet controller that a session has ended.
|
||||||
|
// Cleans up the per-session ServiceAccount (audit trail SA).
|
||||||
|
// Non-fatal: fleet controller handles this async.
|
||||||
|
func (f *FleetClient) Release(ctx context.Context, sessionID string) error {
|
||||||
|
req, err := http.NewRequestWithContext(ctx, http.MethodDelete,
|
||||||
|
f.baseURL+"/sessions/"+sessionID, bytes.NewReader(nil))
|
||||||
|
if err != nil {
|
||||||
|
return fmt.Errorf("build request: %w", err)
|
||||||
|
}
|
||||||
|
resp, err := f.client.Do(req)
|
||||||
|
if err != nil {
|
||||||
|
return fmt.Errorf("fleet controller unreachable: %w", err)
|
||||||
|
}
|
||||||
|
defer resp.Body.Close()
|
||||||
|
if resp.StatusCode >= 300 {
|
||||||
|
body, _ := io.ReadAll(io.LimitReader(resp.Body, 512))
|
||||||
|
return fmt.Errorf("fleet controller returned %d: %s", resp.StatusCode, body)
|
||||||
|
}
|
||||||
|
return nil
|
||||||
|
}
|
||||||
@@ -0,0 +1,634 @@
|
|||||||
|
// Package sessions provides the internal handler for the sessions service adapter.
|
||||||
|
// It bridges the api-gw REST surface to poimen-graphql for all reads and writes:
|
||||||
|
//
|
||||||
|
// - Reads: hybridSearch / entity queries
|
||||||
|
// - Writes: upsertEntity mutation (direct, no LLM extraction)
|
||||||
|
//
|
||||||
|
// All IDs passed by the caller (session_id, episode_id) go as query params because
|
||||||
|
// the ServiceAdapter dispatcher does not substitute URL path parameters.
|
||||||
|
// episode_id is sha256(session_id + milestone_id + attempt) — callers compute it the
|
||||||
|
// same way; the handler accepts it on write and returns it on read.
|
||||||
|
package sessions
|
||||||
|
|
||||||
|
import (
|
||||||
|
"bytes"
|
||||||
|
"context"
|
||||||
|
"crypto/sha256"
|
||||||
|
"encoding/json"
|
||||||
|
"fmt"
|
||||||
|
"log"
|
||||||
|
"net/http"
|
||||||
|
"os"
|
||||||
|
"time"
|
||||||
|
|
||||||
|
"github.com/google/uuid"
|
||||||
|
|
||||||
|
"forgejo.riotpiao.com/rock/homelab-frontend/internal/tracing"
|
||||||
|
)
|
||||||
|
|
||||||
|
const (
|
||||||
|
defaultGraphQLURL = "http://poimen-graphql.poimen.svc.cluster.local:8090/graphql"
|
||||||
|
)
|
||||||
|
|
||||||
|
// Handler handles session REST endpoints via internal service calls.
|
||||||
|
// All traffic (reads and writes) goes to graphqlURL (poimen-graphql).
|
||||||
|
// Sandbox creates/deletes pi-session pods via the SandboxFleet controller
|
||||||
|
// (primary) with SandboxManager as direct-k8s fallback.
|
||||||
|
type Handler struct {
|
||||||
|
graphqlURL string
|
||||||
|
client *http.Client
|
||||||
|
sandbox *SandboxManager
|
||||||
|
fleet *FleetClient
|
||||||
|
}
|
||||||
|
|
||||||
|
// NewHandler creates a Handler reading upstream URL from env var.
|
||||||
|
// GRAPHQL_UPSTREAM_URL — default: defaultGraphQLURL
|
||||||
|
func NewHandler() *Handler {
|
||||||
|
gql := os.Getenv("GRAPHQL_UPSTREAM_URL")
|
||||||
|
if gql == "" {
|
||||||
|
gql = defaultGraphQLURL
|
||||||
|
}
|
||||||
|
h := &Handler{
|
||||||
|
graphqlURL: gql,
|
||||||
|
client: &http.Client{Timeout: 30 * time.Second, Transport: tracing.NewTransport(nil)},
|
||||||
|
sandbox: NewSandboxManager(),
|
||||||
|
fleet: NewFleetClient(),
|
||||||
|
}
|
||||||
|
h.startIdleReaper()
|
||||||
|
return h
|
||||||
|
}
|
||||||
|
|
||||||
|
// startIdleReaper launches a background goroutine that deletes exited pi-session
|
||||||
|
// pods (bridge.js idle-timeout triggered) while preserving their PVCs.
|
||||||
|
func (h *Handler) startIdleReaper() {
|
||||||
|
intervalStr := os.Getenv("SANDBOX_REAPER_INTERVAL_MINUTES")
|
||||||
|
interval := defaultReaperInterval
|
||||||
|
if intervalStr != "" {
|
||||||
|
var min int
|
||||||
|
if n, _ := fmt.Sscanf(intervalStr, "%d", &min); n == 1 && min > 0 {
|
||||||
|
interval = time.Duration(min) * time.Minute
|
||||||
|
}
|
||||||
|
}
|
||||||
|
go func() {
|
||||||
|
ticker := time.NewTicker(interval)
|
||||||
|
defer ticker.Stop()
|
||||||
|
for range ticker.C {
|
||||||
|
h.sandbox.ReapIdlePods(context.Background())
|
||||||
|
}
|
||||||
|
}()
|
||||||
|
}
|
||||||
|
|
||||||
|
// ServeHTTP routes by X-Upstream-Path set by the dispatcher.
|
||||||
|
func (h *Handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
|
||||||
|
path := r.Header.Get("X-Upstream-Path")
|
||||||
|
switch path {
|
||||||
|
case "/v1/sessions/create":
|
||||||
|
h.handleCreate(w, r)
|
||||||
|
case "/v1/sessions/list":
|
||||||
|
h.handleList(w, r)
|
||||||
|
case "/v1/sessions/get":
|
||||||
|
h.handleGet(w, r)
|
||||||
|
case "/v1/sessions/patch":
|
||||||
|
h.handlePatch(w, r)
|
||||||
|
case "/v1/sessions/milestones/create":
|
||||||
|
h.handleCreateMilestone(w, r)
|
||||||
|
case "/v1/sessions/episodes/list":
|
||||||
|
h.handleListEpisodes(w, r)
|
||||||
|
case "/v1/sessions/episodes/get":
|
||||||
|
h.handleGetEpisode(w, r)
|
||||||
|
case "/v1/sessions/sandbox":
|
||||||
|
h.handleSandbox(w, r)
|
||||||
|
default:
|
||||||
|
writeJSON(w, http.StatusNotFound, map[string]string{"error": "sessions endpoint not found"})
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// ── request / response types ──────────────────────────────────────────────────
|
||||||
|
|
||||||
|
type createSessionReq struct {
|
||||||
|
Name string `json:"name"`
|
||||||
|
Goal string `json:"goal,omitempty"`
|
||||||
|
Cwd string `json:"cwd,omitempty"`
|
||||||
|
}
|
||||||
|
|
||||||
|
type sessionResp struct {
|
||||||
|
SessionID string `json:"session_id"`
|
||||||
|
Name string `json:"name"`
|
||||||
|
Goal string `json:"goal,omitempty"`
|
||||||
|
Status string `json:"status"`
|
||||||
|
MemoryProject string `json:"memory_project"`
|
||||||
|
Cwd string `json:"cwd,omitempty"`
|
||||||
|
CreatedAt string `json:"created_at"`
|
||||||
|
}
|
||||||
|
|
||||||
|
type patchSessionReq struct {
|
||||||
|
Name *string `json:"name,omitempty"`
|
||||||
|
Goal *string `json:"goal,omitempty"`
|
||||||
|
Status *string `json:"status,omitempty"`
|
||||||
|
}
|
||||||
|
|
||||||
|
type createMilestoneReq struct {
|
||||||
|
Title string `json:"title"`
|
||||||
|
OrderIdx int `json:"order_idx"`
|
||||||
|
ParentID *string `json:"parent_id,omitempty"`
|
||||||
|
}
|
||||||
|
|
||||||
|
type milestoneResp struct {
|
||||||
|
MilestoneID string `json:"milestone_id"`
|
||||||
|
SessionID string `json:"session_id"`
|
||||||
|
Title string `json:"title"`
|
||||||
|
Status string `json:"status"`
|
||||||
|
OrderIdx int `json:"order_idx"`
|
||||||
|
ParentID *string `json:"parent_id,omitempty"`
|
||||||
|
CreatedAt string `json:"created_at"`
|
||||||
|
}
|
||||||
|
|
||||||
|
// ── handlers ──────────────────────────────────────────────────────────────────
|
||||||
|
|
||||||
|
// POST /v1/sessions (X-Resource: sessions)
|
||||||
|
func (h *Handler) handleCreate(w http.ResponseWriter, r *http.Request) {
|
||||||
|
var req createSessionReq
|
||||||
|
if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
|
||||||
|
writeJSON(w, http.StatusBadRequest, map[string]string{"error": "invalid JSON"})
|
||||||
|
return
|
||||||
|
}
|
||||||
|
if req.Name == "" {
|
||||||
|
writeJSON(w, http.StatusBadRequest, map[string]string{"error": "name is required"})
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
sessionID := uuid.New().String()
|
||||||
|
project := "session-" + sessionID
|
||||||
|
now := time.Now().UTC().Format(time.RFC3339)
|
||||||
|
|
||||||
|
summaryBytes, err := json.Marshal(map[string]any{
|
||||||
|
"id": sessionID,
|
||||||
|
"name": req.Name,
|
||||||
|
"goal": req.Goal,
|
||||||
|
"status": "active",
|
||||||
|
"cwd": req.Cwd,
|
||||||
|
"memory_project": project,
|
||||||
|
"created_at": now,
|
||||||
|
})
|
||||||
|
if err != nil {
|
||||||
|
writeJSON(w, http.StatusInternalServerError, map[string]string{"error": "marshal failed: " + err.Error()})
|
||||||
|
return
|
||||||
|
}
|
||||||
|
if _, err := h.upsertEntity(r, project, sessionID, "SESSION", string(summaryBytes)); err != nil {
|
||||||
|
writeJSON(w, http.StatusBadGateway, map[string]string{"error": "upsert failed: " + err.Error()})
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
writeJSON(w, http.StatusCreated, sessionResp{
|
||||||
|
SessionID: sessionID,
|
||||||
|
Name: req.Name,
|
||||||
|
Goal: req.Goal,
|
||||||
|
Status: "active",
|
||||||
|
MemoryProject: project,
|
||||||
|
Cwd: req.Cwd,
|
||||||
|
CreatedAt: now,
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
// GET /v1/sessions (X-Resource: sessions) ?status=active&cwd=/path
|
||||||
|
func (h *Handler) handleList(w http.ResponseWriter, r *http.Request) {
|
||||||
|
q := r.URL.Query()
|
||||||
|
status := q.Get("status")
|
||||||
|
cwd := q.Get("cwd")
|
||||||
|
limitStr := q.Get("limit")
|
||||||
|
limit := 50
|
||||||
|
if limitStr != "" {
|
||||||
|
fmt.Sscanf(limitStr, "%d", &limit) //nolint:errcheck
|
||||||
|
}
|
||||||
|
if limit > 200 {
|
||||||
|
limit = 200
|
||||||
|
}
|
||||||
|
|
||||||
|
searchQuery := "session"
|
||||||
|
if status != "" {
|
||||||
|
searchQuery = fmt.Sprintf("session status:%s", status)
|
||||||
|
}
|
||||||
|
if cwd != "" {
|
||||||
|
searchQuery = fmt.Sprintf("%s cwd:%s", searchQuery, cwd)
|
||||||
|
}
|
||||||
|
|
||||||
|
// entities query (not hybridSearch) — project arg is optional, so it queries
|
||||||
|
// across ALL projects. Required because each session lives in its own
|
||||||
|
// session-<uuid> project; there is no single project to hybridSearch against.
|
||||||
|
gqlQuery := `query($q:String!,$limit:Int!){
|
||||||
|
entities(query:$q,limit:$limit,entityType:SESSION){
|
||||||
|
entities{id name summary projectId tCreated}
|
||||||
|
searchMode
|
||||||
|
}
|
||||||
|
}`
|
||||||
|
vars := map[string]any{"q": searchQuery, "limit": limit}
|
||||||
|
result, err := h.gql(r, gqlQuery, vars)
|
||||||
|
if err != nil {
|
||||||
|
writeJSON(w, http.StatusBadGateway, map[string]string{"error": "graphql query failed: " + err.Error()})
|
||||||
|
return
|
||||||
|
}
|
||||||
|
writeJSON(w, http.StatusOK, result)
|
||||||
|
}
|
||||||
|
|
||||||
|
// GET /v1/sessions/:id (X-Resource: session) ?session_id=<uuid>
|
||||||
|
func (h *Handler) handleGet(w http.ResponseWriter, r *http.Request) {
|
||||||
|
sessionID := r.URL.Query().Get("session_id")
|
||||||
|
if sessionID == "" {
|
||||||
|
writeJSON(w, http.StatusBadRequest, map[string]string{"error": "session_id query param required"})
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
project := "session-" + sessionID
|
||||||
|
gqlQuery := `query($id:String!,$project:String!){
|
||||||
|
hybridSearch(query:$id,project:$project,limit:1,entityType:SESSION){
|
||||||
|
entities{id name summary tCreated}
|
||||||
|
}
|
||||||
|
}`
|
||||||
|
vars := map[string]any{"id": sessionID, "project": project}
|
||||||
|
result, err := h.gql(r, gqlQuery, vars)
|
||||||
|
if err != nil {
|
||||||
|
writeJSON(w, http.StatusBadGateway, map[string]string{"error": "graphql query failed: " + err.Error()})
|
||||||
|
return
|
||||||
|
}
|
||||||
|
writeJSON(w, http.StatusOK, result)
|
||||||
|
}
|
||||||
|
|
||||||
|
// PATCH /v1/sessions/:id (X-Resource: session) ?session_id=<uuid>
|
||||||
|
func (h *Handler) handlePatch(w http.ResponseWriter, r *http.Request) {
|
||||||
|
sessionID := r.URL.Query().Get("session_id")
|
||||||
|
if sessionID == "" {
|
||||||
|
writeJSON(w, http.StatusBadRequest, map[string]string{"error": "session_id query param required"})
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
var req patchSessionReq
|
||||||
|
if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
|
||||||
|
writeJSON(w, http.StatusBadRequest, map[string]string{"error": "invalid JSON"})
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
now := time.Now().UTC().Format(time.RFC3339)
|
||||||
|
project := "session-" + sessionID
|
||||||
|
|
||||||
|
// Read current summary so the upsert carries all existing fields.
|
||||||
|
current, err := h.fetchEntitySummary(r, project, sessionID, "SESSION")
|
||||||
|
if err != nil {
|
||||||
|
writeJSON(w, http.StatusBadGateway, map[string]string{"error": "fetch failed: " + err.Error()})
|
||||||
|
return
|
||||||
|
}
|
||||||
|
current["updated_at"] = now
|
||||||
|
if req.Name != nil {
|
||||||
|
current["name"] = *req.Name
|
||||||
|
}
|
||||||
|
if req.Goal != nil {
|
||||||
|
current["goal"] = *req.Goal
|
||||||
|
}
|
||||||
|
if req.Status != nil {
|
||||||
|
current["status"] = *req.Status
|
||||||
|
if *req.Status == "archived" {
|
||||||
|
current["archived_at"] = now
|
||||||
|
}
|
||||||
|
}
|
||||||
|
summaryBytes, err := json.Marshal(current)
|
||||||
|
if err != nil {
|
||||||
|
writeJSON(w, http.StatusInternalServerError, map[string]string{"error": "marshal failed: " + err.Error()})
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
if _, err := h.upsertEntity(r, project, sessionID, "SESSION", string(summaryBytes)); err != nil {
|
||||||
|
writeJSON(w, http.StatusBadGateway, map[string]string{"error": "upsert failed: " + err.Error()})
|
||||||
|
return
|
||||||
|
}
|
||||||
|
writeJSON(w, http.StatusOK, map[string]string{"session_id": sessionID, "status": "updated"})
|
||||||
|
}
|
||||||
|
|
||||||
|
// POST /v1/sessions/:id/milestones (X-Resource: milestones) ?session_id=<uuid>
|
||||||
|
func (h *Handler) handleCreateMilestone(w http.ResponseWriter, r *http.Request) {
|
||||||
|
sessionID := r.URL.Query().Get("session_id")
|
||||||
|
if sessionID == "" {
|
||||||
|
writeJSON(w, http.StatusBadRequest, map[string]string{"error": "session_id query param required"})
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
var req createMilestoneReq
|
||||||
|
if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
|
||||||
|
writeJSON(w, http.StatusBadRequest, map[string]string{"error": "invalid JSON"})
|
||||||
|
return
|
||||||
|
}
|
||||||
|
if req.Title == "" {
|
||||||
|
writeJSON(w, http.StatusBadRequest, map[string]string{"error": "title is required"})
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
milestoneID := uuid.New().String()
|
||||||
|
now := time.Now().UTC().Format(time.RFC3339)
|
||||||
|
project := "session-" + sessionID
|
||||||
|
|
||||||
|
msContent := map[string]any{
|
||||||
|
"type": "MILESTONE",
|
||||||
|
"id": milestoneID,
|
||||||
|
"session_id": sessionID,
|
||||||
|
"title": req.Title,
|
||||||
|
"status": "pending",
|
||||||
|
"order_idx": req.OrderIdx,
|
||||||
|
"created_at": now,
|
||||||
|
}
|
||||||
|
if req.ParentID != nil {
|
||||||
|
msContent["parent_id"] = *req.ParentID
|
||||||
|
}
|
||||||
|
summaryBytes, err := json.Marshal(msContent)
|
||||||
|
if err != nil {
|
||||||
|
writeJSON(w, http.StatusInternalServerError, map[string]string{"error": "marshal failed: " + err.Error()})
|
||||||
|
return
|
||||||
|
}
|
||||||
|
if _, err := h.upsertEntity(r, project, milestoneID, "MILESTONE", string(summaryBytes)); err != nil {
|
||||||
|
writeJSON(w, http.StatusBadGateway, map[string]string{"error": "upsert failed: " + err.Error()})
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
writeJSON(w, http.StatusCreated, milestoneResp{
|
||||||
|
MilestoneID: milestoneID,
|
||||||
|
SessionID: sessionID,
|
||||||
|
Title: req.Title,
|
||||||
|
Status: "pending",
|
||||||
|
OrderIdx: req.OrderIdx,
|
||||||
|
ParentID: req.ParentID,
|
||||||
|
CreatedAt: now,
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
// GET /v1/sessions/:id/episodes (X-Resource: episodes) ?session_id=<uuid>
|
||||||
|
func (h *Handler) handleListEpisodes(w http.ResponseWriter, r *http.Request) {
|
||||||
|
sessionID := r.URL.Query().Get("session_id")
|
||||||
|
if sessionID == "" {
|
||||||
|
writeJSON(w, http.StatusBadRequest, map[string]string{"error": "session_id query param required"})
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
limitStr := r.URL.Query().Get("limit")
|
||||||
|
limit := 50
|
||||||
|
if limitStr != "" {
|
||||||
|
fmt.Sscanf(limitStr, "%d", &limit) //nolint:errcheck
|
||||||
|
}
|
||||||
|
if limit > 200 {
|
||||||
|
limit = 200
|
||||||
|
}
|
||||||
|
|
||||||
|
project := "session-" + sessionID
|
||||||
|
gqlQuery := `query($project:String!,$limit:Int!){
|
||||||
|
hybridSearch(query:"episode",project:$project,limit:$limit,entityType:EPISODE){
|
||||||
|
entities{id name summary tCreated}
|
||||||
|
searchMode
|
||||||
|
}
|
||||||
|
}`
|
||||||
|
vars := map[string]any{"project": project, "limit": limit}
|
||||||
|
result, err := h.gql(r, gqlQuery, vars)
|
||||||
|
if err != nil {
|
||||||
|
writeJSON(w, http.StatusBadGateway, map[string]string{"error": "graphql query failed: " + err.Error()})
|
||||||
|
return
|
||||||
|
}
|
||||||
|
writeJSON(w, http.StatusOK, result)
|
||||||
|
}
|
||||||
|
|
||||||
|
// GET /v1/sessions/:id/episodes/:eid (X-Resource: episode)
|
||||||
|
// ?session_id=<uuid>&episode_id=<sha256>
|
||||||
|
func (h *Handler) handleGetEpisode(w http.ResponseWriter, r *http.Request) {
|
||||||
|
q := r.URL.Query()
|
||||||
|
sessionID := q.Get("session_id")
|
||||||
|
episodeID := q.Get("episode_id")
|
||||||
|
if sessionID == "" || episodeID == "" {
|
||||||
|
writeJSON(w, http.StatusBadRequest, map[string]string{
|
||||||
|
"error": "session_id and episode_id query params required",
|
||||||
|
})
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
gqlQuery := `query($id:ID!){
|
||||||
|
entity(id:$id){
|
||||||
|
id name entityType summary tCreated
|
||||||
|
edges{
|
||||||
|
relationType fact confidence
|
||||||
|
target{name entityType}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}`
|
||||||
|
vars := map[string]any{"id": episodeID}
|
||||||
|
result, err := h.gql(r, gqlQuery, vars)
|
||||||
|
if err != nil {
|
||||||
|
writeJSON(w, http.StatusBadGateway, map[string]string{"error": "graphql query failed: " + err.Error()})
|
||||||
|
return
|
||||||
|
}
|
||||||
|
writeJSON(w, http.StatusOK, result)
|
||||||
|
}
|
||||||
|
|
||||||
|
// ── helpers ───────────────────────────────────────────────────────────────────
|
||||||
|
|
||||||
|
// upsertEntity runs a GraphQL upsertEntity mutation on poimen-graphql.
|
||||||
|
// Writes are direct (no LLM extraction). Returns the entity id assigned by the server.
|
||||||
|
// The name field is used as the stable upsert key within (project, entityType).
|
||||||
|
func (h *Handler) upsertEntity(r *http.Request, project, name, entityType, summary string) (string, error) {
|
||||||
|
mutation := `mutation($input:UpsertEntityInput!){
|
||||||
|
upsertEntity(input:$input){ id created }
|
||||||
|
}`
|
||||||
|
vars := map[string]any{
|
||||||
|
"input": map[string]any{
|
||||||
|
"project": project,
|
||||||
|
"name": name,
|
||||||
|
"entityType": entityType,
|
||||||
|
"summary": summary,
|
||||||
|
},
|
||||||
|
}
|
||||||
|
result, err := h.gql(r, mutation, vars)
|
||||||
|
if err != nil {
|
||||||
|
return "", err
|
||||||
|
}
|
||||||
|
data, _ := result.(map[string]any)
|
||||||
|
d, _ := data["data"].(map[string]any)
|
||||||
|
entity, _ := d["upsertEntity"].(map[string]any)
|
||||||
|
id, _ := entity["id"].(string)
|
||||||
|
return id, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// fetchEntitySummary fetches the current summary blob of an entity by (project, name, entityType).
|
||||||
|
// Returns a parsed map so callers can merge patch fields before re-upserting.
|
||||||
|
// Returns empty map (not error) if no entity found.
|
||||||
|
func (h *Handler) fetchEntitySummary(r *http.Request, project, name, entityType string) (map[string]any, error) {
|
||||||
|
gqlQuery := `query($q:String!,$project:String!,$entityType:GqlEntityType!,$limit:Int!){
|
||||||
|
hybridSearch(query:$q,project:$project,limit:$limit,entityType:$entityType){
|
||||||
|
entities{id name summary tCreated}
|
||||||
|
}
|
||||||
|
}`
|
||||||
|
vars := map[string]any{
|
||||||
|
"q": name,
|
||||||
|
"project": project,
|
||||||
|
"entityType": entityType,
|
||||||
|
"limit": 1,
|
||||||
|
}
|
||||||
|
result, err := h.gql(r, gqlQuery, vars)
|
||||||
|
if err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
data, _ := result.(map[string]any)
|
||||||
|
d, _ := data["data"].(map[string]any)
|
||||||
|
search, _ := d["hybridSearch"].(map[string]any)
|
||||||
|
entities, _ := search["entities"].([]any)
|
||||||
|
if len(entities) == 0 {
|
||||||
|
return map[string]any{}, nil
|
||||||
|
}
|
||||||
|
entity, _ := entities[0].(map[string]any)
|
||||||
|
summaryStr, _ := entity["summary"].(string)
|
||||||
|
var summary map[string]any
|
||||||
|
if summaryStr != "" {
|
||||||
|
if err := json.Unmarshal([]byte(summaryStr), &summary); err != nil {
|
||||||
|
return map[string]any{}, nil // corrupt summary — start fresh
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if summary == nil {
|
||||||
|
summary = map[string]any{}
|
||||||
|
}
|
||||||
|
return summary, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// gql executes a GraphQL query against poimen-graphql, forwarding the caller's JWT.
|
||||||
|
func (h *Handler) gql(r *http.Request, query string, variables map[string]any) (any, error) {
|
||||||
|
payload, err := json.Marshal(map[string]any{
|
||||||
|
"query": query,
|
||||||
|
"variables": variables,
|
||||||
|
})
|
||||||
|
if err != nil {
|
||||||
|
return nil, fmt.Errorf("marshal: %w", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
req, err := http.NewRequestWithContext(r.Context(), http.MethodPost,
|
||||||
|
h.graphqlURL, bytes.NewReader(payload))
|
||||||
|
if err != nil {
|
||||||
|
return nil, fmt.Errorf("build request: %w", err)
|
||||||
|
}
|
||||||
|
req.Header.Set("Content-Type", "application/json")
|
||||||
|
if auth := r.Header.Get("Authorization"); auth != "" {
|
||||||
|
req.Header.Set("Authorization", auth)
|
||||||
|
}
|
||||||
|
|
||||||
|
resp, err := h.client.Do(req)
|
||||||
|
if err != nil {
|
||||||
|
return nil, fmt.Errorf("http: %w", err)
|
||||||
|
}
|
||||||
|
defer resp.Body.Close()
|
||||||
|
|
||||||
|
if resp.StatusCode >= 300 {
|
||||||
|
return nil, fmt.Errorf("graphql returned %d", resp.StatusCode)
|
||||||
|
}
|
||||||
|
|
||||||
|
var result map[string]any
|
||||||
|
if err := json.NewDecoder(resp.Body).Decode(&result); err != nil {
|
||||||
|
return nil, fmt.Errorf("decode: %w", err)
|
||||||
|
}
|
||||||
|
if errs, ok := result["errors"]; ok && errs != nil {
|
||||||
|
b, _ := json.Marshal(errs)
|
||||||
|
return nil, fmt.Errorf("graphql errors: %s", b)
|
||||||
|
}
|
||||||
|
return result, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// handleSandbox dispatches POST (create/resume) and DELETE (teardown).
|
||||||
|
// Both take ?session_id= as a query param.
|
||||||
|
func (h *Handler) handleSandbox(w http.ResponseWriter, r *http.Request) {
|
||||||
|
sessionID := r.URL.Query().Get("session_id")
|
||||||
|
if sessionID == "" {
|
||||||
|
writeJSON(w, http.StatusBadRequest, map[string]string{"error": "session_id query param required"})
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
switch r.Method {
|
||||||
|
case http.MethodPost:
|
||||||
|
h.handleSandboxCreate(w, r, sessionID)
|
||||||
|
case http.MethodDelete:
|
||||||
|
h.handleSandboxDelete(w, r, sessionID)
|
||||||
|
default:
|
||||||
|
writeJSON(w, http.StatusMethodNotAllowed, map[string]string{"error": "method not allowed"})
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// POST ?session_id= — assign session to SandboxFleet pod (primary) or
|
||||||
|
// create a dedicated pod via SandboxManager (fallback).
|
||||||
|
func (h *Handler) handleSandboxCreate(w http.ResponseWriter, r *http.Request, sessionID string) {
|
||||||
|
// Primary: fleet controller assigns session to a pre-warmed pod.
|
||||||
|
if assign, err := h.fleet.Assign(r.Context(), sessionID); err == nil {
|
||||||
|
svcDns := assign.PodSvcDns + fmt.Sprintf(":%d", sandboxPort)
|
||||||
|
log.Printf("sandbox: fleet assigned session=%s pod=%s dns=%s", sessionID, assign.PodName, svcDns)
|
||||||
|
go h.updateSandboxSummary(r, sessionID, assign.PodName, "", svcDns)
|
||||||
|
writeJSON(w, http.StatusOK, map[string]any{
|
||||||
|
"session_id": sessionID,
|
||||||
|
"pod_name": assign.PodName,
|
||||||
|
"svc_dns": svcDns,
|
||||||
|
"resumed": false,
|
||||||
|
"fleet": true,
|
||||||
|
})
|
||||||
|
return
|
||||||
|
} else {
|
||||||
|
log.Printf("sandbox: fleet unavailable for session=%s (%v), falling back to direct k8s", sessionID, err)
|
||||||
|
}
|
||||||
|
|
||||||
|
// Fallback: create a dedicated pod via SandboxManager.
|
||||||
|
result, err := h.sandbox.CreateOrResume(r.Context(), sessionID)
|
||||||
|
if err != nil {
|
||||||
|
log.Printf("sandbox create/resume session=%s: %v", sessionID, err)
|
||||||
|
writeJSON(w, http.StatusBadGateway, map[string]string{"error": "sandbox error: " + err.Error()})
|
||||||
|
return
|
||||||
|
}
|
||||||
|
go h.updateSandboxSummary(r, sessionID, result.PodName, result.PVCName, result.SvcDNS)
|
||||||
|
log.Printf("sandbox: direct pod session=%s pod=%s resumed=%v dns=%s", sessionID, result.PodName, result.Resumed, result.SvcDNS)
|
||||||
|
writeJSON(w, http.StatusOK, result)
|
||||||
|
}
|
||||||
|
|
||||||
|
// updateSandboxSummary merges pod/svc metadata into the session memory_entity (fire-and-forget).
|
||||||
|
// Uses context.Background() because r.Context() is cancelled once the HTTP handler returns.
|
||||||
|
func (h *Handler) updateSandboxSummary(r *http.Request, sessionID, podName, pvcName, svcDns string) {
|
||||||
|
// Build a detached request carrying just the auth header.
|
||||||
|
detached, _ := http.NewRequestWithContext(context.Background(), r.Method, r.URL.String(), nil)
|
||||||
|
if auth := r.Header.Get("Authorization"); auth != "" {
|
||||||
|
detached.Header.Set("Authorization", auth)
|
||||||
|
}
|
||||||
|
current, _ := h.fetchEntitySummary(detached, "session-"+sessionID, sessionID, "SESSION")
|
||||||
|
if current == nil {
|
||||||
|
current = map[string]any{}
|
||||||
|
}
|
||||||
|
current["pod_name"] = podName
|
||||||
|
current["svc_dns"] = svcDns
|
||||||
|
current["last_active_at"] = time.Now().UTC().Format(time.RFC3339)
|
||||||
|
if pvcName != "" {
|
||||||
|
current["pvc_name"] = pvcName
|
||||||
|
}
|
||||||
|
if b, err := json.Marshal(current); err == nil {
|
||||||
|
h.upsertEntity(r, "session-"+sessionID, sessionID, "SESSION", string(b)) //nolint:errcheck
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// DELETE ?session_id= — release from fleet + delete dedicated pod+svc+pvc if any.
|
||||||
|
func (h *Handler) handleSandboxDelete(w http.ResponseWriter, r *http.Request, sessionID string) {
|
||||||
|
// Release from fleet (cleans up per-session SA; non-fatal if session was fleet-based).
|
||||||
|
if err := h.fleet.Release(r.Context(), sessionID); err != nil {
|
||||||
|
log.Printf("sandbox: fleet release session=%s: %v (continuing)", sessionID, err)
|
||||||
|
}
|
||||||
|
// Also attempt direct k8s cleanup (no-ops if resources don't exist).
|
||||||
|
if err := h.sandbox.Delete(r.Context(), sessionID); err != nil {
|
||||||
|
log.Printf("sandbox delete session=%s: %v", sessionID, err)
|
||||||
|
}
|
||||||
|
log.Printf("sandbox: released session=%s", sessionID)
|
||||||
|
writeJSON(w, http.StatusOK, map[string]string{"session_id": sessionID, "status": "deleted"})
|
||||||
|
}
|
||||||
|
|
||||||
|
// EpisodeID computes the deterministic episode ID per §4.14 Q6:
|
||||||
|
// sha256(session_id + milestone_id + attempt).
|
||||||
|
// Exported so callers can verify the value the server would produce.
|
||||||
|
func EpisodeID(sessionID, milestoneID string, attempt int) string {
|
||||||
|
h := sha256.New()
|
||||||
|
h.Write([]byte(fmt.Sprintf("%s%s%d", sessionID, milestoneID, attempt)))
|
||||||
|
return fmt.Sprintf("%x", h.Sum(nil))
|
||||||
|
}
|
||||||
|
|
||||||
|
func writeJSON(w http.ResponseWriter, status int, data any) {
|
||||||
|
w.Header().Set("Content-Type", "application/json")
|
||||||
|
w.WriteHeader(status)
|
||||||
|
json.NewEncoder(w).Encode(data) //nolint:errcheck
|
||||||
|
}
|
||||||
@@ -0,0 +1,594 @@
|
|||||||
|
package sessions
|
||||||
|
|
||||||
|
import (
|
||||||
|
"crypto/sha256"
|
||||||
|
"encoding/json"
|
||||||
|
"fmt"
|
||||||
|
"io"
|
||||||
|
"net/http"
|
||||||
|
"net/http/httptest"
|
||||||
|
"strings"
|
||||||
|
"testing"
|
||||||
|
"time"
|
||||||
|
)
|
||||||
|
|
||||||
|
// ── stub servers ──────────────────────────────────────────────────────────────
|
||||||
|
|
||||||
|
// newGQLServer returns a test server that routes by inspecting the request body's
|
||||||
|
// "query" field and records captured request bodies.
|
||||||
|
//
|
||||||
|
// - "upsertEntity" → mutation response
|
||||||
|
// - "hybridSearch" → search response (get, patch, episodes handlers)
|
||||||
|
// - "entities" → list response (sessions list handler)
|
||||||
|
// - "entity" → single entity response
|
||||||
|
func newGQLServer(t *testing.T, captured *[]byte) *httptest.Server {
|
||||||
|
t.Helper()
|
||||||
|
return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||||
|
body, _ := io.ReadAll(r.Body)
|
||||||
|
if captured != nil {
|
||||||
|
*captured = body
|
||||||
|
}
|
||||||
|
w.Header().Set("Content-Type", "application/json")
|
||||||
|
bs := string(body)
|
||||||
|
switch {
|
||||||
|
case strings.Contains(bs, "upsertEntity"):
|
||||||
|
fmt.Fprint(w, `{"data":{"upsertEntity":{"id":"test-entity-id","created":true}}}`)
|
||||||
|
case strings.Contains(bs, "hybridSearch"):
|
||||||
|
// get, patch, episodes handlers still use hybridSearch
|
||||||
|
fmt.Fprint(w, `{"data":{"hybridSearch":{"entities":[]}}}`)
|
||||||
|
case strings.Contains(bs, "entities"):
|
||||||
|
// sessions list handler uses entities query
|
||||||
|
fmt.Fprint(w, `{"data":{"entities":{"entities":[]}}}`)
|
||||||
|
case strings.Contains(bs, "entity"):
|
||||||
|
fmt.Fprint(w, `{"data":{"entity":null}}`)
|
||||||
|
default:
|
||||||
|
fmt.Fprint(w, `{"data":{}}`)
|
||||||
|
}
|
||||||
|
}))
|
||||||
|
}
|
||||||
|
|
||||||
|
// buildHandler wires a Handler with the supplied graphql URL.
|
||||||
|
func buildHandler(gqlURL string) *Handler {
|
||||||
|
return &Handler{
|
||||||
|
graphqlURL: gqlURL,
|
||||||
|
client: &http.Client{Timeout: 5 * time.Second},
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// do fires a request through the handler ServeHTTP. path sets X-Upstream-Path.
|
||||||
|
func do(h *Handler, method, path, upstreamPath, body, query string) *httptest.ResponseRecorder {
|
||||||
|
var bodyReader io.Reader
|
||||||
|
if body != "" {
|
||||||
|
bodyReader = strings.NewReader(body)
|
||||||
|
} else {
|
||||||
|
bodyReader = http.NoBody
|
||||||
|
}
|
||||||
|
url := "/"
|
||||||
|
if query != "" {
|
||||||
|
url = "/?" + query
|
||||||
|
}
|
||||||
|
req := httptest.NewRequest(method, url, bodyReader)
|
||||||
|
req.Header.Set("X-Upstream-Path", upstreamPath)
|
||||||
|
if body != "" {
|
||||||
|
req.Header.Set("Content-Type", "application/json")
|
||||||
|
}
|
||||||
|
_ = path // kept for readability in callers
|
||||||
|
rr := httptest.NewRecorder()
|
||||||
|
h.ServeHTTP(rr, req)
|
||||||
|
return rr
|
||||||
|
}
|
||||||
|
|
||||||
|
// jsonBody decodes response body into map.
|
||||||
|
func jsonBody(t *testing.T, rr *httptest.ResponseRecorder) map[string]any {
|
||||||
|
t.Helper()
|
||||||
|
var out map[string]any
|
||||||
|
if err := json.Unmarshal(rr.Body.Bytes(), &out); err != nil {
|
||||||
|
t.Fatalf("decode body: %v — raw: %s", err, rr.Body.String())
|
||||||
|
}
|
||||||
|
return out
|
||||||
|
}
|
||||||
|
|
||||||
|
// ── 1. handleCreate ───────────────────────────────────────────────────────────
|
||||||
|
|
||||||
|
func TestHandleCreate(t *testing.T) {
|
||||||
|
t.Run("missing name → 400", func(t *testing.T) {
|
||||||
|
gql := newGQLServer(t, nil)
|
||||||
|
defer gql.Close()
|
||||||
|
h := buildHandler(gql.URL)
|
||||||
|
|
||||||
|
rr := do(h, http.MethodPost, "/v1/sessions/create", "/v1/sessions/create",
|
||||||
|
`{"goal":"some goal"}`, "")
|
||||||
|
if rr.Code != http.StatusBadRequest {
|
||||||
|
t.Fatalf("want 400, got %d", rr.Code)
|
||||||
|
}
|
||||||
|
})
|
||||||
|
|
||||||
|
t.Run("bad JSON → 400", func(t *testing.T) {
|
||||||
|
gql := newGQLServer(t, nil)
|
||||||
|
defer gql.Close()
|
||||||
|
h := buildHandler(gql.URL)
|
||||||
|
|
||||||
|
rr := do(h, http.MethodPost, "/v1/sessions/create", "/v1/sessions/create",
|
||||||
|
`not json`, "")
|
||||||
|
if rr.Code != http.StatusBadRequest {
|
||||||
|
t.Fatalf("want 400, got %d", rr.Code)
|
||||||
|
}
|
||||||
|
})
|
||||||
|
|
||||||
|
t.Run("upsert graphql error → 502", func(t *testing.T) {
|
||||||
|
gqlErr := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||||
|
w.Header().Set("Content-Type", "application/json")
|
||||||
|
w.WriteHeader(200)
|
||||||
|
fmt.Fprint(w, `{"errors":[{"message":"resolver error"}]}`)
|
||||||
|
}))
|
||||||
|
defer gqlErr.Close()
|
||||||
|
h := buildHandler(gqlErr.URL)
|
||||||
|
|
||||||
|
rr := do(h, http.MethodPost, "/v1/sessions/create", "/v1/sessions/create",
|
||||||
|
`{"name":"fail-session"}`, "")
|
||||||
|
if rr.Code != http.StatusBadGateway {
|
||||||
|
t.Fatalf("want 502, got %d", rr.Code)
|
||||||
|
}
|
||||||
|
if !strings.Contains(rr.Body.String(), "upsert failed:") {
|
||||||
|
t.Fatalf("want 'upsert failed:' in body, got: %s", rr.Body.String())
|
||||||
|
}
|
||||||
|
})
|
||||||
|
|
||||||
|
t.Run("upsert 200 → 201 with session_id, status=active, memory_project", func(t *testing.T) {
|
||||||
|
gql := newGQLServer(t, nil)
|
||||||
|
defer gql.Close()
|
||||||
|
h := buildHandler(gql.URL)
|
||||||
|
|
||||||
|
rr := do(h, http.MethodPost, "/v1/sessions/create", "/v1/sessions/create",
|
||||||
|
`{"name":"my-session","goal":"finish task"}`, "")
|
||||||
|
if rr.Code != http.StatusCreated {
|
||||||
|
t.Fatalf("want 201, got %d — body: %s", rr.Code, rr.Body.String())
|
||||||
|
}
|
||||||
|
body := jsonBody(t, rr)
|
||||||
|
if _, ok := body["session_id"]; !ok {
|
||||||
|
t.Fatal("missing session_id")
|
||||||
|
}
|
||||||
|
if body["status"] != "active" {
|
||||||
|
t.Fatalf("want status=active, got %v", body["status"])
|
||||||
|
}
|
||||||
|
proj, _ := body["memory_project"].(string)
|
||||||
|
if !strings.HasPrefix(proj, "session-") {
|
||||||
|
t.Fatalf("memory_project must start with 'session-', got %s", proj)
|
||||||
|
}
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
// ── 2. handleList ─────────────────────────────────────────────────────────────
|
||||||
|
|
||||||
|
func TestHandleList(t *testing.T) {
|
||||||
|
t.Run("no params → 200, gql called with query containing 'session'", func(t *testing.T) {
|
||||||
|
var captured []byte
|
||||||
|
gql := newGQLServer(t, &captured)
|
||||||
|
defer gql.Close()
|
||||||
|
h := buildHandler(gql.URL)
|
||||||
|
|
||||||
|
rr := do(h, http.MethodGet, "/v1/sessions/list", "/v1/sessions/list", "", "")
|
||||||
|
if rr.Code != http.StatusOK {
|
||||||
|
t.Fatalf("want 200, got %d", rr.Code)
|
||||||
|
}
|
||||||
|
if !strings.Contains(string(captured), "session") {
|
||||||
|
t.Fatalf("gql payload missing 'session': %s", captured)
|
||||||
|
}
|
||||||
|
})
|
||||||
|
|
||||||
|
t.Run("status=paused → gql query contains 'status:paused'", func(t *testing.T) {
|
||||||
|
var captured []byte
|
||||||
|
gql := newGQLServer(t, &captured)
|
||||||
|
defer gql.Close()
|
||||||
|
h := buildHandler(gql.URL)
|
||||||
|
|
||||||
|
rr := do(h, http.MethodGet, "/v1/sessions/list", "/v1/sessions/list", "", "status=paused")
|
||||||
|
if rr.Code != http.StatusOK {
|
||||||
|
t.Fatalf("want 200, got %d", rr.Code)
|
||||||
|
}
|
||||||
|
if !strings.Contains(string(captured), "status:paused") {
|
||||||
|
t.Fatalf("gql payload missing 'status:paused': %s", captured)
|
||||||
|
}
|
||||||
|
})
|
||||||
|
|
||||||
|
t.Run("graphql errors body → 502", func(t *testing.T) {
|
||||||
|
gql2 := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||||
|
w.WriteHeader(500)
|
||||||
|
fmt.Fprint(w, `{"errors":[{"message":"resolver error"}]}`)
|
||||||
|
}))
|
||||||
|
defer gql2.Close()
|
||||||
|
h2 := buildHandler(gql2.URL)
|
||||||
|
|
||||||
|
rr := do(h2, http.MethodGet, "/v1/sessions/list", "/v1/sessions/list", "", "")
|
||||||
|
if rr.Code != http.StatusBadGateway {
|
||||||
|
t.Fatalf("want 502, got %d — body: %s", rr.Code, rr.Body.String())
|
||||||
|
}
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
// ── 3. handleGet ─────────────────────────────────────────────────────────────
|
||||||
|
|
||||||
|
func TestHandleGet(t *testing.T) {
|
||||||
|
t.Run("missing session_id → 400", func(t *testing.T) {
|
||||||
|
gql := newGQLServer(t, nil)
|
||||||
|
defer gql.Close()
|
||||||
|
h := buildHandler(gql.URL)
|
||||||
|
|
||||||
|
rr := do(h, http.MethodGet, "/v1/sessions/get", "/v1/sessions/get", "", "")
|
||||||
|
if rr.Code != http.StatusBadRequest {
|
||||||
|
t.Fatalf("want 400, got %d", rr.Code)
|
||||||
|
}
|
||||||
|
})
|
||||||
|
|
||||||
|
t.Run("success → 200, gql called with session-scoped project", func(t *testing.T) {
|
||||||
|
var captured []byte
|
||||||
|
gql := newGQLServer(t, &captured)
|
||||||
|
defer gql.Close()
|
||||||
|
h := buildHandler(gql.URL)
|
||||||
|
|
||||||
|
sid := "abc-123"
|
||||||
|
rr := do(h, http.MethodGet, "/v1/sessions/get", "/v1/sessions/get", "", "session_id="+sid)
|
||||||
|
if rr.Code != http.StatusOK {
|
||||||
|
t.Fatalf("want 200, got %d — body: %s", rr.Code, rr.Body.String())
|
||||||
|
}
|
||||||
|
if !strings.Contains(string(captured), "session-"+sid) {
|
||||||
|
t.Fatalf("gql payload missing scoped project 'session-%s': %s", sid, captured)
|
||||||
|
}
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
// ── 4. handlePatch ────────────────────────────────────────────────────────────
|
||||||
|
|
||||||
|
func TestHandlePatch(t *testing.T) {
|
||||||
|
t.Run("missing session_id → 400", func(t *testing.T) {
|
||||||
|
gql := newGQLServer(t, nil)
|
||||||
|
defer gql.Close()
|
||||||
|
h := buildHandler(gql.URL)
|
||||||
|
|
||||||
|
rr := do(h, http.MethodPatch, "/v1/sessions/patch", "/v1/sessions/patch",
|
||||||
|
`{"status":"paused"}`, "")
|
||||||
|
if rr.Code != http.StatusBadRequest {
|
||||||
|
t.Fatalf("want 400, got %d", rr.Code)
|
||||||
|
}
|
||||||
|
})
|
||||||
|
|
||||||
|
t.Run("bad JSON → 400", func(t *testing.T) {
|
||||||
|
gql := newGQLServer(t, nil)
|
||||||
|
defer gql.Close()
|
||||||
|
h := buildHandler(gql.URL)
|
||||||
|
|
||||||
|
rr := do(h, http.MethodPatch, "/v1/sessions/patch", "/v1/sessions/patch",
|
||||||
|
`not json`, "session_id=abc-123")
|
||||||
|
if rr.Code != http.StatusBadRequest {
|
||||||
|
t.Fatalf("want 400, got %d", rr.Code)
|
||||||
|
}
|
||||||
|
})
|
||||||
|
|
||||||
|
t.Run("upsert success → 200 with session_id, status=updated", func(t *testing.T) {
|
||||||
|
var secondCall []byte
|
||||||
|
callCount := 0
|
||||||
|
gqlServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
|
||||||
|
body, _ := io.ReadAll(req.Body)
|
||||||
|
bodyStr := string(body)
|
||||||
|
w.Header().Set("Content-Type", "application/json")
|
||||||
|
callCount++
|
||||||
|
if strings.Contains(bodyStr, "hybridSearch") {
|
||||||
|
summary := `{"id":"sess-1","name":"abc-123","status":"active","goal":"old-goal"}`
|
||||||
|
fmt.Fprintf(w, `{"data":{"hybridSearch":{"entities":[{"id":"eid","name":"abc-123","summary":%q,"tCreated":"2026-01-01T00:00:00Z"}]}}}`, summary)
|
||||||
|
} else {
|
||||||
|
secondCall = body
|
||||||
|
fmt.Fprint(w, `{"data":{"upsertEntity":{"id":"eid","created":false}}}`)
|
||||||
|
}
|
||||||
|
}))
|
||||||
|
defer gqlServer.Close()
|
||||||
|
h := buildHandler(gqlServer.URL)
|
||||||
|
|
||||||
|
sid := "abc-123"
|
||||||
|
rr := do(h, http.MethodPatch, "/v1/sessions/patch", "/v1/sessions/patch",
|
||||||
|
`{"status":"paused"}`, "session_id="+sid)
|
||||||
|
if rr.Code != http.StatusOK {
|
||||||
|
t.Fatalf("want 200, got %d — body: %s", rr.Code, rr.Body.String())
|
||||||
|
}
|
||||||
|
body := jsonBody(t, rr)
|
||||||
|
if body["session_id"] != sid {
|
||||||
|
t.Fatalf("want session_id=%s, got %v", sid, body["session_id"])
|
||||||
|
}
|
||||||
|
if body["status"] != "updated" {
|
||||||
|
t.Fatalf("want status=updated, got %v", body["status"])
|
||||||
|
}
|
||||||
|
if callCount != 2 {
|
||||||
|
t.Fatalf("want 2 gql calls (hybridSearch+upsert), got %d", callCount)
|
||||||
|
}
|
||||||
|
if !strings.Contains(string(secondCall), "paused") {
|
||||||
|
t.Fatalf("upsert payload missing merged status: %s", secondCall)
|
||||||
|
}
|
||||||
|
})
|
||||||
|
|
||||||
|
t.Run("upsert error → 502", func(t *testing.T) {
|
||||||
|
gqlServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
|
||||||
|
body, _ := io.ReadAll(req.Body)
|
||||||
|
w.Header().Set("Content-Type", "application/json")
|
||||||
|
if strings.Contains(string(body), "hybridSearch") {
|
||||||
|
fmt.Fprint(w, `{"data":{"hybridSearch":{"entities":[]}}}`)
|
||||||
|
} else {
|
||||||
|
fmt.Fprint(w, `{"errors":[{"message":"write error"}]}`)
|
||||||
|
}
|
||||||
|
}))
|
||||||
|
defer gqlServer.Close()
|
||||||
|
h := buildHandler(gqlServer.URL)
|
||||||
|
|
||||||
|
rr := do(h, http.MethodPatch, "/v1/sessions/patch", "/v1/sessions/patch",
|
||||||
|
`{"status":"paused"}`, "session_id=abc-123")
|
||||||
|
if rr.Code != http.StatusBadGateway {
|
||||||
|
t.Fatalf("want 502, got %d — body: %s", rr.Code, rr.Body.String())
|
||||||
|
}
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
// ── 5. handleCreateMilestone ─────────────────────────────────────────────────
|
||||||
|
|
||||||
|
func TestHandleCreateMilestone(t *testing.T) {
|
||||||
|
t.Run("missing session_id → 400", func(t *testing.T) {
|
||||||
|
gql := newGQLServer(t, nil)
|
||||||
|
defer gql.Close()
|
||||||
|
h := buildHandler(gql.URL)
|
||||||
|
|
||||||
|
rr := do(h, http.MethodPost, "/v1/sessions/milestones/create", "/v1/sessions/milestones/create",
|
||||||
|
`{"title":"step 1"}`, "")
|
||||||
|
if rr.Code != http.StatusBadRequest {
|
||||||
|
t.Fatalf("want 400, got %d", rr.Code)
|
||||||
|
}
|
||||||
|
})
|
||||||
|
|
||||||
|
t.Run("missing title → 400", func(t *testing.T) {
|
||||||
|
gql := newGQLServer(t, nil)
|
||||||
|
defer gql.Close()
|
||||||
|
h := buildHandler(gql.URL)
|
||||||
|
|
||||||
|
rr := do(h, http.MethodPost, "/v1/sessions/milestones/create", "/v1/sessions/milestones/create",
|
||||||
|
`{"order_idx":1}`, "session_id=abc-123")
|
||||||
|
if rr.Code != http.StatusBadRequest {
|
||||||
|
t.Fatalf("want 400, got %d", rr.Code)
|
||||||
|
}
|
||||||
|
})
|
||||||
|
|
||||||
|
t.Run("upsert error → 502", func(t *testing.T) {
|
||||||
|
gqlErr := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||||
|
w.Header().Set("Content-Type", "application/json")
|
||||||
|
fmt.Fprint(w, `{"errors":[{"message":"write error"}]}`)
|
||||||
|
}))
|
||||||
|
defer gqlErr.Close()
|
||||||
|
h := buildHandler(gqlErr.URL)
|
||||||
|
|
||||||
|
rr := do(h, http.MethodPost, "/v1/sessions/milestones/create", "/v1/sessions/milestones/create",
|
||||||
|
`{"title":"step 1","order_idx":0}`, "session_id=abc-123")
|
||||||
|
if rr.Code != http.StatusBadGateway {
|
||||||
|
t.Fatalf("want 502, got %d — body: %s", rr.Code, rr.Body.String())
|
||||||
|
}
|
||||||
|
})
|
||||||
|
|
||||||
|
t.Run("success → 201 with milestone_id, session_id, status=pending", func(t *testing.T) {
|
||||||
|
gql := newGQLServer(t, nil)
|
||||||
|
defer gql.Close()
|
||||||
|
h := buildHandler(gql.URL)
|
||||||
|
|
||||||
|
sid := "abc-123"
|
||||||
|
rr := do(h, http.MethodPost, "/v1/sessions/milestones/create", "/v1/sessions/milestones/create",
|
||||||
|
`{"title":"phase 1","order_idx":0}`, "session_id="+sid)
|
||||||
|
if rr.Code != http.StatusCreated {
|
||||||
|
t.Fatalf("want 201, got %d — body: %s", rr.Code, rr.Body.String())
|
||||||
|
}
|
||||||
|
body := jsonBody(t, rr)
|
||||||
|
if _, ok := body["milestone_id"]; !ok {
|
||||||
|
t.Fatal("missing milestone_id")
|
||||||
|
}
|
||||||
|
if body["session_id"] != sid {
|
||||||
|
t.Fatalf("want session_id=%s, got %v", sid, body["session_id"])
|
||||||
|
}
|
||||||
|
if body["status"] != "pending" {
|
||||||
|
t.Fatalf("want status=pending, got %v", body["status"])
|
||||||
|
}
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
// ── 6. handleListEpisodes ─────────────────────────────────────────────────────
|
||||||
|
|
||||||
|
func TestHandleListEpisodes(t *testing.T) {
|
||||||
|
t.Run("missing session_id → 400", func(t *testing.T) {
|
||||||
|
gql := newGQLServer(t, nil)
|
||||||
|
defer gql.Close()
|
||||||
|
h := buildHandler(gql.URL)
|
||||||
|
|
||||||
|
rr := do(h, http.MethodGet, "/v1/sessions/episodes/list", "/v1/sessions/episodes/list",
|
||||||
|
"", "")
|
||||||
|
if rr.Code != http.StatusBadRequest {
|
||||||
|
t.Fatalf("want 400, got %d", rr.Code)
|
||||||
|
}
|
||||||
|
})
|
||||||
|
|
||||||
|
t.Run("limit > 200 clamped → gql called with limit=200", func(t *testing.T) {
|
||||||
|
var captured []byte
|
||||||
|
gql := newGQLServer(t, &captured)
|
||||||
|
defer gql.Close()
|
||||||
|
h := buildHandler(gql.URL)
|
||||||
|
|
||||||
|
rr := do(h, http.MethodGet, "/v1/sessions/episodes/list", "/v1/sessions/episodes/list",
|
||||||
|
"", "session_id=abc-123&limit=999")
|
||||||
|
if rr.Code != http.StatusOK {
|
||||||
|
t.Fatalf("want 200, got %d — body: %s", rr.Code, rr.Body.String())
|
||||||
|
}
|
||||||
|
// Verify GQL was called (body non-empty).
|
||||||
|
if len(captured) == 0 {
|
||||||
|
t.Fatal("gql not called")
|
||||||
|
}
|
||||||
|
// GQL was called; handler clamps limit to 200.
|
||||||
|
})
|
||||||
|
|
||||||
|
t.Run("success → 200", func(t *testing.T) {
|
||||||
|
gql := newGQLServer(t, nil)
|
||||||
|
defer gql.Close()
|
||||||
|
h := buildHandler(gql.URL)
|
||||||
|
|
||||||
|
rr := do(h, http.MethodGet, "/v1/sessions/episodes/list", "/v1/sessions/episodes/list",
|
||||||
|
"", "session_id=abc-123")
|
||||||
|
if rr.Code != http.StatusOK {
|
||||||
|
t.Fatalf("want 200, got %d — body: %s", rr.Code, rr.Body.String())
|
||||||
|
}
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
// ── 7. handleGetEpisode ───────────────────────────────────────────────────────
|
||||||
|
|
||||||
|
func TestHandleGetEpisode(t *testing.T) {
|
||||||
|
t.Run("missing session_id → 400", func(t *testing.T) {
|
||||||
|
gql := newGQLServer(t, nil)
|
||||||
|
defer gql.Close()
|
||||||
|
h := buildHandler(gql.URL)
|
||||||
|
|
||||||
|
rr := do(h, http.MethodGet, "/v1/sessions/episodes/get", "/v1/sessions/episodes/get",
|
||||||
|
"", "episode_id=deadbeef")
|
||||||
|
if rr.Code != http.StatusBadRequest {
|
||||||
|
t.Fatalf("want 400, got %d", rr.Code)
|
||||||
|
}
|
||||||
|
})
|
||||||
|
|
||||||
|
t.Run("missing episode_id → 400", func(t *testing.T) {
|
||||||
|
gql := newGQLServer(t, nil)
|
||||||
|
defer gql.Close()
|
||||||
|
h := buildHandler(gql.URL)
|
||||||
|
|
||||||
|
rr := do(h, http.MethodGet, "/v1/sessions/episodes/get", "/v1/sessions/episodes/get",
|
||||||
|
"", "session_id=abc-123")
|
||||||
|
if rr.Code != http.StatusBadRequest {
|
||||||
|
t.Fatalf("want 400, got %d", rr.Code)
|
||||||
|
}
|
||||||
|
})
|
||||||
|
|
||||||
|
t.Run("success → 200", func(t *testing.T) {
|
||||||
|
gql := newGQLServer(t, nil)
|
||||||
|
defer gql.Close()
|
||||||
|
h := buildHandler(gql.URL)
|
||||||
|
|
||||||
|
rr := do(h, http.MethodGet, "/v1/sessions/episodes/get", "/v1/sessions/episodes/get",
|
||||||
|
"", "session_id=abc-123&episode_id=deadbeef")
|
||||||
|
if rr.Code != http.StatusOK {
|
||||||
|
t.Fatalf("want 200, got %d — body: %s", rr.Code, rr.Body.String())
|
||||||
|
}
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
// ── 8. EpisodeID helper ───────────────────────────────────────────────────────
|
||||||
|
|
||||||
|
func TestEpisodeID(t *testing.T) {
|
||||||
|
sid := "session-abc"
|
||||||
|
mid := "milestone-xyz"
|
||||||
|
attempt := 1
|
||||||
|
h := sha256.New()
|
||||||
|
h.Write([]byte(fmt.Sprintf("%s%s%d", sid, mid, attempt)))
|
||||||
|
expected := fmt.Sprintf("%x", h.Sum(nil))
|
||||||
|
|
||||||
|
got := EpisodeID(sid, mid, attempt)
|
||||||
|
if got != expected {
|
||||||
|
t.Fatalf("EpisodeID mismatch: want %s, got %s", expected, got)
|
||||||
|
}
|
||||||
|
|
||||||
|
// Sanity: different inputs differ.
|
||||||
|
if EpisodeID(sid, mid, 2) == got {
|
||||||
|
t.Fatal("attempt 2 must differ from attempt 1")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// ── 9. unknown path → 404 ────────────────────────────────────────────────────
|
||||||
|
|
||||||
|
func TestUnknownPath(t *testing.T) {
|
||||||
|
gql := newGQLServer(t, nil)
|
||||||
|
defer gql.Close()
|
||||||
|
h := buildHandler(gql.URL)
|
||||||
|
|
||||||
|
rr := do(h, http.MethodGet, "/v1/sessions/unknown", "/v1/sessions/unknown", "", "")
|
||||||
|
if rr.Code != http.StatusNotFound {
|
||||||
|
t.Fatalf("want 404, got %d", rr.Code)
|
||||||
|
}
|
||||||
|
body := jsonBody(t, rr)
|
||||||
|
if body["error"] == nil {
|
||||||
|
t.Fatal("missing error field in 404 response")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// ── table-driven summary for handleCreate (extra variants) ───────────────────
|
||||||
|
|
||||||
|
func TestHandleCreate_Table(t *testing.T) {
|
||||||
|
tests := []struct {
|
||||||
|
name string
|
||||||
|
body string
|
||||||
|
wantStatus int
|
||||||
|
}{
|
||||||
|
{"valid", `{"name":"s1"}`, http.StatusCreated},
|
||||||
|
{"empty name", `{"name":""}`, http.StatusBadRequest},
|
||||||
|
{"empty_name_string", `{"name":""}`, http.StatusBadRequest},
|
||||||
|
}
|
||||||
|
|
||||||
|
for _, tc := range tests {
|
||||||
|
t.Run(tc.name, func(t *testing.T) {
|
||||||
|
gql := newGQLServer(t, nil)
|
||||||
|
defer gql.Close()
|
||||||
|
h := buildHandler(gql.URL)
|
||||||
|
|
||||||
|
rr := do(h, http.MethodPost, "/v1/sessions/create", "/v1/sessions/create", tc.body, "")
|
||||||
|
if rr.Code != tc.wantStatus {
|
||||||
|
t.Fatalf("want %d, got %d — body: %s", tc.wantStatus, rr.Code, rr.Body.String())
|
||||||
|
}
|
||||||
|
})
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// ── table-driven summary for handlePatch ─────────────────────────────────────
|
||||||
|
|
||||||
|
func TestHandlePatch_Table(t *testing.T) {
|
||||||
|
validSID := "session_id=abc-123"
|
||||||
|
tests := []struct {
|
||||||
|
name string
|
||||||
|
query string
|
||||||
|
body string
|
||||||
|
wantStatus int
|
||||||
|
}{
|
||||||
|
{"ok", validSID, `{"status":"paused"}`, http.StatusOK},
|
||||||
|
{"no session_id", "", `{"status":"paused"}`, http.StatusBadRequest},
|
||||||
|
{"bad json", validSID, `{bad}`, http.StatusBadRequest},
|
||||||
|
}
|
||||||
|
for _, tc := range tests {
|
||||||
|
t.Run(tc.name, func(t *testing.T) {
|
||||||
|
gql := newGQLServer(t, nil)
|
||||||
|
defer gql.Close()
|
||||||
|
h := buildHandler(gql.URL)
|
||||||
|
|
||||||
|
rr := do(h, http.MethodPatch, "/v1/sessions/patch", "/v1/sessions/patch",
|
||||||
|
tc.body, tc.query)
|
||||||
|
if rr.Code != tc.wantStatus {
|
||||||
|
t.Fatalf("want %d, got %d — body: %s", tc.wantStatus, rr.Code, rr.Body.String())
|
||||||
|
}
|
||||||
|
})
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// ── verify upsertEntity mutation payload ─────────────────────────────────────
|
||||||
|
|
||||||
|
func TestHandleCreate_IngestPayload(t *testing.T) {
|
||||||
|
var captured []byte
|
||||||
|
gql := newGQLServer(t, &captured)
|
||||||
|
defer gql.Close()
|
||||||
|
h := buildHandler(gql.URL)
|
||||||
|
|
||||||
|
rr := do(h, http.MethodPost, "/v1/sessions/create", "/v1/sessions/create",
|
||||||
|
`{"name":"test-session","goal":"do stuff"}`, "")
|
||||||
|
if rr.Code != http.StatusCreated {
|
||||||
|
t.Fatalf("want 201, got %d", rr.Code)
|
||||||
|
}
|
||||||
|
bs := string(captured)
|
||||||
|
if !strings.Contains(bs, "upsertEntity") {
|
||||||
|
t.Fatalf("gql body missing 'upsertEntity': %s", bs)
|
||||||
|
}
|
||||||
|
if !strings.Contains(bs, "SESSION") {
|
||||||
|
t.Fatalf("gql body missing 'SESSION': %s", bs)
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,492 @@
|
|||||||
|
// Package sessions — SandboxManager provisions per-session k8s resources:
|
||||||
|
//
|
||||||
|
// PVC workspace-session-<uuid> Longhorn 1Gi
|
||||||
|
// Pod pi-session-<uuid> single container, /workspace <- PVC
|
||||||
|
// Svc pi-session-<uuid> ClusterIP :8787
|
||||||
|
// NetPol pi-session-sandbox shared for all pi-session pods (created once)
|
||||||
|
//
|
||||||
|
// All k8s API calls use raw HTTP + in-cluster SA token (no client-go).
|
||||||
|
// Same pattern as internal/webhook/promoter.go.
|
||||||
|
package sessions
|
||||||
|
|
||||||
|
import (
|
||||||
|
"bytes"
|
||||||
|
"context"
|
||||||
|
"crypto/tls"
|
||||||
|
"crypto/x509"
|
||||||
|
"encoding/json"
|
||||||
|
"fmt"
|
||||||
|
"io"
|
||||||
|
"log"
|
||||||
|
"net/http"
|
||||||
|
"os"
|
||||||
|
"strings"
|
||||||
|
"time"
|
||||||
|
)
|
||||||
|
|
||||||
|
const (
|
||||||
|
sandboxNamespace = "poimen"
|
||||||
|
sandboxImage = "forgejo.riotpiao.com/riotpiao-poimen/poimen-pi-session-runtime:latest"
|
||||||
|
sandboxPort = 8787
|
||||||
|
defaultPVCSize = "1Gi"
|
||||||
|
defaultStorageClass = "longhorn"
|
||||||
|
defaultReadyTimeout = 30 * time.Second
|
||||||
|
defaultReaperInterval = 5 * time.Minute
|
||||||
|
sandboxSATokenPath = "/var/run/secrets/kubernetes.io/serviceaccount/token"
|
||||||
|
sandboxCACertPath = "/var/run/secrets/kubernetes.io/serviceaccount/ca.crt"
|
||||||
|
)
|
||||||
|
|
||||||
|
// SandboxManager manages the k8s lifecycle of pi-session Pods, PVCs, Services.
|
||||||
|
type SandboxManager struct {
|
||||||
|
k8sHost string
|
||||||
|
k8sClient *http.Client
|
||||||
|
pvcSize string
|
||||||
|
readyTimeout time.Duration
|
||||||
|
}
|
||||||
|
|
||||||
|
// SandboxResult is returned by CreateOrResume.
|
||||||
|
type SandboxResult struct {
|
||||||
|
SessionID string `json:"session_id"`
|
||||||
|
PodName string `json:"pod_name"`
|
||||||
|
PVCName string `json:"pvc_name"`
|
||||||
|
SvcDNS string `json:"svc_dns"`
|
||||||
|
Resumed bool `json:"resumed"` // true = pod already existed
|
||||||
|
}
|
||||||
|
|
||||||
|
// NewSandboxManager creates a SandboxManager from env + in-cluster SA files.
|
||||||
|
func NewSandboxManager() *SandboxManager {
|
||||||
|
host := os.Getenv("K8S_API_HOST")
|
||||||
|
if host == "" {
|
||||||
|
host = "https://kubernetes.default.svc"
|
||||||
|
}
|
||||||
|
pvcSize := os.Getenv("SANDBOX_PVC_SIZE")
|
||||||
|
if pvcSize == "" {
|
||||||
|
pvcSize = defaultPVCSize
|
||||||
|
}
|
||||||
|
caCert, _ := os.ReadFile(sandboxCACertPath)
|
||||||
|
return &SandboxManager{
|
||||||
|
k8sHost: host,
|
||||||
|
k8sClient: buildSandboxK8sClient(caCert),
|
||||||
|
pvcSize: pvcSize,
|
||||||
|
readyTimeout: defaultReadyTimeout,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func buildSandboxK8sClient(caCert []byte) *http.Client {
|
||||||
|
transport := &http.Transport{}
|
||||||
|
if len(caCert) > 0 {
|
||||||
|
pool := x509.NewCertPool()
|
||||||
|
pool.AppendCertsFromPEM(caCert)
|
||||||
|
transport.TLSClientConfig = &tls.Config{RootCAs: pool}
|
||||||
|
} else {
|
||||||
|
log.Printf("sandbox: WARNING no SA CA cert at %s — using InsecureSkipVerify", sandboxCACertPath)
|
||||||
|
transport.TLSClientConfig = &tls.Config{InsecureSkipVerify: true} //nolint:gosec
|
||||||
|
}
|
||||||
|
return &http.Client{Timeout: 30 * time.Second, Transport: transport}
|
||||||
|
}
|
||||||
|
|
||||||
|
func sandboxSAToken() string {
|
||||||
|
b, err := os.ReadFile(sandboxSATokenPath)
|
||||||
|
if err != nil {
|
||||||
|
log.Printf("sandbox: WARNING failed to read SA token %s: %v", sandboxSATokenPath, err)
|
||||||
|
return ""
|
||||||
|
}
|
||||||
|
return strings.TrimSpace(string(b))
|
||||||
|
}
|
||||||
|
|
||||||
|
func pvcName(sessionID string) string { return "workspace-session-" + sessionID }
|
||||||
|
func podName(sessionID string) string { return "pi-session-" + sessionID }
|
||||||
|
func svcName(sessionID string) string { return "pi-session-" + sessionID }
|
||||||
|
func svcDNS(sessionID string) string {
|
||||||
|
return fmt.Sprintf("%s.%s.svc.cluster.local:%d", svcName(sessionID), sandboxNamespace, sandboxPort)
|
||||||
|
}
|
||||||
|
|
||||||
|
// ── public API ────────────────────────────────────────────────────────────
|
||||||
|
|
||||||
|
// CreateOrResume provisions or resumes a sandbox for sessionID.
|
||||||
|
// Idempotent: if the pod is already Running, returns its DNS immediately.
|
||||||
|
func (s *SandboxManager) CreateOrResume(ctx context.Context, sessionID string) (*SandboxResult, error) {
|
||||||
|
res := &SandboxResult{
|
||||||
|
SessionID: sessionID,
|
||||||
|
PodName: podName(sessionID),
|
||||||
|
PVCName: pvcName(sessionID),
|
||||||
|
SvcDNS: svcDNS(sessionID),
|
||||||
|
}
|
||||||
|
|
||||||
|
// 1. Ensure PVC exists
|
||||||
|
if err := s.ensurePVC(ctx, sessionID); err != nil {
|
||||||
|
return nil, fmt.Errorf("ensure PVC: %w", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
// 2. Check if pod already running (resume case)
|
||||||
|
phase, err := s.getPodPhase(ctx, sessionID)
|
||||||
|
if err != nil {
|
||||||
|
return nil, fmt.Errorf("check pod: %w", err)
|
||||||
|
}
|
||||||
|
if phase == "Running" {
|
||||||
|
res.Resumed = true
|
||||||
|
return res, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// 3. Ensure NetworkPolicy for all pi-session pods exists
|
||||||
|
if err := s.ensureNetworkPolicy(ctx); err != nil {
|
||||||
|
log.Printf("sandbox: WARNING failed to ensure network policy: %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
// 4. Create Pod + Service
|
||||||
|
if err := s.createPod(ctx, sessionID); err != nil {
|
||||||
|
return nil, fmt.Errorf("create pod: %w", err)
|
||||||
|
}
|
||||||
|
if err := s.ensureService(ctx, sessionID); err != nil {
|
||||||
|
return nil, fmt.Errorf("ensure service: %w", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
// 5. Wait for Ready
|
||||||
|
if err := s.waitPodReady(ctx, sessionID); err != nil {
|
||||||
|
return nil, fmt.Errorf("wait ready: %w", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
return res, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// Delete removes the Pod, Service, and PVC for sessionID.
|
||||||
|
func (s *SandboxManager) Delete(ctx context.Context, sessionID string) error {
|
||||||
|
var errs []string
|
||||||
|
for _, fn := range []func(context.Context, string) error{
|
||||||
|
s.deletePod, s.deleteService, s.deletePVC,
|
||||||
|
} {
|
||||||
|
if err := fn(ctx, sessionID); err != nil && !isNotFound(err) {
|
||||||
|
errs = append(errs, err.Error())
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if len(errs) > 0 {
|
||||||
|
return fmt.Errorf("delete sandbox: %s", strings.Join(errs, "; "))
|
||||||
|
}
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// ReapIdlePods lists all pi-session-* pods and deletes those that have exited
|
||||||
|
// (Succeeded/Failed) — bridge.js self-exits after idle timeout.
|
||||||
|
// The PVC is preserved so the next CreateOrResume reattaches the same workspace.
|
||||||
|
func (s *SandboxManager) ReapIdlePods(ctx context.Context) {
|
||||||
|
pods, err := s.listSandboxPods(ctx)
|
||||||
|
if err != nil {
|
||||||
|
log.Printf("sandbox reaper: list pods: %v", err)
|
||||||
|
return
|
||||||
|
}
|
||||||
|
for _, podAny := range pods {
|
||||||
|
pod, _ := podAny.(map[string]any)
|
||||||
|
phase, _ := pod["status"].(map[string]any)["phase"].(string)
|
||||||
|
name, _ := pod["metadata"].(map[string]any)["name"].(string)
|
||||||
|
if phase == "Succeeded" || phase == "Failed" {
|
||||||
|
log.Printf("sandbox reaper: deleting exited pod %s (phase=%s)", name, phase)
|
||||||
|
if err := s.k8sDelete(ctx, "/api/v1/namespaces/"+sandboxNamespace+"/pods/"+name); err != nil && !isNotFound(err) {
|
||||||
|
log.Printf("sandbox reaper: delete pod %s: %v", name, err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// ── k8s resource helpers ─────────────────────────────────────────────────
|
||||||
|
|
||||||
|
func (s *SandboxManager) ensurePVC(ctx context.Context, sessionID string) error {
|
||||||
|
name := pvcName(sessionID)
|
||||||
|
_, err := s.k8sGet(ctx, "/api/v1/namespaces/"+sandboxNamespace+"/persistentvolumeclaims/"+name)
|
||||||
|
if err == nil {
|
||||||
|
return nil // already exists
|
||||||
|
}
|
||||||
|
if !isNotFound(err) {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
// Create
|
||||||
|
body := map[string]any{
|
||||||
|
"apiVersion": "v1",
|
||||||
|
"kind": "PersistentVolumeClaim",
|
||||||
|
"metadata": map[string]any{
|
||||||
|
"name": name, "namespace": sandboxNamespace,
|
||||||
|
"labels": map[string]string{"app": "pi-session", "session-id": sessionID},
|
||||||
|
},
|
||||||
|
"spec": map[string]any{
|
||||||
|
"accessModes": []string{"ReadWriteOnce"},
|
||||||
|
"storageClassName": defaultStorageClass,
|
||||||
|
"resources": map[string]any{
|
||||||
|
"requests": map[string]string{"storage": s.pvcSize},
|
||||||
|
},
|
||||||
|
},
|
||||||
|
}
|
||||||
|
return s.k8sCreate(ctx, "/api/v1/namespaces/"+sandboxNamespace+"/persistentvolumeclaims", body)
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *SandboxManager) getPodPhase(ctx context.Context, sessionID string) (string, error) {
|
||||||
|
pod, err := s.k8sGet(ctx, "/api/v1/namespaces/"+sandboxNamespace+"/pods/"+podName(sessionID))
|
||||||
|
if isNotFound(err) {
|
||||||
|
return "", nil
|
||||||
|
}
|
||||||
|
if err != nil {
|
||||||
|
return "", err
|
||||||
|
}
|
||||||
|
phase, _ := pod["status"].(map[string]any)["phase"].(string)
|
||||||
|
return phase, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *SandboxManager) createPod(ctx context.Context, sessionID string) error {
|
||||||
|
body := map[string]any{
|
||||||
|
"apiVersion": "v1",
|
||||||
|
"kind": "Pod",
|
||||||
|
"metadata": map[string]any{
|
||||||
|
"name": podName(sessionID), "namespace": sandboxNamespace,
|
||||||
|
"labels": map[string]string{"app": "pi-session", "session-id": sessionID},
|
||||||
|
},
|
||||||
|
"spec": map[string]any{
|
||||||
|
"restartPolicy": "Never",
|
||||||
|
"automountServiceAccountToken": false,
|
||||||
|
"securityContext": map[string]any{
|
||||||
|
"runAsNonRoot": true,
|
||||||
|
"runAsUser": int64(1000),
|
||||||
|
"seccompProfile": map[string]string{"type": "RuntimeDefault"},
|
||||||
|
},
|
||||||
|
"containers": []map[string]any{{
|
||||||
|
"name": "pi-session",
|
||||||
|
"image": sandboxImage,
|
||||||
|
"ports": []map[string]any{{"containerPort": sandboxPort}},
|
||||||
|
"securityContext": map[string]any{
|
||||||
|
"readOnlyRootFilesystem": true,
|
||||||
|
"allowPrivilegeEscalation": false,
|
||||||
|
"capabilities": map[string]any{"drop": []string{"ALL"}},
|
||||||
|
},
|
||||||
|
"resources": map[string]any{
|
||||||
|
"requests": map[string]string{"cpu": "500m", "memory": "512Mi"},
|
||||||
|
"limits": map[string]string{"cpu": "1", "memory": "1Gi"},
|
||||||
|
},
|
||||||
|
"readinessProbe": map[string]any{
|
||||||
|
"tcpSocket": map[string]any{"port": sandboxPort},
|
||||||
|
"initialDelaySeconds": 3,
|
||||||
|
"periodSeconds": 2,
|
||||||
|
"failureThreshold": 15,
|
||||||
|
},
|
||||||
|
"volumeMounts": []map[string]any{
|
||||||
|
{"name": "workspace", "mountPath": "/workspace"},
|
||||||
|
{"name": "tmp", "mountPath": "/tmp"},
|
||||||
|
},
|
||||||
|
}},
|
||||||
|
"volumes": []map[string]any{
|
||||||
|
{"name": "workspace", "persistentVolumeClaim": map[string]string{"claimName": pvcName(sessionID)}},
|
||||||
|
{"name": "tmp", "emptyDir": map[string]any{}},
|
||||||
|
},
|
||||||
|
},
|
||||||
|
}
|
||||||
|
return s.k8sCreate(ctx, "/api/v1/namespaces/"+sandboxNamespace+"/pods", body)
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *SandboxManager) ensureService(ctx context.Context, sessionID string) error {
|
||||||
|
name := svcName(sessionID)
|
||||||
|
_, err := s.k8sGet(ctx, "/api/v1/namespaces/"+sandboxNamespace+"/services/"+name)
|
||||||
|
if err == nil {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
if !isNotFound(err) {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
body := map[string]any{
|
||||||
|
"apiVersion": "v1",
|
||||||
|
"kind": "Service",
|
||||||
|
"metadata": map[string]any{
|
||||||
|
"name": name, "namespace": sandboxNamespace,
|
||||||
|
"labels": map[string]string{"app": "pi-session", "session-id": sessionID},
|
||||||
|
},
|
||||||
|
"spec": map[string]any{
|
||||||
|
"type": "ClusterIP",
|
||||||
|
"selector": map[string]string{"app": "pi-session", "session-id": sessionID},
|
||||||
|
"ports": []map[string]any{{
|
||||||
|
"port": sandboxPort, "targetPort": sandboxPort, "protocol": "TCP",
|
||||||
|
}},
|
||||||
|
},
|
||||||
|
}
|
||||||
|
return s.k8sCreate(ctx, "/api/v1/namespaces/"+sandboxNamespace+"/services", body)
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *SandboxManager) ensureNetworkPolicy(ctx context.Context) error {
|
||||||
|
const npName = "pi-session-sandbox"
|
||||||
|
_, err := s.k8sGet(ctx, "/apis/networking.k8s.io/v1/namespaces/"+sandboxNamespace+"/networkpolicies/"+npName)
|
||||||
|
if err == nil {
|
||||||
|
return nil // already exists
|
||||||
|
}
|
||||||
|
if !isNotFound(err) {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
body := map[string]any{
|
||||||
|
"apiVersion": "networking.k8s.io/v1",
|
||||||
|
"kind": "NetworkPolicy",
|
||||||
|
"metadata": map[string]any{
|
||||||
|
"name": npName, "namespace": sandboxNamespace,
|
||||||
|
"labels": map[string]string{"app": "pi-session"},
|
||||||
|
},
|
||||||
|
"spec": map[string]any{
|
||||||
|
"podSelector": map[string]any{"matchLabels": map[string]string{"app": "pi-session"}},
|
||||||
|
"policyTypes": []string{"Ingress", "Egress"},
|
||||||
|
"ingress": []map[string]any{{
|
||||||
|
"from": []map[string]any{{
|
||||||
|
"podSelector": map[string]any{"matchLabels": map[string]string{"app": "poimen-frontend"}},
|
||||||
|
}},
|
||||||
|
"ports": []map[string]any{{"protocol": "TCP", "port": sandboxPort}},
|
||||||
|
}},
|
||||||
|
"egress": []map[string]any{
|
||||||
|
// DNS
|
||||||
|
{
|
||||||
|
"to": []map[string]any{{"namespaceSelector": map[string]any{"matchLabels": map[string]string{"kubernetes.io/metadata.name": "kube-system"}}}},
|
||||||
|
"ports": []map[string]any{{"protocol": "UDP", "port": 53}, {"protocol": "TCP", "port": 53}},
|
||||||
|
},
|
||||||
|
// api-gateway (LLM inference)
|
||||||
|
{
|
||||||
|
"to": []map[string]any{{"namespaceSelector": map[string]any{"matchLabels": map[string]string{"kubernetes.io/metadata.name": "api"}}}},
|
||||||
|
"ports": []map[string]any{{"protocol": "TCP", "port": 8080}},
|
||||||
|
},
|
||||||
|
},
|
||||||
|
},
|
||||||
|
}
|
||||||
|
return s.k8sCreate(ctx, "/apis/networking.k8s.io/v1/namespaces/"+sandboxNamespace+"/networkpolicies", body)
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *SandboxManager) waitPodReady(ctx context.Context, sessionID string) error {
|
||||||
|
deadline := time.Now().Add(s.readyTimeout)
|
||||||
|
for time.Now().Before(deadline) {
|
||||||
|
pod, err := s.k8sGet(ctx, "/api/v1/namespaces/"+sandboxNamespace+"/pods/"+podName(sessionID))
|
||||||
|
if err != nil {
|
||||||
|
return fmt.Errorf("poll pod: %w", err)
|
||||||
|
}
|
||||||
|
status, _ := pod["status"].(map[string]any)
|
||||||
|
phase, _ := status["phase"].(string)
|
||||||
|
if phase == "Failed" || phase == "Succeeded" {
|
||||||
|
return fmt.Errorf("pod entered terminal phase %s before ready", phase)
|
||||||
|
}
|
||||||
|
if phase == "Running" {
|
||||||
|
for _, cond := range toSlice(status["conditions"]) {
|
||||||
|
if c, ok := cond.(map[string]any); ok {
|
||||||
|
if c["type"] == "Ready" && c["status"] == "True" {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
select {
|
||||||
|
case <-time.After(2 * time.Second):
|
||||||
|
case <-ctx.Done():
|
||||||
|
return ctx.Err()
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return fmt.Errorf("pod not ready after %v", s.readyTimeout)
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *SandboxManager) deletePod(ctx context.Context, sessionID string) error {
|
||||||
|
return s.k8sDelete(ctx, "/api/v1/namespaces/"+sandboxNamespace+"/pods/"+podName(sessionID))
|
||||||
|
}
|
||||||
|
func (s *SandboxManager) deleteService(ctx context.Context, sessionID string) error {
|
||||||
|
return s.k8sDelete(ctx, "/api/v1/namespaces/"+sandboxNamespace+"/services/"+svcName(sessionID))
|
||||||
|
}
|
||||||
|
func (s *SandboxManager) deletePVC(ctx context.Context, sessionID string) error {
|
||||||
|
return s.k8sDelete(ctx, "/api/v1/namespaces/"+sandboxNamespace+"/persistentvolumeclaims/"+pvcName(sessionID))
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *SandboxManager) listSandboxPods(ctx context.Context) ([]any, error) {
|
||||||
|
body, err := s.k8sGet(ctx, "/api/v1/namespaces/"+sandboxNamespace+"/pods?labelSelector=app%3Dpi-session")
|
||||||
|
if err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
items, _ := body["items"].([]any)
|
||||||
|
return items, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// ── raw k8s HTTP ─────────────────────────────────────────────────────────
|
||||||
|
|
||||||
|
func (s *SandboxManager) k8sGet(ctx context.Context, path string) (map[string]any, error) {
|
||||||
|
req, err := http.NewRequestWithContext(ctx, http.MethodGet, s.k8sHost+path, nil)
|
||||||
|
if err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
req.Header.Set("Authorization", "Bearer "+sandboxSAToken())
|
||||||
|
resp, err := s.k8sClient.Do(req)
|
||||||
|
if err != nil {
|
||||||
|
return nil, fmt.Errorf("k8s GET %s: %w", path, err)
|
||||||
|
}
|
||||||
|
defer resp.Body.Close()
|
||||||
|
if resp.StatusCode == 404 {
|
||||||
|
return nil, &k8sNotFound{path: path}
|
||||||
|
}
|
||||||
|
if resp.StatusCode >= 400 {
|
||||||
|
b, _ := io.ReadAll(io.LimitReader(resp.Body, 512))
|
||||||
|
return nil, fmt.Errorf("k8s GET %s returned %d: %s", path, resp.StatusCode, b)
|
||||||
|
}
|
||||||
|
var result map[string]any
|
||||||
|
if err := json.NewDecoder(resp.Body).Decode(&result); err != nil {
|
||||||
|
return nil, fmt.Errorf("decode: %w", err)
|
||||||
|
}
|
||||||
|
return result, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *SandboxManager) k8sCreate(ctx context.Context, path string, body any) error {
|
||||||
|
payload, err := json.Marshal(body)
|
||||||
|
if err != nil {
|
||||||
|
return fmt.Errorf("marshal: %w", err)
|
||||||
|
}
|
||||||
|
req, err := http.NewRequestWithContext(ctx, http.MethodPost, s.k8sHost+path, bytes.NewReader(payload))
|
||||||
|
if err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
req.Header.Set("Authorization", "Bearer "+sandboxSAToken())
|
||||||
|
req.Header.Set("Content-Type", "application/json")
|
||||||
|
resp, err := s.k8sClient.Do(req)
|
||||||
|
if err != nil {
|
||||||
|
return fmt.Errorf("k8s POST %s: %w", path, err)
|
||||||
|
}
|
||||||
|
defer resp.Body.Close()
|
||||||
|
if resp.StatusCode == 409 { // AlreadyExists — idempotent
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
if resp.StatusCode < 200 || resp.StatusCode >= 300 {
|
||||||
|
b, _ := io.ReadAll(io.LimitReader(resp.Body, 512))
|
||||||
|
return fmt.Errorf("k8s POST %s returned %d: %s", path, resp.StatusCode, b)
|
||||||
|
}
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *SandboxManager) k8sDelete(ctx context.Context, path string) error {
|
||||||
|
req, err := http.NewRequestWithContext(ctx, http.MethodDelete, s.k8sHost+path, nil)
|
||||||
|
if err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
req.Header.Set("Authorization", "Bearer "+sandboxSAToken())
|
||||||
|
resp, err := s.k8sClient.Do(req)
|
||||||
|
if err != nil {
|
||||||
|
return fmt.Errorf("k8s DELETE %s: %w", path, err)
|
||||||
|
}
|
||||||
|
defer resp.Body.Close()
|
||||||
|
if resp.StatusCode == 404 {
|
||||||
|
return &k8sNotFound{path: path}
|
||||||
|
}
|
||||||
|
if resp.StatusCode >= 400 {
|
||||||
|
b, _ := io.ReadAll(io.LimitReader(resp.Body, 512))
|
||||||
|
return fmt.Errorf("k8s DELETE %s returned %d: %s", path, resp.StatusCode, b)
|
||||||
|
}
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// ── error helpers ─────────────────────────────────────────────────────────
|
||||||
|
|
||||||
|
type k8sNotFound struct{ path string }
|
||||||
|
|
||||||
|
func (e *k8sNotFound) Error() string { return "not found: " + e.path }
|
||||||
|
|
||||||
|
func isNotFound(err error) bool {
|
||||||
|
if err == nil {
|
||||||
|
return false
|
||||||
|
}
|
||||||
|
_, ok := err.(*k8sNotFound)
|
||||||
|
return ok
|
||||||
|
}
|
||||||
|
|
||||||
|
func toSlice(v any) []any {
|
||||||
|
if s, ok := v.([]any); ok {
|
||||||
|
return s
|
||||||
|
}
|
||||||
|
return nil
|
||||||
|
}
|
||||||
@@ -0,0 +1,271 @@
|
|||||||
|
package sessions
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"encoding/json"
|
||||||
|
"fmt"
|
||||||
|
"net/http"
|
||||||
|
"net/http/httptest"
|
||||||
|
"sync/atomic"
|
||||||
|
"testing"
|
||||||
|
"time"
|
||||||
|
)
|
||||||
|
|
||||||
|
// ── helpers ───────────────────────────────────────────────────────────────────
|
||||||
|
|
||||||
|
// runningReadyPod returns minimal k8s pod JSON: phase Running + Ready condition.
|
||||||
|
func runningReadyPod(name string) map[string]any {
|
||||||
|
return map[string]any{
|
||||||
|
"metadata": map[string]any{"name": name},
|
||||||
|
"status": map[string]any{
|
||||||
|
"phase": "Running",
|
||||||
|
"conditions": []any{
|
||||||
|
map[string]any{"type": "Ready", "status": "True"},
|
||||||
|
},
|
||||||
|
},
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// respondK8sJSON writes a JSON body with the given status code.
|
||||||
|
func respondK8sJSON(w http.ResponseWriter, status int, body any) {
|
||||||
|
w.Header().Set("Content-Type", "application/json")
|
||||||
|
w.WriteHeader(status)
|
||||||
|
json.NewEncoder(w).Encode(body) //nolint:errcheck
|
||||||
|
}
|
||||||
|
|
||||||
|
// newTestSM builds a SandboxManager pointed at the given test server URL.
|
||||||
|
func newTestSM(url string) *SandboxManager {
|
||||||
|
return &SandboxManager{
|
||||||
|
k8sHost: url,
|
||||||
|
k8sClient: &http.Client{Timeout: 5 * time.Second},
|
||||||
|
pvcSize: "1Gi",
|
||||||
|
readyTimeout: 3 * time.Second,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// ── TestSandboxCreateOrResume_CreateNew ───────────────────────────────────────
|
||||||
|
|
||||||
|
// No existing resources → creates PVC, NetworkPolicy, Pod, Service; Resumed=false.
|
||||||
|
func TestSandboxCreateOrResume_CreateNew(t *testing.T) {
|
||||||
|
id := "newid"
|
||||||
|
var podGetCount int32
|
||||||
|
|
||||||
|
mux := http.NewServeMux()
|
||||||
|
|
||||||
|
// PVC: GET → 404, POST → 201
|
||||||
|
mux.HandleFunc("GET /api/v1/namespaces/poimen/persistentvolumeclaims/{name}", func(w http.ResponseWriter, _ *http.Request) {
|
||||||
|
w.WriteHeader(http.StatusNotFound)
|
||||||
|
})
|
||||||
|
mux.HandleFunc("POST /api/v1/namespaces/poimen/persistentvolumeclaims", func(w http.ResponseWriter, _ *http.Request) {
|
||||||
|
w.WriteHeader(http.StatusCreated)
|
||||||
|
})
|
||||||
|
|
||||||
|
// NetworkPolicy (fixed name "pi-session-sandbox"): GET → 404, POST → 201
|
||||||
|
mux.HandleFunc("GET /apis/networking.k8s.io/v1/namespaces/poimen/networkpolicies/{name}", func(w http.ResponseWriter, _ *http.Request) {
|
||||||
|
w.WriteHeader(http.StatusNotFound)
|
||||||
|
})
|
||||||
|
mux.HandleFunc("POST /apis/networking.k8s.io/v1/namespaces/poimen/networkpolicies", func(w http.ResponseWriter, _ *http.Request) {
|
||||||
|
w.WriteHeader(http.StatusCreated)
|
||||||
|
})
|
||||||
|
|
||||||
|
// Pod: first GET (getPodPhase) → 404; second GET (waitPodReady) → Running+Ready; POST → 201
|
||||||
|
mux.HandleFunc("GET /api/v1/namespaces/poimen/pods/{name}", func(w http.ResponseWriter, _ *http.Request) {
|
||||||
|
if atomic.AddInt32(&podGetCount, 1) == 1 {
|
||||||
|
w.WriteHeader(http.StatusNotFound)
|
||||||
|
return
|
||||||
|
}
|
||||||
|
respondK8sJSON(w, http.StatusOK, runningReadyPod(podName(id)))
|
||||||
|
})
|
||||||
|
mux.HandleFunc("POST /api/v1/namespaces/poimen/pods", func(w http.ResponseWriter, _ *http.Request) {
|
||||||
|
w.WriteHeader(http.StatusCreated)
|
||||||
|
})
|
||||||
|
|
||||||
|
// Service: GET → 404, POST → 201
|
||||||
|
mux.HandleFunc("GET /api/v1/namespaces/poimen/services/{name}", func(w http.ResponseWriter, _ *http.Request) {
|
||||||
|
w.WriteHeader(http.StatusNotFound)
|
||||||
|
})
|
||||||
|
mux.HandleFunc("POST /api/v1/namespaces/poimen/services", func(w http.ResponseWriter, _ *http.Request) {
|
||||||
|
w.WriteHeader(http.StatusCreated)
|
||||||
|
})
|
||||||
|
|
||||||
|
srv := httptest.NewServer(mux)
|
||||||
|
defer srv.Close()
|
||||||
|
|
||||||
|
result, err := newTestSM(srv.URL).CreateOrResume(context.Background(), id)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("CreateOrResume: %v", err)
|
||||||
|
}
|
||||||
|
if want := svcDNS(id); result.SvcDNS != want {
|
||||||
|
t.Errorf("SvcDNS: want %q, got %q", want, result.SvcDNS)
|
||||||
|
}
|
||||||
|
if result.Resumed {
|
||||||
|
t.Error("want Resumed=false, got true")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// ── TestSandboxCreateOrResume_Resume ─────────────────────────────────────────
|
||||||
|
|
||||||
|
// PVC exists + pod Running → Resumed=true, no POST calls.
|
||||||
|
func TestSandboxCreateOrResume_Resume(t *testing.T) {
|
||||||
|
id := "resumeid"
|
||||||
|
|
||||||
|
mux := http.NewServeMux()
|
||||||
|
|
||||||
|
// PVC exists
|
||||||
|
mux.HandleFunc("GET /api/v1/namespaces/poimen/persistentvolumeclaims/{name}", func(w http.ResponseWriter, _ *http.Request) {
|
||||||
|
respondK8sJSON(w, http.StatusOK, map[string]any{"metadata": map[string]any{"name": pvcName(id)}})
|
||||||
|
})
|
||||||
|
|
||||||
|
// Pod exists and Running (getPodPhase uses GET /{name})
|
||||||
|
mux.HandleFunc("GET /api/v1/namespaces/poimen/pods/{name}", func(w http.ResponseWriter, _ *http.Request) {
|
||||||
|
respondK8sJSON(w, http.StatusOK, map[string]any{
|
||||||
|
"metadata": map[string]any{"name": podName(id)},
|
||||||
|
"status": map[string]any{"phase": "Running"},
|
||||||
|
})
|
||||||
|
})
|
||||||
|
|
||||||
|
// Any POST → 500 to assert no creates are issued
|
||||||
|
mux.HandleFunc("POST /", func(w http.ResponseWriter, r *http.Request) {
|
||||||
|
t.Errorf("unexpected POST to %s", r.URL.Path)
|
||||||
|
w.WriteHeader(http.StatusInternalServerError)
|
||||||
|
})
|
||||||
|
|
||||||
|
srv := httptest.NewServer(mux)
|
||||||
|
defer srv.Close()
|
||||||
|
|
||||||
|
result, err := newTestSM(srv.URL).CreateOrResume(context.Background(), id)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("CreateOrResume: %v", err)
|
||||||
|
}
|
||||||
|
if !result.Resumed {
|
||||||
|
t.Error("want Resumed=true, got false")
|
||||||
|
}
|
||||||
|
if want := svcDNS(id); result.SvcDNS != want {
|
||||||
|
t.Errorf("SvcDNS: want %q, got %q", want, result.SvcDNS)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// ── TestSandboxDelete ─────────────────────────────────────────────────────────
|
||||||
|
|
||||||
|
// All DELETEs succeed → no error returned.
|
||||||
|
func TestSandboxDelete(t *testing.T) {
|
||||||
|
id := "deleteid"
|
||||||
|
|
||||||
|
mux := http.NewServeMux()
|
||||||
|
mux.HandleFunc("DELETE /api/v1/namespaces/poimen/pods/{name}", func(w http.ResponseWriter, _ *http.Request) {
|
||||||
|
w.WriteHeader(http.StatusOK)
|
||||||
|
})
|
||||||
|
mux.HandleFunc("DELETE /api/v1/namespaces/poimen/services/{name}", func(w http.ResponseWriter, _ *http.Request) {
|
||||||
|
w.WriteHeader(http.StatusOK)
|
||||||
|
})
|
||||||
|
mux.HandleFunc("DELETE /api/v1/namespaces/poimen/persistentvolumeclaims/{name}", func(w http.ResponseWriter, _ *http.Request) {
|
||||||
|
w.WriteHeader(http.StatusOK)
|
||||||
|
})
|
||||||
|
|
||||||
|
srv := httptest.NewServer(mux)
|
||||||
|
defer srv.Close()
|
||||||
|
|
||||||
|
if err := newTestSM(srv.URL).Delete(context.Background(), id); err != nil {
|
||||||
|
t.Errorf("Delete: %v", err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// ── TestSandboxDelete_NotFound ────────────────────────────────────────────────
|
||||||
|
|
||||||
|
// All DELETEs return 404 → isNotFound errors swallowed → no error returned.
|
||||||
|
func TestSandboxDelete_NotFound(t *testing.T) {
|
||||||
|
id := "ghostid"
|
||||||
|
|
||||||
|
mux := http.NewServeMux()
|
||||||
|
mux.HandleFunc("DELETE /api/v1/namespaces/poimen/pods/{name}", func(w http.ResponseWriter, _ *http.Request) {
|
||||||
|
w.WriteHeader(http.StatusNotFound)
|
||||||
|
})
|
||||||
|
mux.HandleFunc("DELETE /api/v1/namespaces/poimen/services/{name}", func(w http.ResponseWriter, _ *http.Request) {
|
||||||
|
w.WriteHeader(http.StatusNotFound)
|
||||||
|
})
|
||||||
|
mux.HandleFunc("DELETE /api/v1/namespaces/poimen/persistentvolumeclaims/{name}", func(w http.ResponseWriter, _ *http.Request) {
|
||||||
|
w.WriteHeader(http.StatusNotFound)
|
||||||
|
})
|
||||||
|
|
||||||
|
srv := httptest.NewServer(mux)
|
||||||
|
defer srv.Close()
|
||||||
|
|
||||||
|
if err := newTestSM(srv.URL).Delete(context.Background(), id); err != nil {
|
||||||
|
t.Errorf("Delete_NotFound: want nil, got %v", err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// ── TestSandboxReapIdlePods ───────────────────────────────────────────────────
|
||||||
|
|
||||||
|
// pod-a=Succeeded → deleted; pod-b=Running → skipped.
|
||||||
|
func TestSandboxReapIdlePods(t *testing.T) {
|
||||||
|
var deletedA int32
|
||||||
|
|
||||||
|
mux := http.NewServeMux()
|
||||||
|
|
||||||
|
// List all pi-session pods (query string not part of mux path matching)
|
||||||
|
mux.HandleFunc("GET /api/v1/namespaces/poimen/pods", func(w http.ResponseWriter, _ *http.Request) {
|
||||||
|
respondK8sJSON(w, http.StatusOK, map[string]any{
|
||||||
|
"items": []any{
|
||||||
|
map[string]any{
|
||||||
|
"metadata": map[string]any{"name": "pod-a"},
|
||||||
|
"status": map[string]any{"phase": "Succeeded"},
|
||||||
|
},
|
||||||
|
map[string]any{
|
||||||
|
"metadata": map[string]any{"name": "pod-b"},
|
||||||
|
"status": map[string]any{"phase": "Running"},
|
||||||
|
},
|
||||||
|
},
|
||||||
|
})
|
||||||
|
})
|
||||||
|
|
||||||
|
mux.HandleFunc("DELETE /api/v1/namespaces/poimen/pods/{name}", func(w http.ResponseWriter, r *http.Request) {
|
||||||
|
name := r.PathValue("name")
|
||||||
|
if name == "pod-b" {
|
||||||
|
t.Errorf("pod-b (Running) must NOT be deleted")
|
||||||
|
}
|
||||||
|
if name == "pod-a" {
|
||||||
|
atomic.AddInt32(&deletedA, 1)
|
||||||
|
}
|
||||||
|
w.WriteHeader(http.StatusOK)
|
||||||
|
})
|
||||||
|
|
||||||
|
srv := httptest.NewServer(mux)
|
||||||
|
defer srv.Close()
|
||||||
|
|
||||||
|
newTestSM(srv.URL).ReapIdlePods(context.Background())
|
||||||
|
|
||||||
|
if got := atomic.LoadInt32(&deletedA); got != 1 {
|
||||||
|
t.Errorf("pod-a delete count: want 1, got %d", got)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// ── TestSandboxNameHelpers ────────────────────────────────────────────────────
|
||||||
|
|
||||||
|
func TestSandboxNameHelpers(t *testing.T) {
|
||||||
|
cases := []struct{ fn, want, got string }{
|
||||||
|
{"pvcName", "workspace-session-abc", pvcName("abc")},
|
||||||
|
{"podName", "pi-session-abc", podName("abc")},
|
||||||
|
{"svcName", "pi-session-abc", svcName("abc")},
|
||||||
|
{"svcDNS", "pi-session-abc.poimen.svc.cluster.local:8787", svcDNS("abc")},
|
||||||
|
}
|
||||||
|
for _, c := range cases {
|
||||||
|
if c.got != c.want {
|
||||||
|
t.Errorf("%s: want %q, got %q", c.fn, c.want, c.got)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// ── TestSandboxIsNotFound ─────────────────────────────────────────────────────
|
||||||
|
|
||||||
|
func TestSandboxIsNotFound(t *testing.T) {
|
||||||
|
if !isNotFound(&k8sNotFound{}) {
|
||||||
|
t.Error("isNotFound(&k8sNotFound{}) must be true")
|
||||||
|
}
|
||||||
|
if isNotFound(fmt.Errorf("other")) {
|
||||||
|
t.Error("isNotFound(fmt.Errorf) must be false")
|
||||||
|
}
|
||||||
|
if isNotFound(nil) {
|
||||||
|
t.Error("isNotFound(nil) must be false")
|
||||||
|
}
|
||||||
|
}
|
||||||
+23
-45
@@ -4,7 +4,6 @@ import (
|
|||||||
"encoding/json"
|
"encoding/json"
|
||||||
"fmt"
|
"fmt"
|
||||||
"net/http"
|
"net/http"
|
||||||
"strings"
|
|
||||||
"time"
|
"time"
|
||||||
)
|
)
|
||||||
|
|
||||||
@@ -33,26 +32,29 @@ func (p *PoimenHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
|
|||||||
p.handleExecute(w, r)
|
p.handleExecute(w, r)
|
||||||
case path == "/v1/poimen/executions" && r.Method == http.MethodGet:
|
case path == "/v1/poimen/executions" && r.Method == http.MethodGet:
|
||||||
p.handleListExecutions(w, r)
|
p.handleListExecutions(w, r)
|
||||||
case strings.HasPrefix(path, "/v1/poimen/executions/") && strings.HasSuffix(path, "/stream") && r.Method == http.MethodGet:
|
case path == "/v1/poimen/executions/stream" && r.Method == http.MethodGet:
|
||||||
p.handleStream(w, r)
|
p.handleStream(w, r)
|
||||||
case strings.HasPrefix(path, "/v1/poimen/executions/") && strings.HasSuffix(path, "/signal") && r.Method == http.MethodPost:
|
case path == "/v1/poimen/executions/signal" && r.Method == http.MethodPost:
|
||||||
p.handleSignal(w, r)
|
p.handleSignal(w, r)
|
||||||
case strings.HasPrefix(path, "/v1/poimen/executions/") && strings.HasSuffix(path, "/history") && r.Method == http.MethodGet:
|
case path == "/v1/poimen/executions/history" && r.Method == http.MethodGet:
|
||||||
p.handleHistory(w, r)
|
p.handleHistory(w, r)
|
||||||
case strings.HasPrefix(path, "/v1/poimen/executions/") && r.Method == http.MethodDelete:
|
case path == "/v1/poimen/executions/detail" && r.Method == http.MethodDelete:
|
||||||
p.handleTerminate(w, r)
|
p.handleTerminate(w, r)
|
||||||
case strings.HasPrefix(path, "/v1/poimen/executions/") && r.Method == http.MethodGet:
|
case path == "/v1/poimen/executions/detail" && r.Method == http.MethodGet:
|
||||||
p.handleDescribe(w, r)
|
p.handleDescribe(w, r)
|
||||||
default:
|
default:
|
||||||
writeJSON(w, http.StatusNotFound, map[string]string{"error": "endpoint not found"})
|
writeJSON(w, http.StatusNotFound, map[string]string{"error": "endpoint not found"})
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
// extractWorkflowID extracts the workflow ID from paths like /v1/poimen/executions/{id}/...
|
// getWorkflowID reads the workflow ID from query param (preferred) or path.
|
||||||
func extractWorkflowID(path string) string {
|
// The ServiceAdapter sets X-Upstream-Path to a template like
|
||||||
parts := strings.Split(strings.TrimPrefix(path, "/v1/poimen/executions/"), "/")
|
// /v1/poimen/executions/{id}/stream — the {id} is literal, not substituted.
|
||||||
if len(parts) > 0 && parts[0] != "" {
|
// So workflow_id must come from a query parameter.
|
||||||
return parts[0]
|
func getWorkflowID(r *http.Request) string {
|
||||||
|
// Query param is the primary source
|
||||||
|
if id := r.URL.Query().Get("workflow_id"); id != "" {
|
||||||
|
return id
|
||||||
}
|
}
|
||||||
return ""
|
return ""
|
||||||
}
|
}
|
||||||
@@ -176,13 +178,9 @@ func (p *PoimenHandler) handleDescribe(w http.ResponseWriter, r *http.Request) {
|
|||||||
return
|
return
|
||||||
}
|
}
|
||||||
|
|
||||||
path := r.Header.Get("X-Upstream-Path")
|
workflowID := getWorkflowID(r)
|
||||||
if path == "" {
|
|
||||||
path = r.URL.Path
|
|
||||||
}
|
|
||||||
workflowID := extractWorkflowID(path)
|
|
||||||
if workflowID == "" {
|
if workflowID == "" {
|
||||||
writeJSON(w, http.StatusBadRequest, map[string]string{"error": "workflow_id required in path"})
|
writeJSON(w, http.StatusBadRequest, map[string]string{"error": "workflow_id query param required"})
|
||||||
return
|
return
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -209,13 +207,9 @@ func (p *PoimenHandler) handleHistory(w http.ResponseWriter, r *http.Request) {
|
|||||||
return
|
return
|
||||||
}
|
}
|
||||||
|
|
||||||
path := r.Header.Get("X-Upstream-Path")
|
workflowID := getWorkflowID(r)
|
||||||
if path == "" {
|
|
||||||
path = r.URL.Path
|
|
||||||
}
|
|
||||||
workflowID := extractWorkflowID(path)
|
|
||||||
if workflowID == "" {
|
if workflowID == "" {
|
||||||
writeJSON(w, http.StatusBadRequest, map[string]string{"error": "workflow_id required in path"})
|
writeJSON(w, http.StatusBadRequest, map[string]string{"error": "workflow_id query param required"})
|
||||||
return
|
return
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -248,15 +242,9 @@ func (p *PoimenHandler) handleSignal(w http.ResponseWriter, r *http.Request) {
|
|||||||
return
|
return
|
||||||
}
|
}
|
||||||
|
|
||||||
path := r.Header.Get("X-Upstream-Path")
|
workflowID := getWorkflowID(r)
|
||||||
if path == "" {
|
|
||||||
path = r.URL.Path
|
|
||||||
}
|
|
||||||
// Remove /signal suffix to get workflow ID
|
|
||||||
cleanPath := strings.TrimSuffix(path, "/signal")
|
|
||||||
workflowID := extractWorkflowID(cleanPath)
|
|
||||||
if workflowID == "" {
|
if workflowID == "" {
|
||||||
writeJSON(w, http.StatusBadRequest, map[string]string{"error": "workflow_id required in path"})
|
writeJSON(w, http.StatusBadRequest, map[string]string{"error": "workflow_id query param required"})
|
||||||
return
|
return
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -290,13 +278,9 @@ func (p *PoimenHandler) handleTerminate(w http.ResponseWriter, r *http.Request)
|
|||||||
return
|
return
|
||||||
}
|
}
|
||||||
|
|
||||||
path := r.Header.Get("X-Upstream-Path")
|
workflowID := getWorkflowID(r)
|
||||||
if path == "" {
|
|
||||||
path = r.URL.Path
|
|
||||||
}
|
|
||||||
workflowID := extractWorkflowID(path)
|
|
||||||
if workflowID == "" {
|
if workflowID == "" {
|
||||||
writeJSON(w, http.StatusBadRequest, map[string]string{"error": "workflow_id required in path"})
|
writeJSON(w, http.StatusBadRequest, map[string]string{"error": "workflow_id query param required"})
|
||||||
return
|
return
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -317,15 +301,9 @@ func (p *PoimenHandler) handleTerminate(w http.ResponseWriter, r *http.Request)
|
|||||||
// --- Stream (delegate to existing SSE handler) ---
|
// --- Stream (delegate to existing SSE handler) ---
|
||||||
|
|
||||||
func (p *PoimenHandler) handleStream(w http.ResponseWriter, r *http.Request) {
|
func (p *PoimenHandler) handleStream(w http.ResponseWriter, r *http.Request) {
|
||||||
path := r.Header.Get("X-Upstream-Path")
|
workflowID := getWorkflowID(r)
|
||||||
if path == "" {
|
|
||||||
path = r.URL.Path
|
|
||||||
}
|
|
||||||
// Remove /stream suffix
|
|
||||||
cleanPath := strings.TrimSuffix(path, "/stream")
|
|
||||||
workflowID := extractWorkflowID(cleanPath)
|
|
||||||
if workflowID == "" {
|
if workflowID == "" {
|
||||||
writeJSON(w, http.StatusBadRequest, map[string]string{"error": "workflow_id required in path"})
|
writeJSON(w, http.StatusBadRequest, map[string]string{"error": "workflow_id query param required"})
|
||||||
return
|
return
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
@@ -11,6 +11,7 @@ import (
|
|||||||
"net/http"
|
"net/http"
|
||||||
"os"
|
"os"
|
||||||
"strings"
|
"strings"
|
||||||
|
"time"
|
||||||
|
|
||||||
"forgejo.riotpiao.com/rock/homelab-frontend/internal/notification"
|
"forgejo.riotpiao.com/rock/homelab-frontend/internal/notification"
|
||||||
)
|
)
|
||||||
@@ -18,8 +19,9 @@ import (
|
|||||||
// ForgejoHandler receives Forgejo webhook payloads and forwards them to Gotify.
|
// ForgejoHandler receives Forgejo webhook payloads and forwards them to Gotify.
|
||||||
// Forgejo sends a Gitea-compatible JSON payload with X-Gitea-Signature-256 header.
|
// Forgejo sends a Gitea-compatible JSON payload with X-Gitea-Signature-256 header.
|
||||||
type ForgejoHandler struct {
|
type ForgejoHandler struct {
|
||||||
secret string
|
secret string
|
||||||
gotify *notification.GotifyClient
|
gotify *notification.GotifyClient
|
||||||
|
promoter *ImagePromoter
|
||||||
}
|
}
|
||||||
|
|
||||||
// NewForgejoHandler creates a handler from environment variables.
|
// NewForgejoHandler creates a handler from environment variables.
|
||||||
@@ -35,8 +37,9 @@ func NewForgejoHandler() *ForgejoHandler {
|
|||||||
}
|
}
|
||||||
|
|
||||||
return &ForgejoHandler{
|
return &ForgejoHandler{
|
||||||
secret: os.Getenv("FORGEJO_WEBHOOK_SECRET"),
|
secret: os.Getenv("FORGEJO_WEBHOOK_SECRET"),
|
||||||
gotify: client,
|
gotify: client,
|
||||||
|
promoter: NewImagePromoter(),
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -65,17 +68,23 @@ func (h *ForgejoHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
event := r.Header.Get("X-Gitea-Event")
|
||||||
|
if event == "" {
|
||||||
|
event = r.Header.Get("X-GitHub-Event")
|
||||||
|
}
|
||||||
|
|
||||||
|
// package events trigger image promotion; handled independently of Gotify
|
||||||
|
if event == "package" {
|
||||||
|
h.handlePackagePromotion(w, r, body)
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
if h.gotify == nil {
|
if h.gotify == nil {
|
||||||
log.Printf("forgejo webhook received but Gotify not configured (GOTIFY_URL/GOTIFY_APP_TOKEN missing)")
|
log.Printf("forgejo webhook received but Gotify not configured (GOTIFY_URL/GOTIFY_APP_TOKEN missing)")
|
||||||
w.WriteHeader(http.StatusOK)
|
w.WriteHeader(http.StatusOK)
|
||||||
return
|
return
|
||||||
}
|
}
|
||||||
|
|
||||||
event := r.Header.Get("X-Gitea-Event")
|
|
||||||
if event == "" {
|
|
||||||
event = r.Header.Get("X-GitHub-Event")
|
|
||||||
}
|
|
||||||
|
|
||||||
title, message, priority := h.formatMessage(event, body)
|
title, message, priority := h.formatMessage(event, body)
|
||||||
if title == "" {
|
if title == "" {
|
||||||
// Unhandled event type — ack and ignore
|
// Unhandled event type — ack and ignore
|
||||||
@@ -99,6 +108,81 @@ func (h *ForgejoHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
|
|||||||
w.WriteHeader(http.StatusOK)
|
w.WriteHeader(http.StatusOK)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// handlePackagePromotion processes a Forgejo package webhook event.
|
||||||
|
// It promotes container images by patching ArgoCD Application CRs.
|
||||||
|
// Sends a Gotify notification with the result (success or failure).
|
||||||
|
func (h *ForgejoHandler) handlePackagePromotion(w http.ResponseWriter, r *http.Request, body []byte) {
|
||||||
|
// Generate a short request ID for correlating log lines from concurrent promotions.
|
||||||
|
reqID := fmt.Sprintf("%08x", time.Now().UnixNano()&0xFFFFFFFF)
|
||||||
|
|
||||||
|
var payload map[string]interface{}
|
||||||
|
if err := json.Unmarshal(body, &payload); err != nil {
|
||||||
|
log.Printf("forgejo webhook [%s]: package event: invalid JSON: %v", reqID, err)
|
||||||
|
w.WriteHeader(http.StatusOK)
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
action := fmt.Sprintf("%v", payload["action"])
|
||||||
|
if action != "created" {
|
||||||
|
// Ignore deletions and other actions
|
||||||
|
w.WriteHeader(http.StatusOK)
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
pkg, _ := payload["package"].(map[string]interface{})
|
||||||
|
pkgType := fmt.Sprintf("%v", pkg["type"])
|
||||||
|
if pkgType != "container" {
|
||||||
|
// Only promote container images
|
||||||
|
w.WriteHeader(http.StatusOK)
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
owner := jsonStr(pkg, "owner", "login")
|
||||||
|
name, _ := pkg["name"].(string)
|
||||||
|
version, _ := pkg["version"].(string)
|
||||||
|
|
||||||
|
if owner == "" || name == "" || version == "" {
|
||||||
|
log.Printf("forgejo webhook [%s]: package event: missing owner/name/version in payload", reqID)
|
||||||
|
w.WriteHeader(http.StatusOK)
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
image := fmt.Sprintf("forgejo.riotpiao.com/%s/%s", owner, name)
|
||||||
|
log.Printf("forgejo webhook [%s]: package event action=%s image=%s tag=%s", reqID, action, image, version)
|
||||||
|
|
||||||
|
var title, message string
|
||||||
|
var priority int
|
||||||
|
|
||||||
|
if h.promoter == nil {
|
||||||
|
title = fmt.Sprintf("📦 Image pushed: %s:%s", image, version)
|
||||||
|
message = "Image promotion disabled — REGISTRY_PAT not configured"
|
||||||
|
priority = 3
|
||||||
|
log.Printf("forgejo webhook [%s]: promoter disabled, image=%s tag=%s", reqID, image, version)
|
||||||
|
} else if err := h.promoter.Promote(r.Context(), image, version); err != nil {
|
||||||
|
title = fmt.Sprintf("❌ Promotion failed: %s:%s", image, version)
|
||||||
|
message = err.Error()
|
||||||
|
priority = 8
|
||||||
|
log.Printf("forgejo webhook [%s]: image promotion failed: %v", reqID, err)
|
||||||
|
} else {
|
||||||
|
title = fmt.Sprintf("🚀 Promoted: %s:%s", image, version)
|
||||||
|
message = "ArgoCD Application patched — redeployment triggered"
|
||||||
|
priority = 5
|
||||||
|
log.Printf("forgejo webhook [%s]: promoted image=%s tag=%s", reqID, image, version)
|
||||||
|
}
|
||||||
|
|
||||||
|
if h.gotify != nil && title != "" {
|
||||||
|
if _, err := h.gotify.SendMessage(notification.GotifyMessage{
|
||||||
|
Title: title,
|
||||||
|
Message: message,
|
||||||
|
Priority: priority,
|
||||||
|
}); err != nil {
|
||||||
|
log.Printf("forgejo webhook: failed to send gotify notification: %v", err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
w.WriteHeader(http.StatusOK)
|
||||||
|
}
|
||||||
|
|
||||||
// verifySignature checks X-Gitea-Signature-256: sha256=<hex>
|
// verifySignature checks X-Gitea-Signature-256: sha256=<hex>
|
||||||
func (h *ForgejoHandler) verifySignature(body []byte, sig string) bool {
|
func (h *ForgejoHandler) verifySignature(body []byte, sig string) bool {
|
||||||
sig = strings.TrimPrefix(sig, "sha256=")
|
sig = strings.TrimPrefix(sig, "sha256=")
|
||||||
|
|||||||
@@ -0,0 +1,424 @@
|
|||||||
|
// Package webhook — ImagePromoter patches ArgoCD Application CRs when a
|
||||||
|
// Forgejo container package webhook fires. It replaces ArgoCD Image Updater
|
||||||
|
// with a webhook-driven flow:
|
||||||
|
//
|
||||||
|
// 1. Forgejo fires X-Gitea-Event: package (action: created, type: container)
|
||||||
|
// 2. ForgejoHandler calls ImagePromoter.Promote(image, tag)
|
||||||
|
// 3. Promoter fetches the digest from Forgejo registry v2 API
|
||||||
|
// 4. Promoter patches the ArgoCD Application CR via the k8s API server
|
||||||
|
// using the in-cluster ServiceAccount token.
|
||||||
|
//
|
||||||
|
// No external dependencies — raw HTTP only.
|
||||||
|
package webhook
|
||||||
|
|
||||||
|
import (
|
||||||
|
"bytes"
|
||||||
|
"context"
|
||||||
|
"crypto/tls"
|
||||||
|
"crypto/x509"
|
||||||
|
"encoding/json"
|
||||||
|
"fmt"
|
||||||
|
"io"
|
||||||
|
"log"
|
||||||
|
"net/http"
|
||||||
|
"os"
|
||||||
|
"strings"
|
||||||
|
"time"
|
||||||
|
)
|
||||||
|
|
||||||
|
const (
|
||||||
|
k8sAPIHost = "https://kubernetes.default.svc"
|
||||||
|
argoNamespace = "argocd"
|
||||||
|
saTokenPath = "/var/run/secrets/kubernetes.io/serviceaccount/token"
|
||||||
|
saCACertPath = "/var/run/secrets/kubernetes.io/serviceaccount/ca.crt"
|
||||||
|
defaultRegBase = "http://forgejo-gitea-http.cicd.svc.cluster.local:3000"
|
||||||
|
defaultRegUser = "rock"
|
||||||
|
promoterTimeout = 20 * time.Second
|
||||||
|
)
|
||||||
|
|
||||||
|
// promotionEntry maps a fully-qualified container image to an ArgoCD Application.
|
||||||
|
type promotionEntry struct {
|
||||||
|
Image string // e.g. "forgejo.riotpiao.com/rock/api-gateway"
|
||||||
|
Application string // ArgoCD Application name in namespace argocd
|
||||||
|
SourceIndex int // index into spec.sources[]; -1 = spec.source (singular, helm)
|
||||||
|
Type string // "kustomize" or "helm"
|
||||||
|
}
|
||||||
|
|
||||||
|
// promotionTable is the static image-to-application mapping.
|
||||||
|
// Order matters for poimen-toolkit: two images share the same source index.
|
||||||
|
var promotionTable = []promotionEntry{
|
||||||
|
{"forgejo.riotpiao.com/rock/api-gateway", "api-gw", 0, "kustomize"},
|
||||||
|
{"forgejo.riotpiao.com/rock/portfolio", "portfolio", 0, "kustomize"},
|
||||||
|
{"forgejo.riotpiao.com/riotpiao-poimen/poimen-workflows", "poimen", 0, "kustomize"},
|
||||||
|
{"forgejo.riotpiao.com/riotpiao-poimen/poimen-memory", "poimen", 1, "kustomize"},
|
||||||
|
// sourceIndex 2 in poimen is a non-container source (chart/config repo) — not promoted.
|
||||||
|
{"forgejo.riotpiao.com/riotpiao-poimen/poimen-frontend", "poimen", 3, "kustomize"},
|
||||||
|
{"forgejo.riotpiao.com/riotpiao-poimen/poimen-toolkit-controller", "poimen-toolkit", 0, "kustomize"},
|
||||||
|
{"forgejo.riotpiao.com/riotpiao-poimen/pi-worker", "poimen-toolkit", 0, "kustomize"},
|
||||||
|
{"forgejo.riotpiao.com/rock/forgejo-runner-golang", "forgejo-runner-golang", -1, "helm"},
|
||||||
|
{"forgejo.riotpiao.com/rock/forgejo-runner-node", "forgejo-runner-node", -1, "helm"},
|
||||||
|
{"forgejo.riotpiao.com/rock/forgejo-runner-rust", "forgejo-runner-rust", -1, "helm"},
|
||||||
|
}
|
||||||
|
|
||||||
|
// ImagePromoter fetches digests and patches ArgoCD Applications.
|
||||||
|
type ImagePromoter struct {
|
||||||
|
registryBase string
|
||||||
|
registryUser string
|
||||||
|
registryPAT string
|
||||||
|
k8sHost string // k8s API server base URL; defaults to k8sAPIHost
|
||||||
|
k8sClient *http.Client // mTLS with SA CA cert
|
||||||
|
regClient *http.Client // plain HTTP to internal Forgejo
|
||||||
|
}
|
||||||
|
|
||||||
|
// NewImagePromoter builds an ImagePromoter from environment and in-cluster SA files.
|
||||||
|
// Returns nil (not an error) when REGISTRY_PAT is empty — promoter is disabled.
|
||||||
|
func NewImagePromoter() *ImagePromoter {
|
||||||
|
pat := os.Getenv("REGISTRY_PAT")
|
||||||
|
if pat == "" {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
regBase := os.Getenv("REGISTRY_BASE")
|
||||||
|
if regBase == "" {
|
||||||
|
regBase = defaultRegBase
|
||||||
|
}
|
||||||
|
regUser := os.Getenv("REGISTRY_USER")
|
||||||
|
if regUser == "" {
|
||||||
|
regUser = defaultRegUser
|
||||||
|
}
|
||||||
|
|
||||||
|
caCert, _ := os.ReadFile(saCACertPath)
|
||||||
|
return &ImagePromoter{
|
||||||
|
registryBase: regBase,
|
||||||
|
registryUser: regUser,
|
||||||
|
registryPAT: pat,
|
||||||
|
k8sClient: buildK8sClient(caCert),
|
||||||
|
regClient: &http.Client{Timeout: promoterTimeout},
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// buildK8sClient creates an HTTP client that trusts the in-cluster CA cert.
|
||||||
|
func buildK8sClient(caCert []byte) *http.Client {
|
||||||
|
transport := &http.Transport{}
|
||||||
|
if len(caCert) > 0 {
|
||||||
|
pool := x509.NewCertPool()
|
||||||
|
pool.AppendCertsFromPEM(caCert)
|
||||||
|
transport.TLSClientConfig = &tls.Config{RootCAs: pool}
|
||||||
|
} else {
|
||||||
|
log.Printf("image-promoter: WARNING SA CA cert not found at %s — using InsecureSkipVerify (dev mode only)", saCACertPath)
|
||||||
|
transport.TLSClientConfig = &tls.Config{InsecureSkipVerify: true} //nolint:gosec
|
||||||
|
}
|
||||||
|
return &http.Client{Timeout: promoterTimeout, Transport: transport}
|
||||||
|
}
|
||||||
|
|
||||||
|
// saToken reads the current in-cluster ServiceAccount token.
|
||||||
|
// Re-reads on every call to handle token rotation.
|
||||||
|
func saToken() string {
|
||||||
|
b, err := os.ReadFile(saTokenPath)
|
||||||
|
if err != nil {
|
||||||
|
log.Printf("image-promoter: WARNING failed to read SA token %s: %v", saTokenPath, err)
|
||||||
|
return ""
|
||||||
|
}
|
||||||
|
return strings.TrimSpace(string(b))
|
||||||
|
}
|
||||||
|
|
||||||
|
// Promote fetches the digest for image:tag and patches the corresponding
|
||||||
|
// ArgoCD Application CR. Returns an error if the image is unknown.
|
||||||
|
// Retries up to 3 times on k8s 409 Conflict (optimistic concurrency for apps
|
||||||
|
// like poimen-toolkit that share a source index across multiple images).
|
||||||
|
func (p *ImagePromoter) Promote(ctx context.Context, image, tag string) error {
|
||||||
|
entry := lookupEntry(image)
|
||||||
|
if entry == nil {
|
||||||
|
return fmt.Errorf("image %q not in promotion table — skipping", image)
|
||||||
|
}
|
||||||
|
|
||||||
|
// Parse owner/name from "forgejo.riotpiao.com/owner/name"
|
||||||
|
withoutHost := strings.TrimPrefix(image, "forgejo.riotpiao.com/")
|
||||||
|
parts := strings.SplitN(withoutHost, "/", 2)
|
||||||
|
if len(parts) != 2 {
|
||||||
|
return fmt.Errorf("cannot parse owner/name from image %q", image)
|
||||||
|
}
|
||||||
|
owner, name := parts[0], parts[1]
|
||||||
|
|
||||||
|
digest, err := p.fetchDigest(ctx, owner, name, tag)
|
||||||
|
if err != nil {
|
||||||
|
return fmt.Errorf("fetch digest for %s:%s: %w", image, tag, err)
|
||||||
|
}
|
||||||
|
|
||||||
|
log.Printf("image-promoter: digest %s → %s application=%s sourceIndex=%d",
|
||||||
|
image, digest, entry.Application, entry.SourceIndex)
|
||||||
|
|
||||||
|
// Retry on 409 Conflict: concurrent promotions for the same (app, sourceIndex)
|
||||||
|
// (e.g. poimen-toolkit-controller and pi-worker) can race on the same source slot.
|
||||||
|
for attempt := 1; attempt <= 3; attempt++ {
|
||||||
|
err = p.patchApplication(ctx, *entry, image, digest)
|
||||||
|
if err == nil {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
if !strings.Contains(err.Error(), "409") || attempt == 3 {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
log.Printf("image-promoter: 409 conflict on PATCH %s (attempt %d/3), retrying",
|
||||||
|
entry.Application, attempt)
|
||||||
|
select {
|
||||||
|
case <-time.After(time.Duration(attempt) * 150 * time.Millisecond):
|
||||||
|
case <-ctx.Done():
|
||||||
|
return ctx.Err()
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
|
||||||
|
// lookupEntry returns the promotion entry for the given image, or nil if not found.
|
||||||
|
func lookupEntry(image string) *promotionEntry {
|
||||||
|
for i := range promotionTable {
|
||||||
|
if promotionTable[i].Image == image {
|
||||||
|
return &promotionTable[i]
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// fetchDigest queries the Forgejo OCI registry v2 API for the manifest digest.
|
||||||
|
// Uses HTTP Basic auth with the configured PAT.
|
||||||
|
func (p *ImagePromoter) fetchDigest(ctx context.Context, owner, name, tag string) (string, error) {
|
||||||
|
url := fmt.Sprintf("%s/v2/%s/%s/manifests/%s", p.registryBase, owner, name, tag)
|
||||||
|
|
||||||
|
req, err := http.NewRequestWithContext(ctx, http.MethodHead, url, nil)
|
||||||
|
if err != nil {
|
||||||
|
return "", fmt.Errorf("build request: %w", err)
|
||||||
|
}
|
||||||
|
req.SetBasicAuth(p.registryUser, p.registryPAT)
|
||||||
|
req.Header.Add("Accept", "application/vnd.docker.distribution.manifest.v2+json")
|
||||||
|
req.Header.Add("Accept", "application/vnd.oci.image.manifest.v1+json")
|
||||||
|
|
||||||
|
resp, err := p.regClient.Do(req)
|
||||||
|
if err != nil {
|
||||||
|
return "", fmt.Errorf("http: %w", err)
|
||||||
|
}
|
||||||
|
defer resp.Body.Close()
|
||||||
|
|
||||||
|
if resp.StatusCode != http.StatusOK {
|
||||||
|
return "", fmt.Errorf("registry returned %d for %s/%s:%s", resp.StatusCode, owner, name, tag)
|
||||||
|
}
|
||||||
|
|
||||||
|
digest := resp.Header.Get("Docker-Content-Digest")
|
||||||
|
if digest == "" {
|
||||||
|
return "", fmt.Errorf("Docker-Content-Digest header missing from registry response")
|
||||||
|
}
|
||||||
|
return digest, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// patchApplication performs a read-modify-write on the ArgoCD Application CR.
|
||||||
|
// It reads the current spec to preserve existing image entries, updates the
|
||||||
|
// relevant entry, then sends a JSON Patch to the k8s API server.
|
||||||
|
func (p *ImagePromoter) patchApplication(ctx context.Context, entry promotionEntry, image, digest string) error {
|
||||||
|
app, err := p.getApplication(ctx, entry.Application)
|
||||||
|
if err != nil {
|
||||||
|
return fmt.Errorf("get application %s: %w", entry.Application, err)
|
||||||
|
}
|
||||||
|
|
||||||
|
var patch []map[string]any
|
||||||
|
|
||||||
|
switch entry.Type {
|
||||||
|
case "kustomize":
|
||||||
|
images := getKustomizeImages(app, entry.SourceIndex)
|
||||||
|
newImages := updateImageEntry(images, image, digest)
|
||||||
|
// Patch the full kustomize block (not just /kustomize/images) so the
|
||||||
|
// op succeeds even when the kustomize key is absent from the source.
|
||||||
|
// These apps only use kustomize overrides for image pinning, so
|
||||||
|
// replacing the whole block is safe.
|
||||||
|
patch = []map[string]any{{
|
||||||
|
"op": "add",
|
||||||
|
"path": fmt.Sprintf("/spec/sources/%d/kustomize", entry.SourceIndex),
|
||||||
|
"value": map[string]any{"images": newImages},
|
||||||
|
}}
|
||||||
|
|
||||||
|
case "helm":
|
||||||
|
// Helm runners use spec.source (singular). Patch helm parameters
|
||||||
|
// to set runner.image.repository and runner.image.tag with digest.
|
||||||
|
existing := getHelmParameters(app)
|
||||||
|
updates := []map[string]string{
|
||||||
|
{"name": "runner.image.repository", "value": image},
|
||||||
|
{"name": "runner.image.tag", "value": "latest@" + digest},
|
||||||
|
}
|
||||||
|
merged := mergeHelmParams(existing, updates)
|
||||||
|
patch = []map[string]any{
|
||||||
|
// "add" works whether parameters key exists or not (replace fails on absent key)
|
||||||
|
{"op": "add", "path": "/spec/source/helm/parameters", "value": merged},
|
||||||
|
}
|
||||||
|
|
||||||
|
default:
|
||||||
|
return fmt.Errorf("unknown promotion type %q for %s", entry.Type, entry.Application)
|
||||||
|
}
|
||||||
|
|
||||||
|
return p.applyJSONPatch(ctx, entry.Application, patch)
|
||||||
|
}
|
||||||
|
|
||||||
|
// k8sBase returns the configured k8s host or the default in-cluster address.
|
||||||
|
func (p *ImagePromoter) k8sBase() string {
|
||||||
|
if p.k8sHost != "" {
|
||||||
|
return p.k8sHost
|
||||||
|
}
|
||||||
|
return k8sAPIHost
|
||||||
|
}
|
||||||
|
|
||||||
|
// getApplication fetches the current ArgoCD Application CR from the k8s API.
|
||||||
|
func (p *ImagePromoter) getApplication(ctx context.Context, name string) (map[string]any, error) {
|
||||||
|
url := fmt.Sprintf("%s/apis/argoproj.io/v1alpha1/namespaces/%s/applications/%s",
|
||||||
|
p.k8sBase(), argoNamespace, name)
|
||||||
|
|
||||||
|
req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil)
|
||||||
|
if err != nil {
|
||||||
|
return nil, fmt.Errorf("build request: %w", err)
|
||||||
|
}
|
||||||
|
req.Header.Set("Authorization", "Bearer "+saToken())
|
||||||
|
|
||||||
|
resp, err := p.k8sClient.Do(req)
|
||||||
|
if err != nil {
|
||||||
|
return nil, fmt.Errorf("http: %w", err)
|
||||||
|
}
|
||||||
|
defer resp.Body.Close()
|
||||||
|
|
||||||
|
if resp.StatusCode != http.StatusOK {
|
||||||
|
body, _ := io.ReadAll(io.LimitReader(resp.Body, 512))
|
||||||
|
return nil, fmt.Errorf("k8s GET returned %d: %s", resp.StatusCode, body)
|
||||||
|
}
|
||||||
|
|
||||||
|
var app map[string]any
|
||||||
|
if err := json.NewDecoder(resp.Body).Decode(&app); err != nil {
|
||||||
|
return nil, fmt.Errorf("decode: %w", err)
|
||||||
|
}
|
||||||
|
return app, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// applyJSONPatch sends RFC 6902 JSON Patch to the ArgoCD Application.
|
||||||
|
func (p *ImagePromoter) applyJSONPatch(ctx context.Context, appName string, patch []map[string]any) error {
|
||||||
|
patchBytes, err := json.Marshal(patch)
|
||||||
|
if err != nil {
|
||||||
|
return fmt.Errorf("marshal patch: %w", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
url := fmt.Sprintf("%s/apis/argoproj.io/v1alpha1/namespaces/%s/applications/%s",
|
||||||
|
p.k8sBase(), argoNamespace, appName)
|
||||||
|
|
||||||
|
req, err := http.NewRequestWithContext(ctx, http.MethodPatch, url, bytes.NewReader(patchBytes))
|
||||||
|
if err != nil {
|
||||||
|
return fmt.Errorf("build request: %w", err)
|
||||||
|
}
|
||||||
|
req.Header.Set("Content-Type", "application/json-patch+json")
|
||||||
|
req.Header.Set("Authorization", "Bearer "+saToken())
|
||||||
|
|
||||||
|
resp, err := p.k8sClient.Do(req)
|
||||||
|
if err != nil {
|
||||||
|
return fmt.Errorf("http: %w", err)
|
||||||
|
}
|
||||||
|
defer resp.Body.Close()
|
||||||
|
|
||||||
|
if resp.StatusCode < 200 || resp.StatusCode >= 300 {
|
||||||
|
body, _ := io.ReadAll(io.LimitReader(resp.Body, 512))
|
||||||
|
return fmt.Errorf("k8s PATCH returned %d: %s", resp.StatusCode, body)
|
||||||
|
}
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// ── spec helpers ─────────────────────────────────────────────────────────────
|
||||||
|
|
||||||
|
// getKustomizeImages extracts spec.sources[N].kustomize.images.
|
||||||
|
// Returns nil if the path does not exist.
|
||||||
|
func getKustomizeImages(app map[string]any, sourceIndex int) []string {
|
||||||
|
spec, _ := app["spec"].(map[string]any)
|
||||||
|
sources, _ := spec["sources"].([]any)
|
||||||
|
if sourceIndex < 0 || sourceIndex >= len(sources) {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
src, _ := sources[sourceIndex].(map[string]any)
|
||||||
|
kust, _ := src["kustomize"].(map[string]any)
|
||||||
|
return toStringSlice(kust["images"])
|
||||||
|
}
|
||||||
|
|
||||||
|
// getHelmParameters extracts spec.source.helm.parameters.
|
||||||
|
func getHelmParameters(app map[string]any) []map[string]string {
|
||||||
|
spec, _ := app["spec"].(map[string]any)
|
||||||
|
source, _ := spec["source"].(map[string]any)
|
||||||
|
helm, _ := source["helm"].(map[string]any)
|
||||||
|
raw, _ := helm["parameters"].([]any)
|
||||||
|
|
||||||
|
params := make([]map[string]string, 0, len(raw))
|
||||||
|
for _, item := range raw {
|
||||||
|
m, ok := item.(map[string]any)
|
||||||
|
if !ok {
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
p := map[string]string{}
|
||||||
|
if n, ok := m["name"].(string); ok {
|
||||||
|
p["name"] = n
|
||||||
|
}
|
||||||
|
if v, ok := m["value"].(string); ok {
|
||||||
|
p["value"] = v
|
||||||
|
}
|
||||||
|
params = append(params, p)
|
||||||
|
}
|
||||||
|
return params
|
||||||
|
}
|
||||||
|
|
||||||
|
// updateImageEntry replaces the entry for image in the images slice, or appends
|
||||||
|
// if not present. The new entry format is "{image}@{digest}".
|
||||||
|
// Matching is by image name (prefix before ':' or '@').
|
||||||
|
func updateImageEntry(images []string, image, digest string) []string {
|
||||||
|
newEntry := image + "@" + digest
|
||||||
|
for i, img := range images {
|
||||||
|
// Strip tag/digest suffix to get the bare name
|
||||||
|
bare := img
|
||||||
|
if idx := strings.IndexAny(img, ":@"); idx >= 0 {
|
||||||
|
bare = img[:idx]
|
||||||
|
}
|
||||||
|
if bare == image {
|
||||||
|
result := make([]string, len(images))
|
||||||
|
copy(result, images)
|
||||||
|
result[i] = newEntry
|
||||||
|
return result
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return append(images, newEntry)
|
||||||
|
}
|
||||||
|
|
||||||
|
// mergeHelmParams merges update params into existing, preserving unrelated params.
|
||||||
|
func mergeHelmParams(existing, updates []map[string]string) []map[string]string {
|
||||||
|
idx := make(map[string]string, len(updates))
|
||||||
|
for _, u := range updates {
|
||||||
|
idx[u["name"]] = u["value"]
|
||||||
|
}
|
||||||
|
|
||||||
|
result := make([]map[string]string, 0, len(existing)+len(updates))
|
||||||
|
replaced := make(map[string]bool)
|
||||||
|
for _, p := range existing {
|
||||||
|
if val, ok := idx[p["name"]]; ok {
|
||||||
|
result = append(result, map[string]string{"name": p["name"], "value": val})
|
||||||
|
replaced[p["name"]] = true
|
||||||
|
} else {
|
||||||
|
result = append(result, p)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
for _, u := range updates {
|
||||||
|
if !replaced[u["name"]] {
|
||||||
|
result = append(result, u)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return result
|
||||||
|
}
|
||||||
|
|
||||||
|
// toStringSlice converts []any to []string, skipping non-strings.
|
||||||
|
func toStringSlice(v any) []string {
|
||||||
|
arr, ok := v.([]any)
|
||||||
|
if !ok {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
result := make([]string, 0, len(arr))
|
||||||
|
for _, item := range arr {
|
||||||
|
if s, ok := item.(string); ok {
|
||||||
|
result = append(result, s)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return result
|
||||||
|
}
|
||||||
@@ -0,0 +1,408 @@
|
|||||||
|
package webhook
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"encoding/json"
|
||||||
|
"io"
|
||||||
|
"net/http"
|
||||||
|
"net/http/httptest"
|
||||||
|
"strings"
|
||||||
|
"sync/atomic"
|
||||||
|
"testing"
|
||||||
|
"time"
|
||||||
|
)
|
||||||
|
|
||||||
|
// ── unit: updateImageEntry ────────────────────────────────────────────────────
|
||||||
|
|
||||||
|
func TestUpdateImageEntry(t *testing.T) {
|
||||||
|
t.Run("replaces existing by bare name strips tag", func(t *testing.T) {
|
||||||
|
images := []string{"forgejo.riotpiao.com/rock/api-gateway:old"}
|
||||||
|
got := updateImageEntry(images, "forgejo.riotpiao.com/rock/api-gateway", "sha256:new")
|
||||||
|
if len(got) != 1 || got[0] != "forgejo.riotpiao.com/rock/api-gateway@sha256:new" {
|
||||||
|
t.Fatalf("got %v", got)
|
||||||
|
}
|
||||||
|
})
|
||||||
|
|
||||||
|
t.Run("replaces existing by bare name strips digest", func(t *testing.T) {
|
||||||
|
images := []string{"forgejo.riotpiao.com/rock/api-gateway@sha256:old"}
|
||||||
|
got := updateImageEntry(images, "forgejo.riotpiao.com/rock/api-gateway", "sha256:new")
|
||||||
|
if len(got) != 1 || got[0] != "forgejo.riotpiao.com/rock/api-gateway@sha256:new" {
|
||||||
|
t.Fatalf("got %v", got)
|
||||||
|
}
|
||||||
|
})
|
||||||
|
|
||||||
|
t.Run("appends new entry when image not in list", func(t *testing.T) {
|
||||||
|
images := []string{"forgejo.riotpiao.com/rock/other:latest"}
|
||||||
|
got := updateImageEntry(images, "forgejo.riotpiao.com/rock/api-gateway", "sha256:new")
|
||||||
|
if len(got) != 2 || got[1] != "forgejo.riotpiao.com/rock/api-gateway@sha256:new" {
|
||||||
|
t.Fatalf("got %v", got)
|
||||||
|
}
|
||||||
|
})
|
||||||
|
|
||||||
|
t.Run("preserves other entries when updating one", func(t *testing.T) {
|
||||||
|
images := []string{
|
||||||
|
"forgejo.riotpiao.com/rock/other:1.0",
|
||||||
|
"forgejo.riotpiao.com/rock/api-gateway:old",
|
||||||
|
"forgejo.riotpiao.com/rock/third:2.0",
|
||||||
|
}
|
||||||
|
got := updateImageEntry(images, "forgejo.riotpiao.com/rock/api-gateway", "sha256:new")
|
||||||
|
if len(got) != 3 {
|
||||||
|
t.Fatalf("expected 3 entries got %d: %v", len(got), got)
|
||||||
|
}
|
||||||
|
if got[0] != "forgejo.riotpiao.com/rock/other:1.0" {
|
||||||
|
t.Errorf("first entry changed: %s", got[0])
|
||||||
|
}
|
||||||
|
if got[1] != "forgejo.riotpiao.com/rock/api-gateway@sha256:new" {
|
||||||
|
t.Errorf("second entry wrong: %s", got[1])
|
||||||
|
}
|
||||||
|
if got[2] != "forgejo.riotpiao.com/rock/third:2.0" {
|
||||||
|
t.Errorf("third entry changed: %s", got[2])
|
||||||
|
}
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
// ── unit: mergeHelmParams ─────────────────────────────────────────────────────
|
||||||
|
|
||||||
|
func TestMergeHelmParams(t *testing.T) {
|
||||||
|
t.Run("update existing param preserve unrelated", func(t *testing.T) {
|
||||||
|
existing := []map[string]string{
|
||||||
|
{"name": "runner.image.tag", "value": "old"},
|
||||||
|
{"name": "replicas", "value": "2"},
|
||||||
|
}
|
||||||
|
updates := []map[string]string{
|
||||||
|
{"name": "runner.image.tag", "value": "new"},
|
||||||
|
}
|
||||||
|
got := mergeHelmParams(existing, updates)
|
||||||
|
find := func(name string) string {
|
||||||
|
for _, p := range got {
|
||||||
|
if p["name"] == name {
|
||||||
|
return p["value"]
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return ""
|
||||||
|
}
|
||||||
|
if find("runner.image.tag") != "new" {
|
||||||
|
t.Errorf("tag not updated: %v", got)
|
||||||
|
}
|
||||||
|
if find("replicas") != "2" {
|
||||||
|
t.Errorf("replicas lost: %v", got)
|
||||||
|
}
|
||||||
|
})
|
||||||
|
|
||||||
|
t.Run("append new param not in existing list", func(t *testing.T) {
|
||||||
|
existing := []map[string]string{
|
||||||
|
{"name": "replicas", "value": "2"},
|
||||||
|
}
|
||||||
|
updates := []map[string]string{
|
||||||
|
{"name": "runner.image.tag", "value": "new"},
|
||||||
|
}
|
||||||
|
got := mergeHelmParams(existing, updates)
|
||||||
|
if len(got) != 2 {
|
||||||
|
t.Fatalf("expected 2 params got %d: %v", len(got), got)
|
||||||
|
}
|
||||||
|
})
|
||||||
|
|
||||||
|
t.Run("empty existing list just returns updates", func(t *testing.T) {
|
||||||
|
updates := []map[string]string{
|
||||||
|
{"name": "runner.image.tag", "value": "new"},
|
||||||
|
}
|
||||||
|
got := mergeHelmParams(nil, updates)
|
||||||
|
if len(got) != 1 || got[0]["name"] != "runner.image.tag" {
|
||||||
|
t.Fatalf("got %v", got)
|
||||||
|
}
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
// ── unit: lookupEntry ─────────────────────────────────────────────────────────
|
||||||
|
|
||||||
|
func TestLookupEntry(t *testing.T) {
|
||||||
|
t.Run("known image returns entry", func(t *testing.T) {
|
||||||
|
e := lookupEntry("forgejo.riotpiao.com/rock/api-gateway")
|
||||||
|
if e == nil {
|
||||||
|
t.Fatal("expected entry, got nil")
|
||||||
|
}
|
||||||
|
if e.Application != "api-gw" {
|
||||||
|
t.Errorf("wrong app: %s", e.Application)
|
||||||
|
}
|
||||||
|
})
|
||||||
|
|
||||||
|
t.Run("unknown image returns nil", func(t *testing.T) {
|
||||||
|
e := lookupEntry("forgejo.riotpiao.com/rock/unknown")
|
||||||
|
if e != nil {
|
||||||
|
t.Fatalf("expected nil, got %+v", e)
|
||||||
|
}
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
// ── unit: toStringSlice ───────────────────────────────────────────────────────
|
||||||
|
|
||||||
|
func TestToStringSlice(t *testing.T) {
|
||||||
|
t.Run("skips non-string elements", func(t *testing.T) {
|
||||||
|
got := toStringSlice([]any{"a", "b", 42})
|
||||||
|
if len(got) != 2 || got[0] != "a" || got[1] != "b" {
|
||||||
|
t.Fatalf("got %v", got)
|
||||||
|
}
|
||||||
|
})
|
||||||
|
|
||||||
|
t.Run("nil input returns nil", func(t *testing.T) {
|
||||||
|
got := toStringSlice(nil)
|
||||||
|
if got != nil {
|
||||||
|
t.Fatalf("expected nil, got %v", got)
|
||||||
|
}
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
// ── helpers ───────────────────────────────────────────────────────────────────
|
||||||
|
|
||||||
|
func newPromoter(regURL, k8sURL string) *ImagePromoter {
|
||||||
|
return &ImagePromoter{
|
||||||
|
registryBase: regURL,
|
||||||
|
registryUser: "testuser",
|
||||||
|
registryPAT: "testpat",
|
||||||
|
k8sHost: k8sURL,
|
||||||
|
k8sClient: &http.Client{Timeout: 5 * time.Second},
|
||||||
|
regClient: &http.Client{Timeout: 5 * time.Second},
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func mustJSON(v any) []byte {
|
||||||
|
b, err := json.Marshal(v)
|
||||||
|
if err != nil {
|
||||||
|
panic(err)
|
||||||
|
}
|
||||||
|
return b
|
||||||
|
}
|
||||||
|
|
||||||
|
// ── integration: fetchDigest ──────────────────────────────────────────────────
|
||||||
|
|
||||||
|
func TestFetchDigest_Success(t *testing.T) {
|
||||||
|
reg := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||||
|
w.Header().Set("Docker-Content-Digest", "sha256:abc123")
|
||||||
|
w.WriteHeader(http.StatusOK)
|
||||||
|
}))
|
||||||
|
defer reg.Close()
|
||||||
|
|
||||||
|
ctx := context.Background()
|
||||||
|
p := newPromoter(reg.URL, "")
|
||||||
|
digest, err := p.fetchDigest(ctx, "rock", "api-gateway", "latest")
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("unexpected error: %v", err)
|
||||||
|
}
|
||||||
|
if digest != "sha256:abc123" {
|
||||||
|
t.Errorf("got digest %q", digest)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestFetchDigest_NotFound(t *testing.T) {
|
||||||
|
reg := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||||
|
w.WriteHeader(http.StatusNotFound)
|
||||||
|
}))
|
||||||
|
defer reg.Close()
|
||||||
|
|
||||||
|
ctx := context.Background()
|
||||||
|
p := newPromoter(reg.URL, "")
|
||||||
|
_, err := p.fetchDigest(ctx, "rock", "api-gateway", "latest")
|
||||||
|
if err == nil {
|
||||||
|
t.Fatal("expected error, got nil")
|
||||||
|
}
|
||||||
|
if !strings.Contains(err.Error(), "404") {
|
||||||
|
t.Errorf("error should contain 404: %v", err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestFetchDigest_MissingDigestHeader(t *testing.T) {
|
||||||
|
reg := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||||
|
w.WriteHeader(http.StatusOK)
|
||||||
|
// No Docker-Content-Digest header
|
||||||
|
}))
|
||||||
|
defer reg.Close()
|
||||||
|
|
||||||
|
ctx := context.Background()
|
||||||
|
p := newPromoter(reg.URL, "")
|
||||||
|
_, err := p.fetchDigest(ctx, "rock", "api-gateway", "latest")
|
||||||
|
if err == nil {
|
||||||
|
t.Fatal("expected error, got nil")
|
||||||
|
}
|
||||||
|
if !strings.Contains(err.Error(), "Docker-Content-Digest") {
|
||||||
|
t.Errorf("error should mention Docker-Content-Digest: %v", err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// ── integration: Promote kustomize ───────────────────────────────────────────
|
||||||
|
|
||||||
|
func TestPromote_KustomizeSuccess(t *testing.T) {
|
||||||
|
reg := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||||
|
w.Header().Set("Docker-Content-Digest", "sha256:newdigest")
|
||||||
|
w.WriteHeader(http.StatusOK)
|
||||||
|
}))
|
||||||
|
defer reg.Close()
|
||||||
|
|
||||||
|
kustomizeApp := map[string]any{
|
||||||
|
"spec": map[string]any{
|
||||||
|
"sources": []any{
|
||||||
|
map[string]any{
|
||||||
|
"kustomize": map[string]any{
|
||||||
|
"images": []any{"forgejo.riotpiao.com/rock/api-gateway:latest"},
|
||||||
|
},
|
||||||
|
},
|
||||||
|
},
|
||||||
|
},
|
||||||
|
}
|
||||||
|
|
||||||
|
var getCalled, patchCalled int32
|
||||||
|
k8s := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||||
|
switch r.Method {
|
||||||
|
case http.MethodGet:
|
||||||
|
atomic.AddInt32(&getCalled, 1)
|
||||||
|
w.Header().Set("Content-Type", "application/json")
|
||||||
|
w.WriteHeader(http.StatusOK)
|
||||||
|
w.Write(mustJSON(kustomizeApp))
|
||||||
|
case http.MethodPatch:
|
||||||
|
atomic.AddInt32(&patchCalled, 1)
|
||||||
|
raw, _ := readAll(r.Body)
|
||||||
|
if !strings.Contains(string(raw), `"/spec/sources/0/kustomize"`) {
|
||||||
|
t.Errorf("patch path missing: %s", raw)
|
||||||
|
}
|
||||||
|
if !strings.Contains(string(raw), `"images"`) {
|
||||||
|
t.Errorf("images key missing from kustomize block: %s", raw)
|
||||||
|
}
|
||||||
|
if !strings.Contains(string(raw), "@sha256:newdigest") {
|
||||||
|
t.Errorf("digest missing from patch body: %s", raw)
|
||||||
|
}
|
||||||
|
w.WriteHeader(http.StatusOK)
|
||||||
|
}
|
||||||
|
}))
|
||||||
|
defer k8s.Close()
|
||||||
|
|
||||||
|
p := newPromoter(reg.URL, k8s.URL)
|
||||||
|
ctx := context.Background()
|
||||||
|
if err := p.Promote(ctx, "forgejo.riotpiao.com/rock/api-gateway", "latest"); err != nil {
|
||||||
|
t.Fatalf("Promote error: %v", err)
|
||||||
|
}
|
||||||
|
if atomic.LoadInt32(&getCalled) == 0 {
|
||||||
|
t.Error("k8s GET not called")
|
||||||
|
}
|
||||||
|
if atomic.LoadInt32(&patchCalled) == 0 {
|
||||||
|
t.Error("k8s PATCH not called")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// ── integration: Promote helm ─────────────────────────────────────────────────
|
||||||
|
|
||||||
|
func TestPromote_HelmSuccess(t *testing.T) {
|
||||||
|
reg := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||||
|
w.Header().Set("Docker-Content-Digest", "sha256:helmdigest")
|
||||||
|
w.WriteHeader(http.StatusOK)
|
||||||
|
}))
|
||||||
|
defer reg.Close()
|
||||||
|
|
||||||
|
helmApp := map[string]any{
|
||||||
|
"spec": map[string]any{
|
||||||
|
"source": map[string]any{
|
||||||
|
"helm": map[string]any{
|
||||||
|
"parameters": []any{
|
||||||
|
map[string]any{"name": "runner.image.tag", "value": "old"},
|
||||||
|
},
|
||||||
|
},
|
||||||
|
},
|
||||||
|
},
|
||||||
|
}
|
||||||
|
|
||||||
|
var patchCalled int32
|
||||||
|
k8s := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||||
|
switch r.Method {
|
||||||
|
case http.MethodGet:
|
||||||
|
w.Header().Set("Content-Type", "application/json")
|
||||||
|
w.WriteHeader(http.StatusOK)
|
||||||
|
w.Write(mustJSON(helmApp))
|
||||||
|
case http.MethodPatch:
|
||||||
|
atomic.AddInt32(&patchCalled, 1)
|
||||||
|
raw, _ := readAll(r.Body)
|
||||||
|
if !strings.Contains(string(raw), "/spec/source/helm/parameters") {
|
||||||
|
t.Errorf("patch path missing: %s", raw)
|
||||||
|
}
|
||||||
|
if !strings.Contains(string(raw), "helmdigest") {
|
||||||
|
t.Errorf("digest missing from patch body: %s", raw)
|
||||||
|
}
|
||||||
|
w.WriteHeader(http.StatusOK)
|
||||||
|
}
|
||||||
|
}))
|
||||||
|
defer k8s.Close()
|
||||||
|
|
||||||
|
p := newPromoter(reg.URL, k8s.URL)
|
||||||
|
ctx := context.Background()
|
||||||
|
if err := p.Promote(ctx, "forgejo.riotpiao.com/rock/forgejo-runner-golang", "latest"); err != nil {
|
||||||
|
t.Fatalf("Promote error: %v", err)
|
||||||
|
}
|
||||||
|
if atomic.LoadInt32(&patchCalled) == 0 {
|
||||||
|
t.Error("k8s PATCH not called")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// ── integration: Promote unknown image ───────────────────────────────────────
|
||||||
|
|
||||||
|
func TestPromote_UnknownImage(t *testing.T) {
|
||||||
|
p := newPromoter("http://unused", "http://unused")
|
||||||
|
ctx := context.TODO()
|
||||||
|
err := p.Promote(ctx, "forgejo.riotpiao.com/rock/unknown-image", "latest")
|
||||||
|
if err == nil {
|
||||||
|
t.Fatal("expected error, got nil")
|
||||||
|
}
|
||||||
|
if !strings.Contains(err.Error(), "not in promotion table") {
|
||||||
|
t.Errorf("error should mention 'not in promotion table': %v", err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// ── integration: Promote conflict retry ──────────────────────────────────────
|
||||||
|
|
||||||
|
func TestPromote_Conflict_Retry(t *testing.T) {
|
||||||
|
reg := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||||
|
w.Header().Set("Docker-Content-Digest", "sha256:retrydigest")
|
||||||
|
w.WriteHeader(http.StatusOK)
|
||||||
|
}))
|
||||||
|
defer reg.Close()
|
||||||
|
|
||||||
|
kustomizeApp := map[string]any{
|
||||||
|
"spec": map[string]any{
|
||||||
|
"sources": []any{
|
||||||
|
map[string]any{
|
||||||
|
"kustomize": map[string]any{
|
||||||
|
"images": []any{"forgejo.riotpiao.com/rock/api-gateway:latest"},
|
||||||
|
},
|
||||||
|
},
|
||||||
|
},
|
||||||
|
},
|
||||||
|
}
|
||||||
|
|
||||||
|
var patchCount int32
|
||||||
|
k8s := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||||
|
switch r.Method {
|
||||||
|
case http.MethodGet:
|
||||||
|
w.Header().Set("Content-Type", "application/json")
|
||||||
|
w.WriteHeader(http.StatusOK)
|
||||||
|
w.Write(mustJSON(kustomizeApp))
|
||||||
|
case http.MethodPatch:
|
||||||
|
n := atomic.AddInt32(&patchCount, 1)
|
||||||
|
if n < 3 {
|
||||||
|
w.WriteHeader(http.StatusConflict) // 409
|
||||||
|
return
|
||||||
|
}
|
||||||
|
w.WriteHeader(http.StatusOK)
|
||||||
|
}
|
||||||
|
}))
|
||||||
|
defer k8s.Close()
|
||||||
|
|
||||||
|
p := newPromoter(reg.URL, k8s.URL)
|
||||||
|
ctx := context.Background()
|
||||||
|
if err := p.Promote(ctx, "forgejo.riotpiao.com/rock/api-gateway", "latest"); err != nil {
|
||||||
|
t.Fatalf("Promote error after retries: %v", err)
|
||||||
|
}
|
||||||
|
if got := atomic.LoadInt32(&patchCount); got != 3 {
|
||||||
|
t.Errorf("expected 3 PATCH calls, got %d", got)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// readAll is a test helper to read request body.
|
||||||
|
func readAll(r io.Reader) ([]byte, error) {
|
||||||
|
return io.ReadAll(r)
|
||||||
|
}
|
||||||
@@ -95,6 +95,14 @@ spec:
|
|||||||
name: gotify-webhook-secret
|
name: gotify-webhook-secret
|
||||||
key: forgejo-webhook-secret
|
key: forgejo-webhook-secret
|
||||||
optional: true
|
optional: true
|
||||||
|
- name: REGISTRY_PAT
|
||||||
|
valueFrom:
|
||||||
|
secretKeyRef:
|
||||||
|
name: forgejo-registry-pat
|
||||||
|
key: pat
|
||||||
|
optional: true
|
||||||
|
- name: REGISTRY_BASE
|
||||||
|
value: "http://forgejo-gitea-http.cicd.svc.cluster.local:3000"
|
||||||
volumeMounts:
|
volumeMounts:
|
||||||
- name: config
|
- name: config
|
||||||
mountPath: /etc/gateway
|
mountPath: /etc/gateway
|
||||||
|
|||||||
@@ -9,6 +9,7 @@ resources:
|
|||||||
- deployment.yaml
|
- deployment.yaml
|
||||||
- network-policy.yaml
|
- network-policy.yaml
|
||||||
- gateway-config-secret.yaml
|
- gateway-config-secret.yaml
|
||||||
|
- rbac.yaml
|
||||||
|
|
||||||
# The deployed image tag lives here and nowhere else. CI publishes
|
# The deployed image tag lives here and nowhere else. CI publishes
|
||||||
# forgejo.riotpiao.com/rock/api-gateway:<commit-sha> and tags it as :latest on main.
|
# forgejo.riotpiao.com/rock/api-gateway:<commit-sha> and tags it as :latest on main.
|
||||||
|
|||||||
+44
-3
@@ -21,11 +21,17 @@ spec:
|
|||||||
ports:
|
ports:
|
||||||
- protocol: TCP
|
- protocol: TCP
|
||||||
port: 8080
|
port: 8080
|
||||||
# Allow from Prometheus scraping (if in monitoring namespace)
|
# Allow from Prometheus scraping (monitoring namespace)
|
||||||
|
# MUST select on kubernetes.io/metadata.name, not a plain `name:` label.
|
||||||
|
# Only ingress-nginx carries a plain `name` label in this cluster; every
|
||||||
|
# other namespace has just kubernetes.io/metadata.name, so `name: monitoring`
|
||||||
|
# matched nothing and silently denied every Prometheus scrape of :8080
|
||||||
|
# (surfaced as context-deadline-exceeded on all 3 gateway replicas while the
|
||||||
|
# ServiceMonitor itself was discovered and healthy).
|
||||||
- from:
|
- from:
|
||||||
- namespaceSelector:
|
- namespaceSelector:
|
||||||
matchLabels:
|
matchLabels:
|
||||||
name: monitoring
|
kubernetes.io/metadata.name: monitoring
|
||||||
ports:
|
ports:
|
||||||
- protocol: TCP
|
- protocol: TCP
|
||||||
port: 8080
|
port: 8080
|
||||||
@@ -122,7 +128,7 @@ spec:
|
|||||||
ports:
|
ports:
|
||||||
- protocol: TCP
|
- protocol: TCP
|
||||||
port: 7233
|
port: 7233
|
||||||
# Allow to Poimen (memory/semantic search)
|
# Allow to Poimen (memory:8080 + graphql:8090 + redis:6379 + fleet-controller:8788)
|
||||||
- to:
|
- to:
|
||||||
- namespaceSelector:
|
- namespaceSelector:
|
||||||
matchLabels:
|
matchLabels:
|
||||||
@@ -130,6 +136,20 @@ spec:
|
|||||||
ports:
|
ports:
|
||||||
- protocol: TCP
|
- protocol: TCP
|
||||||
port: 8080
|
port: 8080
|
||||||
|
- protocol: TCP
|
||||||
|
port: 8090
|
||||||
|
- protocol: TCP
|
||||||
|
port: 6379
|
||||||
|
- protocol: TCP
|
||||||
|
port: 8788
|
||||||
|
# Allow to Gotify (push notifications)
|
||||||
|
- to:
|
||||||
|
- namespaceSelector:
|
||||||
|
matchLabels:
|
||||||
|
kubernetes.io/metadata.name: notifications
|
||||||
|
ports:
|
||||||
|
- protocol: TCP
|
||||||
|
port: 80
|
||||||
# Allow to MinIO (S3-compatible storage)
|
# Allow to MinIO (S3-compatible storage)
|
||||||
# Service `minio` listens on port 80 (targetPort 9000).
|
# Service `minio` listens on port 80 (targetPort 9000).
|
||||||
# Headless `minio-cluster-hl` is 9000. Allow both.
|
# Headless `minio-cluster-hl` is 9000. Allow both.
|
||||||
@@ -142,3 +162,24 @@ spec:
|
|||||||
port: 80
|
port: 80
|
||||||
- protocol: TCP
|
- protocol: TCP
|
||||||
port: 9000
|
port: 9000
|
||||||
|
# Allow to Forgejo registry v2 API (image digest lookup for promotion)
|
||||||
|
- to:
|
||||||
|
- namespaceSelector:
|
||||||
|
matchLabels:
|
||||||
|
kubernetes.io/metadata.name: cicd
|
||||||
|
ports:
|
||||||
|
- protocol: TCP
|
||||||
|
port: 3000
|
||||||
|
# Allow to kubernetes API server (ArgoCD Application patching via in-cluster SA).
|
||||||
|
# kube-apiserver runs on host network (192.168.1.162/166:6443) — must use ipBlock.
|
||||||
|
# ClusterIP 10.96.0.1:443 also routes there, covered by the 10.96.0.0/12 CIDR.
|
||||||
|
- to:
|
||||||
|
- ipBlock:
|
||||||
|
cidr: 10.96.0.0/12
|
||||||
|
- ipBlock:
|
||||||
|
cidr: 192.168.1.0/24
|
||||||
|
ports:
|
||||||
|
- protocol: TCP
|
||||||
|
port: 443
|
||||||
|
- protocol: TCP
|
||||||
|
port: 6443
|
||||||
|
|||||||
@@ -0,0 +1,62 @@
|
|||||||
|
# RBAC for api-gateway image promotion.
|
||||||
|
# Grants the api-gateway ServiceAccount get+patch on ArgoCD Applications
|
||||||
|
# so it can update kustomize.images / helm.parameters after a container push.
|
||||||
|
apiVersion: rbac.authorization.k8s.io/v1
|
||||||
|
kind: ClusterRole
|
||||||
|
metadata:
|
||||||
|
name: image-promoter
|
||||||
|
labels:
|
||||||
|
app: api-gateway
|
||||||
|
rules:
|
||||||
|
- apiGroups: ["argoproj.io"]
|
||||||
|
resources: ["applications"]
|
||||||
|
verbs: ["get", "patch"]
|
||||||
|
---
|
||||||
|
apiVersion: rbac.authorization.k8s.io/v1
|
||||||
|
kind: ClusterRoleBinding
|
||||||
|
metadata:
|
||||||
|
name: image-promoter-api-gateway
|
||||||
|
labels:
|
||||||
|
app: api-gateway
|
||||||
|
subjects:
|
||||||
|
- kind: ServiceAccount
|
||||||
|
name: api-gateway
|
||||||
|
namespace: api
|
||||||
|
roleRef:
|
||||||
|
kind: ClusterRole
|
||||||
|
name: image-promoter
|
||||||
|
apiGroup: rbac.authorization.k8s.io
|
||||||
|
---
|
||||||
|
# Role + RoleBinding for pi-session sandbox manager (§4.15.4).
|
||||||
|
# api-gateway SA needs get/create/delete for pods, pvcs, services,
|
||||||
|
# and networkpolicies in the poimen namespace.
|
||||||
|
apiVersion: rbac.authorization.k8s.io/v1
|
||||||
|
kind: Role
|
||||||
|
metadata:
|
||||||
|
name: session-sandbox-manager
|
||||||
|
namespace: poimen
|
||||||
|
labels:
|
||||||
|
app: api-gateway
|
||||||
|
rules:
|
||||||
|
- apiGroups: [""]
|
||||||
|
resources: ["pods", "persistentvolumeclaims", "services"]
|
||||||
|
verbs: ["get", "list", "create", "delete"]
|
||||||
|
- apiGroups: ["networking.k8s.io"]
|
||||||
|
resources: ["networkpolicies"]
|
||||||
|
verbs: ["get", "create", "delete"]
|
||||||
|
---
|
||||||
|
apiVersion: rbac.authorization.k8s.io/v1
|
||||||
|
kind: RoleBinding
|
||||||
|
metadata:
|
||||||
|
name: session-sandbox-manager-api-gateway
|
||||||
|
namespace: poimen
|
||||||
|
labels:
|
||||||
|
app: api-gateway
|
||||||
|
subjects:
|
||||||
|
- kind: ServiceAccount
|
||||||
|
name: api-gateway
|
||||||
|
namespace: api
|
||||||
|
roleRef:
|
||||||
|
kind: Role
|
||||||
|
name: session-sandbox-manager
|
||||||
|
apiGroup: rbac.authorization.k8s.io
|
||||||
Reference in New Issue
Block a user