Compare commits

..
Author SHA1 Message Date
Admin Bot d71c88498b fix(proxy): address review findings on admission control (#64)
- tracing: responseWriter now forwards Flush and exposes Unwrap. Without it
  w.(http.Flusher) failed for every request passing the tracing middleware, so
  the parked-stream keep-alive never started: a parked streaming request wrote
  nothing for up to ADMISSION_MAX_PARK_SECONDS, which is exactly the client
  timeout the keep-alive was added to prevent.
- proxy: once the SSE status line is spent, ModifyResponse frames an upstream
  response >= 400 as an event: error instead of letting a raw error body land
  inside the stream, where a streaming client reads it as an empty completion
  and reports success. net/http keeps the first status, so the upstream status
  cannot be forwarded after the park has written one.
- engine_state: new Snapshot() returns the state and its tick channel under one
  lock. Get and Tick taken separately can miss the poll that lands between the
  two calls, parking a request for an extra interval.
- engine_state: observer_degraded now tracks current parsability rather than
  being raised once on the first poll and never cleared, which pinned it at 1
  after a single truncated response and never reported drift that appeared
  later. markDegraded only touches the exporter when the value changes.
- deployment: ADMISSION_MAX_INFLIGHT is set explicitly to 4, documenting the
  coupling to the engine's --max-num-seqs, instead of silently defaulting to
  VLLM_MAX_NUM_SEQS.
2026-09-30 19:32:10 +08:00
Admin Bot c5a91cf647 feat(proxy): engine-wide LLM admission control via vLLM gauges (#64)
CI / CI (pull_request) Successful in 34m5s
The per-pod semaphore was not a cap: 3 replicas x 4 slots advertised 12
concurrent requests while the engine runs 4, and nothing coordinated the
replicas. Admission is now gated on the engine's own published gauges,
which is global by construction because every replica contends for the
same engine.

- new internal/proxy/engine_state.go: 2s scrape of the backend /metrics,
  hand-rolled exposition parse (no new dependency), unhealthy after 3
  consecutive failed polls, and a startup self-check that raises
  observer_degraded rather than falling back to fail-open silently
- gate: park while kv_cache_usage_perc >= 0.90 or running + waiting >=
  ADMISSION_MAX_INFLIGHT, bounded by ADMISSION_MAX_PARK_SECONDS, then
  503 + Retry-After: 5. Unknown state fails open on purpose
- streaming requests receive SSE keep-alive heartbeats while parked
- router: extract only the top-level model field instead of unmarshalling
  the whole body into a map; body cap from MAX_BODY_BYTES (1 MiB, the
  model path previously inherited the 100 MB default)
- GOMAXPROCS from the cgroup CPU quota: the container is limited to 500m
  while Go defaulted to the host core count and was CFS-throttled
- VLLM_MAX_QUEUE_DEPTH default 20 -> 8
- delete k8s/configmap.yaml: absent from kustomization.yaml, and the live
  ConfigMap of that name carries an unrelated key
- correct the llm-client NetworkPolicy comment; llm-serving-ingress now
  selects by namespace, not by that label

Observers run only for models named in ADMISSION_OBSERVE_MODELS.
reasoning is the only vLLM engine here; ornith and qwen-cpu are Ollama
and embeddings/rerank are TEI, none of which publish vllm:* series, so
observing them would hold observer_degraded at 1 permanently.

Companion change (engine flags + dashboard panels) is
riotpiao-poimen/homelab#51.
2026-09-30 17:39:47 +08:00
rock aac68386d8 fix(proxy): wire LLM latency metrics into the streaming response path (#63)
CI / CI (push) Successful in 34m13s
`NewLLMMetricsCapture`/`ResponseWriterWrapper` had **zero call sites**, so `llm_ttft_seconds`, `llm_itl_seconds` and `llm_tokens_total` were exported as HELP/TYPE stubs with no samples (verified against the live `/metrics`).

These are the only per-request LLM latency series covering the Ollama models (`ornith:35b`, `qwen2.5:3b-instruct`), which expose no `/metrics` of their own — so those models had zero latency visibility.

## Change

Wire the wrapper around the response writer in `ServeHTTP` for model-routed requests, and delete the unused `io.WriteCloser` duplicate that made the gap hard to spot (two near-identical implementations, neither constructed).

## Correctness details

- **Only SSE responses are measured.** A non-streaming completion is written in one `Write`, which would otherwise be recorded as a TTFT sample equal to the whole request latency, with no ITL samples — plausible-looking but meaningless numbers.
- **`requestStart` is captured at the top of `ServeHTTP`** so TTFT measures from request acceptance, not from when the upstream response is copied.
- **Token count skips the `data: [DONE]` terminator.**
- **`Unwrap()`** keeps `http.ResponseController` (Flush/Hijack/deadlines) working through the wrapper, so SSE unbuffering (#33) is unaffected.

## Tests

| Test | Guards |
|---|---|
| `TestResponseWriterWrapperRecordsSSEMetrics` | TTFT=1, ITL>0, tokens=2 (terminator excluded) |
| `TestResponseWriterWrapperIgnoresNonStreaming` | JSON response records nothing |
| `TestResponseWriterWrapperPreservesStreaming` | bytes + Flush pass through |
| `TestHandlerRecordsLLMMetricsForStreaming` | end-to-end through `ServeHTTP` |

The last one fails if the wrapper is ever wired nowhere again — verified by disabling the wiring locally and watching it fail.

## Verification needed post-merge

Confirm samples appear:
```
curl -s https://api.riotpiao.com/metrics | grep -E "llm_(ttft|itl)_seconds_count|llm_tokens_total"
```
2026-09-29 15:46:39 +00:00
rockandAdmin Bot 28753d11ed fix(netpol): allow Prometheus scrape of gateway /metrics (#62)
CI / CI (push) Successful in 32m48s
## Problem

The new `api-gateway` ServiceMonitor discovers all 3 replicas, but every scrape
fails with `context deadline exceeded`:

```
down   api-gateway   api-gateway-6986b7c889-pmt4b   err=Get "http://10.244.0.180:8080/metrics": context deadline
down   api-gateway   api-gateway-6986b7c889-c28dt   err=Get "http://10.244.1.1:8080/metrics": context deadline
down   api-gateway   api-gateway-6986b7c889-552gl   err=Get "http://10.244.2.61:8080/metrics": context deadline
```

## Cause

`network-policy.yaml` allowed Prometheus ingress with:

```yaml
namespaceSelector:
  matchLabels:
    name: monitoring
```

Only `ingress-nginx` carries a plain `name:` label in this cluster. Every other
namespace has just `kubernetes.io/metadata.name`, so this rule matched **nothing**
and silently denied the scrape. (The rule directly above it — `name: ingress-nginx`
— works only because that namespace is the documented exception.)

## Fix

Select on `kubernetes.io/metadata.name: monitoring`.

## Impact

Unblocks gateway metrics: `llm_ttft_seconds`, `llm_itl_seconds`,
`llm_tokens_total`, `queue_full_requests_total`,
`warmup_requests_delayed_total`, `vllm_prefix_cache_hit_tokens_total`.

These are the only per-request LLM latency series that cover the two Ollama
models (`ornith:35b`, `qwen2.5:3b-instruct`), which expose no metrics endpoint
of their own.

---------

Co-authored-by: Admin Bot <[email protected]>
Reviewed-on: #62
2026-09-29 14:40:31 +00:00
rock 3068c68730 Merge pull request 'feat(sessions): SandboxFleet integration — fleet primary, k8s fallback' (#61) from feat/sessions-sandbox-fleet-integration into main
CI / CI (push) Successful in 3m28s
2026-09-29 08:14:00 +00:00
Admin Bot da361f2a21 fix(sessions): nil map panic in updateSandboxSummary goroutine
CI / CI (pull_request) Successful in 2m24s
r.Context() is cancelled when HTTP handler returns; fetchEntitySummary
returns nil on cancelled context. Two fixes:
1. Detached request with context.Background() for the goroutine call
2. Nil guard on current before map assignment
2026-09-29 16:10:43 +08:00
Admin Bot 1e17706367 feat(sessions): SandboxFleet integration — fleet controller primary, SandboxManager fallback
FleetClient calls tool-definition-controller:8788 /sessions/{id}/assign
as the primary sandbox create path. Falls back to direct k8s SandboxManager
if fleet controller is unreachable.

handleSandboxCreate:
  1. fleet.Assign() → podSvcDns from pre-warmed fleet pod (fast, no wait)
  2. fallback: sandbox.CreateOrResume() → dedicated pod (slow, 30s wait)

handleSandboxDelete: fleet.Release() + sandbox.Delete() both run (non-fatal).
updateSandboxSummary: extracted helper, updates session memory_entity with
  pod_name / svc_dns / last_active_at.

NetworkPolicy: added port 8788 to poimen egress (fleet controller REST API).

Response adds fleet:true field when served by fleet.
2026-09-29 16:10:43 +08:00
rock 49c9071a59 Merge pull request 'feat(sessions): sandbox resource — per-session k8s Pod+PVC+Service (§4.15.4)' (#60) from feat/sessions-sandbox-manager into main
CI / CI (push) Successful in 3m22s
2026-09-28 14:02:53 +00:00
Admin Bot 16ddd6d4ad feat(sessions): sandbox resource — per-session k8s Pod+PVC+Service (§4.15.4)
CI / CI (pull_request) Successful in 2m8s
New X-Resource: sandbox on X-Service: sessions

POST ?session_id=  → CreateOrResume
  1. Ensure PVC workspace-session-<id> (Longhorn 1Gi)
  2. If pod already Running → return DNS (resume path, idempotent)
  3. Ensure pi-session-sandbox NetworkPolicy (shared, created once)
  4. Create Pod pi-session-<id> + Service pi-session-<id>:8787
  5. Wait up to 30s for Pod Ready
  6. Fire-and-forget: merge pod_name/pvc_name/svc_dns/last_active_at
     into session memory_entity summary

DELETE ?session_id= → delete Pod + Service + PVC

Idle reaper goroutine (5min interval, configurable via
SANDBOX_REAPER_INTERVAL_MINUTES): lists pi-session-* pods,
deletes Succeeded/Failed ones, preserves PVCs for resume.

Security: runAsNonRoot/1000, readOnlyRootFilesystem, drop ALL
caps, RuntimeDefault seccomp, automountServiceAccountToken:false.
Resources: 500m/512Mi req, 1/1Gi limit.
NetworkPolicy: egress → DNS + api:8080 only; ingress ← poimen-frontend.

RBAC: Role session-sandbox-manager + RoleBinding in poimen ns.
7 unit tests — all pass (mock k8s server).

Image: forgejo.riotpiao.com/riotpiao-poimen/poimen-pi-session-runtime:latest
2026-09-28 19:53:37 +08:00
rock fadc658ac0 Merge pull request 'fix(sessions): use entities query for list (hybridSearch requires project)' (#59) from fix/sessions-list-entities-query into main
CI / CI (push) Successful in 3m18s
2026-09-28 10:43:24 +00:00
Admin Bot bb7a490f5c fix(sessions): use entities query for list — hybridSearch requires project arg
CI / CI (pull_request) Successful in 2m10s
hybridSearch mandates a project scope, but sessions scatter across
per-session projects (session-<uuid>). The entities query accepts
optional project — omit it to list SESSION entities across all projects.

Also adds projectId to the list response so callers can scope follow-up
queries to the right session project.

Fixes: GET X-Service:sessions X-Resource:sessions returning 502
2026-09-28 18:43:00 +08:00
rock 4a3768ac35 Merge pull request 'fix(netpol): add poimen:8090+6379 egress for sessions adapter' (#58) from fix/netpol-poimen-graphql into main
CI / CI (push) Successful in 3m8s
2026-09-28 08:31:09 +00:00
Admin Bot f8882e86ae fix(netpol): add poimen:8090 (graphql) + 6379 (redis) to egress
CI / CI (pull_request) Successful in 2m9s
Sessions adapter calls poimen-graphql on port 8090 for upsertEntity
mutations and hybridSearch reads. Was only allowing 8080 (memory).
2026-09-28 16:30:44 +08:00
rock f9093fe713 Merge pull request 'fix(test): update TestPromote_KustomizeSuccess for kustomize patch path' (#57) from fix/promoter-test-kustomize-path into main
CI / CI (push) Successful in 3m26s
2026-09-28 07:18:56 +00:00
Admin Bot ecaaaa8792 fix(test): update TestPromote_KustomizeSuccess for new kustomize patch path
CI / CI (pull_request) Successful in 2m15s
Patch now targets /spec/sources/0/kustomize (full block) not
/spec/sources/0/kustomize/images (subpath that fails when kustomize
key is absent from the Application CR).
2026-09-28 14:18:32 +07:00
rock 6d304a3e5a Merge pull request 'fix(promoter): correct Forgejo registry service hostname' (#56) from fix/promoter-registry-hostname into main
CI / CI (push) Failing after 2m10s
fix(promoter): correct Forgejo registry hostname + kustomize patch path

Patch /spec/sources/N/kustomize (full block) not /kustomize/images subpath.
forgejo-gitea-http.cicd.svc.cluster.local:3000 correct hostname.
CiliumNetworkPolicy kube-apiserver entity egress added.
2026-09-28 07:16:54 +00:00
Admin Bot 3ca654513f fix(promoter): correct Forgejo registry service hostname
CI / CI (pull_request) Failing after 2m6s
forgejo-gitea-http.cicd.svc.cluster.local (headless svc, port 3000)
not forgejo-gitea (doesn't exist). Also wire REGISTRY_BASE env var
in deployment.yaml so it can be overridden without code change.
2026-09-28 12:34:54 +08:00
rockandAdmin Bot 23fff1006a feat(webhook): Forgejo package webhook → ArgoCD image promotion (#55)
CI / CI (push) Successful in 3m21s
Closes #54

Replaces ArgoCD Image Updater with webhook-driven promotion.

## What changed

| File | Change |
|---|---|
| `internal/webhook/promoter.go` | NEW — ImagePromoter + promotion table (10 entries) |
| `internal/webhook/promoter_test.go` | NEW — 13 tests |
| `internal/webhook/forgejo.go` | package event routed to ImagePromoter |
| `k8s/rbac.yaml` | NEW — ClusterRole image-promoter + ClusterRoleBinding |
| `k8s/deployment.yaml` | REGISTRY_PAT env var |
| `k8s/network-policy.yaml` | egress cicd:3000 + default:443 |
| `k8s/kustomization.yaml` | rbac.yaml added |

## homelab repo changes needed

See comment on issue #54 for exact YAML.

---------

Co-authored-by: Admin Bot <[email protected]>
Reviewed-on: #55
2026-09-28 02:45:41 +00:00
rockandAdmin Bot 50bcd19634 feat(sessions): switch writes to GraphQL upsertEntity mutation (#53)
CI / CI (push) Successful in 3m17s
Closes #52

## Summary

Migrates all session/milestone writes from REST `/memory/ingest` to GraphQL `upsertEntity` mutation — direct, typed, no LLM extraction lag.

## Changes

- **Removed**: `memoryURL` field, `defaultMemoryURL` const, `ingest()` REST helper, `"io"` import
- **Added**: `upsertEntity()` GraphQL mutation helper
- **Added**: `fetchEntitySummary()` helper — reads current entity before patch for read-modify-write
- **handlePatch**: now fetches current summary, merges patch fields, re-upserts full blob — no data loss on partial updates
- **handleCreate / handleCreateMilestone**: `json.Marshal` error now returns 500 instead of silently ignoring
- 29 tests pass, 85.4% coverage

## Mutation used

```graphql
mutation {
  upsertEntity(input: {
    project: "session-<uuid>"
    name: "<sessionID | milestoneID>"
    entityType: SESSION | MILESTONE
    summary: "{...json blob...}"
  }) { id created }
}
```

## Single upstream

All reads and writes now go through `graphqlURL` (poimen-graphql:8090). `MEMORY_UPSTREAM_URL` env var no longer read by this adapter.

---------

Co-authored-by: Admin Bot <[email protected]>
Reviewed-on: #53
2026-09-28 02:24:11 +00:00
rockandAdmin Bot a24e4f7174 feat(sessions): SessionAdapter — REST CRUD for sessions/milestones/episodes (#51)
CI / CI (push) Successful in 3m15s
Closes #50

## Changes

- `internal/sessions/handler.go` — new Handler; 7 endpoints dispatched via `X-Upstream-Path`; writes to `poimen-memory` via `/memory/ingest`; reads from `poimen-graphql` via GraphQL `hybridSearch`/`entity`
- `internal/serviceadapter/session_adapter.go` — `NewSessionAdapter()` + `GetSessionSpec()` with 5 resources
- `cmd/gateway/main.go` — register sessions adapter after memory adapter
- `internal/serviceadapter/adapters.go` — add `sessions` case to factory
- `internal/sessions/handler_test.go` — 30 tests, all pass

## Design

- `X-Service: sessions` + `X-Resource: <sessions|session|milestones|episodes|episode>`
- IDs via query params (`?session_id=`, `?episode_id=`) — dispatcher does not substitute path params
- `episode_id = sha256(session_id+milestone_id+attempt)` per §4.14 Q6
- Content JSON-encoded (injection-safe)
- `tracing.NewTransport` wired for cross-service trace propagation
- GraphQL `errors` array detected, surfaced as 502
- limit capped at 200

---------

Co-authored-by: Admin Bot <[email protected]>
Reviewed-on: #51
2026-09-28 01:53:20 +00:00
rockandAdmin Bot 0bde5dd8fe fix(netpol): allow api-gateway egress to notifications:80 (Gotify) (#48)
CI / CI (push) Successful in 3m7s
NetworkPolicy was blocking api-gateway → gotify.notifications.svc.cluster.local:80.
Adds egress rule to `notifications` namespace on port 80.

Required for `X-Service: notification` + `X-Resource: send-message` to deliver push notifications through the gateway.

---------

Co-authored-by: Admin Bot <[email protected]>
Reviewed-on: #48
2026-09-27 14:13:36 +00:00
rockandAdmin Bot ef5e736673 feat(cag): CAG integration — memory adapter, warmup semaphore, cache-aware routing (#47)
CI / CI (push) Successful in 3m20s
## Summary

Three issues in one PR. Implements Cache-Augmented Generation (CAG) support in the api-gw.

---

### #44 — Memory ServiceAdapter: standing-context resource

- Code-registers `memory` adapter with `standing-context` resource (`POST /memory/standing-context`)
- Fixes dispatcher bug: method-level `Auth.Capability` override now respected (was always using verb-derived capability)
- Auth: `memory:read` required via method-level override; `memory:write` for ingest
- `MEMORY_UPSTREAM_URL` env configures upstream

### #45 — Continuous batching: warmup deprioritisation + per-backend semaphore

- In-memory bounded queue per backend: `VLLM_MAX_QUEUE_DEPTH` slots (default 20)
- Concurrency limiter: `VLLM_MAX_NUM_SEQS` slots (default 4)
- `X-Poimen-Warmup: true` → sleep `WARMUP_DELAY_MS` (default 500ms) before queue entry, strip header, add `X-Priority: low`
- Queue full → 503 immediately, no goroutine pile-up
- Metrics: `warmup_requests_delayed_total`, `queue_full_requests_total` at `/metrics`

### #46 — Cache-aware routing: rendezvous hash + prefix cache hit logging

- `addresses: [...]` YAML field for multi-backend models (backward-compat with `address`)
- `X-Poimen-Cache-Key` header → rendezvous-hash backend selection (same key = same GPU)
- No cache key → round-robin
- Failed backend excluded 30s; next-highest-score fallback
- `vllm-num-cached-tokens` response header captured → `vllm_prefix_cache_hit_tokens_total` counter + structured log

---

## Verification

```
go build ./...    ✅
go vet ./...      ✅
go test -race ./... ✅ (all packages, 0 races)
```

## Reviewer verdicts

- reviewer-sre: SIGN-OFF
- reviewer-quality: SIGN-OFF

## Non-blocking findings (🟡🔵 only)

- `getOrCreateTransport` not goroutine-safe by design; pre-population invariant documented
- `time.Sleep(warmupDelay)` holds goroutine; context-cancel alternative left as TODO
- OTel trace_id not propagated into CAG log events (existing gap)

Closes #44 Closes #45 Closes #46

---------

Co-authored-by: Admin Bot <[email protected]>
Reviewed-on: #47
2026-09-27 11:16:40 +00:00
rock aaa9a5d1cb fix: use query params for workflow_id in Poimen endpoints (#42)
CI / CI (push) Successful in 3m7s
**Root cause:** ServiceAdapter dispatcher sets X-Upstream-Path to template strings like `/v1/poimen/executions/{id}/stream` with literal `{id}`. The handler extracted `{id}` as the workflow ID instead of the actual value.

**Fix:** All handlers read `workflow_id` from query params (`?workflow_id=xxx`). Paths changed to flat resource names without path parameters.

**API contract:**
```
GET  /executions                         — list
GET  /executions/detail?workflow_id=xxx   — describe
GET  /executions/history?workflow_id=xxx  — history
GET  /executions/stream?workflow_id=xxx   — SSE
POST /executions/signal?workflow_id=xxx   — signal
DELETE /executions/detail?workflow_id=xxx — terminate
```

All accessed via `X-Service: poimen` + `X-Resource: <name>` headers.

---------
Reviewed-on: #42
2026-09-26 12:52:19 +00:00
28 changed files with 4558 additions and 243 deletions
+62
View File
@@ -4,9 +4,13 @@ import (
"context"
"fmt"
"log"
"math"
"net/http"
"os"
"os/signal"
"runtime"
"strconv"
"strings"
"syscall"
"forgejo.riotpiao.com/rock/homelab-frontend/internal/auth"
@@ -21,6 +25,10 @@ import (
)
func main() {
// Size the Go scheduler to the container's CPU quota before any server
// goroutine starts.
configureGOMAXPROCS()
ctx := context.Background()
// Initialize OpenTelemetry tracing
@@ -117,6 +125,10 @@ func main() {
// 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)
for _, a := range cfg.Adapters {
if existing := registry.Get(a.ServiceName); existing != nil {
@@ -173,3 +185,53 @@ func main() {
log.Printf("gateway shutdown complete")
os.Exit(0)
}
// configureGOMAXPROCS sets GOMAXPROCS from the cgroup CPU quota so the runtime
// sizes its scheduler to the container limit rather than the host CPU count.
// The default is left alone when no quota is configured.
func configureGOMAXPROCS() {
quota, period := readCPUQuota()
if quota <= 0 || period <= 0 {
return
}
n := int(math.Ceil(float64(quota) / float64(period)))
if n < 1 {
n = 1
}
runtime.GOMAXPROCS(n)
log.Printf("GOMAXPROCS set to %d (cgroup cpu quota %d/%d)", n, quota, period)
}
// readCPUQuota returns the CPU quota and period from cgroup v2, falling back to
// cgroup v1. A quota of 0 means no limit.
func readCPUQuota() (int64, int64) {
// cgroup v2: "<quota> <period>" or "max <period>".
if data, err := os.ReadFile("/sys/fs/cgroup/cpu.max"); err == nil {
fields := strings.Fields(string(data))
if len(fields) >= 2 {
period, perr := strconv.ParseInt(fields[1], 10, 64)
if perr == nil && period > 0 {
if fields[0] == "max" {
return 0, period
}
if quota, qerr := strconv.ParseInt(fields[0], 10, 64); qerr == nil {
return quota, period
}
}
}
}
// cgroup v1: cpu.cfs_quota_us is -1 when unlimited.
period := int64(100000)
if data, err := os.ReadFile("/sys/fs/cgroup/cpu/cpu.cfs_period_us"); err == nil {
if v, perr := strconv.ParseInt(strings.TrimSpace(string(data)), 10, 64); perr == nil && v > 0 {
period = v
}
}
if data, err := os.ReadFile("/sys/fs/cgroup/cpu/cpu.cfs_quota_us"); err == nil {
if v, qerr := strconv.ParseInt(strings.TrimSpace(string(data)), 10, 64); qerr == nil {
return v, period
}
}
return 0, 0
}
+121 -12
View File
@@ -15,7 +15,7 @@ type Metrics struct {
// Request latencies: request_duration_seconds (histogram)
// Stored as cumulative buckets for Prometheus text format
requestDuration map[string]int64 // stores duration samples in milliseconds
requestDuration map[string]int64 // stores duration samples in milliseconds
requestDurationBuckets map[string]map[float64]int64 // histogram buckets
// Bytes counters: gateway_bytes{direction, route, upstream}
@@ -41,22 +41,33 @@ type Metrics struct {
warmupRequestsDelayed int64
prefixCacheHitTokens int64
queueFullRequests int64
// Admission control metrics (#64)
admissionParked int64
admissionRejected map[string]int64
engineStateStale int64
// Engine gauge readings mirrored from vLLM (#64)
vllmRequestsRunning int64
vllmKVUsagePerc float64
observerDegraded int
}
// NewMetrics creates a new Metrics instance.
func NewMetrics() *Metrics {
return &Metrics{
requestTotal: make(map[string]int64),
requestDuration: make(map[string]int64),
requestDurationBuckets: make(map[string]map[float64]int64),
bytesIn: make(map[string]int64),
bytesOut: make(map[string]int64),
upstreamHealth: make(map[string]int),
requestTotal: make(map[string]int64),
requestDuration: make(map[string]int64),
requestDurationBuckets: make(map[string]map[float64]int64),
bytesIn: make(map[string]int64),
bytesOut: make(map[string]int64),
upstreamHealth: make(map[string]int),
streamingResponsesTotal: make(map[string]int64),
streamingByteCount: make(map[string]int64),
ttftMs: make(map[string][]int64),
itlMs: make(map[string][]int64),
tokenCount: make(map[string]int64),
streamingByteCount: make(map[string]int64),
ttftMs: make(map[string][]int64),
itlMs: make(map[string][]int64),
tokenCount: make(map[string]int64),
admissionRejected: make(map[string]int64),
}
}
@@ -133,7 +144,7 @@ func (m *Metrics) RecordStreamingResponse(route, upstream string, totalBytes int
func (m *Metrics) recordDuration(key string, duration time.Duration) {
m.requestDuration[key] += int64(duration.Milliseconds())
if _, ok := m.requestDurationBuckets[key]; !ok {
m.requestDurationBuckets[key] = make(map[float64]int64)
}
@@ -288,6 +299,12 @@ func (m *Metrics) Reset() {
m.warmupRequestsDelayed = 0
m.prefixCacheHitTokens = 0
m.queueFullRequests = 0
m.admissionParked = 0
m.admissionRejected = make(map[string]int64)
m.engineStateStale = 0
m.vllmRequestsRunning = 0
m.vllmKVUsagePerc = 0
m.observerDegraded = 0
}
// IncrQueueFull increments queue_full_requests_total.
@@ -332,6 +349,98 @@ func (m *Metrics) GetPrefixCacheHitTokens() int64 {
return m.prefixCacheHitTokens
}
// IncrAdmissionParked increments admission_parked_total.
func (m *Metrics) IncrAdmissionParked() {
m.mu.Lock()
defer m.mu.Unlock()
m.admissionParked++
}
// GetAdmissionParked returns admission_parked_total.
func (m *Metrics) GetAdmissionParked() int64 {
m.mu.RLock()
defer m.mu.RUnlock()
return m.admissionParked
}
// IncrAdmissionRejected increments admission_rejected_total{reason=reason}.
func (m *Metrics) IncrAdmissionRejected(reason string) {
m.mu.Lock()
defer m.mu.Unlock()
m.admissionRejected[reason]++
}
// GetAdmissionRejected returns a copy of admission_rejected_total by reason.
func (m *Metrics) GetAdmissionRejected() map[string]int64 {
m.mu.RLock()
defer m.mu.RUnlock()
out := make(map[string]int64, len(m.admissionRejected))
for reason, count := range m.admissionRejected {
out[reason] = count
}
return out
}
// IncrEngineStateStale increments engine_state_stale_total.
func (m *Metrics) IncrEngineStateStale() {
m.mu.Lock()
defer m.mu.Unlock()
m.engineStateStale++
}
// GetEngineStateStale returns engine_state_stale_total.
func (m *Metrics) GetEngineStateStale() int64 {
m.mu.RLock()
defer m.mu.RUnlock()
return m.engineStateStale
}
// SetVLLMRequestsRunning sets vllm_requests_running from the gateway's own reading.
func (m *Metrics) SetVLLMRequestsRunning(n int64) {
m.mu.Lock()
defer m.mu.Unlock()
m.vllmRequestsRunning = n
}
// GetVLLMRequestsRunning returns vllm_requests_running.
func (m *Metrics) GetVLLMRequestsRunning() int64 {
m.mu.RLock()
defer m.mu.RUnlock()
return m.vllmRequestsRunning
}
// SetVLLMKVUsagePerc sets vllm_kv_usage_perc from the gateway's own reading.
func (m *Metrics) SetVLLMKVUsagePerc(v float64) {
m.mu.Lock()
defer m.mu.Unlock()
m.vllmKVUsagePerc = v
}
// GetVLLMKVUsagePerc returns vllm_kv_usage_perc.
func (m *Metrics) GetVLLMKVUsagePerc() float64 {
m.mu.RLock()
defer m.mu.RUnlock()
return m.vllmKVUsagePerc
}
// SetObserverDegraded sets observer_degraded (1 = startup self-check failed).
func (m *Metrics) SetObserverDegraded(degraded bool) {
m.mu.Lock()
defer m.mu.Unlock()
if degraded {
m.observerDegraded = 1
} else {
m.observerDegraded = 0
}
}
// GetObserverDegraded returns observer_degraded.
func (m *Metrics) GetObserverDegraded() int {
m.mu.RLock()
defer m.mu.RUnlock()
return m.observerDegraded
}
// globalMetrics is the singleton Metrics instance for gateway-wide telemetry.
var (
globalMetrics *Metrics
+32
View File
@@ -48,6 +48,38 @@ func (p *PrometheusExporter) Export() string {
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()))
lines = append(lines, "# HELP admission_parked_total Requests parked waiting for engine admission")
lines = append(lines, "# TYPE admission_parked_total counter")
lines = append(lines, fmt.Sprintf("admission_parked_total %d", p.metrics.GetAdmissionParked()))
lines = append(lines, "# HELP admission_rejected_total Requests rejected by engine admission control")
lines = append(lines, "# TYPE admission_rejected_total counter")
rejected := p.metrics.GetAdmissionRejected()
reasons := make([]string, 0, len(rejected))
for reason := range rejected {
reasons = append(reasons, reason)
}
sort.Strings(reasons)
for _, reason := range reasons {
lines = append(lines, fmt.Sprintf(`admission_rejected_total{reason="%s"} %d`, reason, rejected[reason]))
}
lines = append(lines, "# HELP engine_state_stale_total Times the engine metrics feed went stale")
lines = append(lines, "# TYPE engine_state_stale_total counter")
lines = append(lines, fmt.Sprintf("engine_state_stale_total %d", p.metrics.GetEngineStateStale()))
lines = append(lines, "# HELP vllm_requests_running vLLM requests running as read by the gateway")
lines = append(lines, "# TYPE vllm_requests_running gauge")
lines = append(lines, fmt.Sprintf("vllm_requests_running %d", p.metrics.GetVLLMRequestsRunning()))
lines = append(lines, "# HELP vllm_kv_usage_perc vLLM KV cache usage fraction as read by the gateway")
lines = append(lines, "# TYPE vllm_kv_usage_perc gauge")
lines = append(lines, fmt.Sprintf("vllm_kv_usage_perc %g", p.metrics.GetVLLMKVUsagePerc()))
lines = append(lines, "# HELP observer_degraded Engine metrics observer startup self-check failed (1) or healthy (0)")
lines = append(lines, "# TYPE observer_degraded gauge")
lines = append(lines, fmt.Sprintf("observer_degraded %d", p.metrics.GetObserverDegraded()))
return strings.Join(lines, "\n") + "\n"
}
+333
View File
@@ -0,0 +1,333 @@
package proxy
import (
"bufio"
"context"
"fmt"
"io"
"math"
"net/http"
"strconv"
"strings"
"sync"
"time"
"forgejo.riotpiao.com/rock/homelab-frontend/internal/logging"
"forgejo.riotpiao.com/rock/homelab-frontend/internal/observability"
)
// Metric names published by the vLLM engine. Admission is gated on these
// rather than on a per-pod semaphore: every gateway replica reads the same
// engine, so the limit is global by construction (#64).
const (
metricRequestsRunning = "vllm:num_requests_running"
metricRequestsWaiting = "vllm:num_requests_waiting"
metricKVCacheUsage = "vllm:kv_cache_usage_perc"
// enginePollInterval is how often each backend's /metrics is scraped.
enginePollInterval = 2 * time.Second
// engineStaleThreshold is the number of consecutive failed polls after
// which the state is marked unhealthy and admission fails open.
engineStaleThreshold = 3
// engineMetricsTimeout bounds a single scrape.
engineMetricsTimeout = 2 * time.Second
)
// EngineState is the gateway's view of one vLLM engine, mirrored from its
// own metrics endpoint.
type EngineState struct {
Running int64
Waiting int64
KVUsage float64
UpdatedAt time.Time
Healthy bool
}
// InFlight returns the number of requests the engine currently holds:
// running plus waiting.
func (s EngineState) InFlight() int64 {
return s.Running + s.Waiting
}
// EngineObserver polls one vLLM backend's /metrics and publishes the gauges
// admission is gated on. Healthy is false until the first successful poll and
// after engineStaleThreshold consecutive failures.
type EngineObserver struct {
addr string
interval time.Duration
client *http.Client
mu sync.Mutex
state EngineState
failures int
degraded bool
tick chan struct{}
}
// NewEngineObserver creates an observer for the backend at addr.
func NewEngineObserver(addr string, interval time.Duration) *EngineObserver {
return &EngineObserver{
addr: addr,
interval: interval,
// Own client: the proxy transports carry socket options and connect
// timeouts tuned for inference, not for a short metrics scrape.
client: &http.Client{Timeout: engineMetricsTimeout},
tick: make(chan struct{}),
}
}
// Start launches the polling goroutine. It returns immediately; the goroutine
// exits when ctx is cancelled.
func (o *EngineObserver) Start(ctx context.Context) {
go func() {
ticker := time.NewTicker(o.interval)
defer ticker.Stop()
// The first scrape waits one interval rather than firing at Start().
// Scraping immediately makes the gateway contact every observed backend
// before any request needs one, which is wasted work and a surprising
// side effect for anything asserting on upstream traffic. State is
// unknown until the first poll, and unknown fails open, so the deferral
// costs at most one interval of ungated admission. The startup
// self-check (decision Q21) still runs, on that first tick.
startup := true
for {
select {
case <-ctx.Done():
return
case <-ticker.C:
o.poll(ctx, startup)
startup = false
}
}
}()
}
// Get returns the most recent engine state.
func (o *EngineObserver) Get() EngineState {
o.mu.Lock()
defer o.mu.Unlock()
return o.state
}
// Tick returns a channel that is closed and replaced after every successful
// poll, waking requests parked on engine admission.
func (o *EngineObserver) Tick() <-chan struct{} {
o.mu.Lock()
defer o.mu.Unlock()
return o.tick
}
// Degraded reports whether the startup self-check found the metrics feed
// unusable (parse failure or missing series).
func (o *EngineObserver) Degraded() bool {
o.mu.Lock()
defer o.mu.Unlock()
return o.degraded
}
// Snapshot returns the latest state together with the tick channel that the
// next successful poll will close. Both are read under one lock on purpose:
// calling Get and Tick separately can miss the wakeup that lands between the
// two calls, parking a request for an extra poll interval.
func (o *EngineObserver) Snapshot() (EngineState, <-chan struct{}) {
o.mu.Lock()
defer o.mu.Unlock()
return o.state, o.tick
}
// markDegraded records whether the metrics feed is currently unusable and
// mirrors that into observer_degraded. It is cleared on the next parseable
// scrape, so one truncated response cannot pin the gauge for the process
// lifetime while metric-name drift after the first poll is still reported.
func (o *EngineObserver) markDegraded(v bool) {
o.mu.Lock()
changed := o.degraded != v
o.degraded = v
o.mu.Unlock()
if changed {
observability.Global().SetObserverDegraded(v)
}
}
// poll scrapes /metrics once. startup marks the first poll, whose failure is a
// self-diagnosis signal rather than a transient blip.
func (o *EngineObserver) poll(ctx context.Context, startup bool) {
req, err := http.NewRequestWithContext(ctx, http.MethodGet, "http://"+o.addr+"/metrics", nil)
if err != nil {
o.recordFailure()
logging.Warnf("engine metrics request build failed", map[string]string{
"address": o.addr,
"err": err.Error(),
})
return
}
resp, err := o.client.Do(req)
if err != nil {
o.recordFailure()
logging.Warnf("engine metrics scrape failed", map[string]string{
"address": o.addr,
"err": err.Error(),
})
return
}
defer resp.Body.Close()
if resp.StatusCode != http.StatusOK {
o.recordFailure()
logging.Warnf("engine metrics scrape returned non-200", map[string]string{
"address": o.addr,
"status": strconv.Itoa(resp.StatusCode),
})
return
}
running, waiting, kvUsage, err := ParseEngineMetrics(resp.Body)
if err != nil {
// Metric-name drift must be loud, never a silent fall-back to fail-open
// admission (decision Q21). The gauge tracks current state, so it clears
// on the next parseable scrape instead of pinning for the pod's life.
o.markDegraded(true)
if startup {
logging.Errorf("engine metrics unusable at startup", err, map[string]string{
"address": o.addr,
})
} else {
logging.Warnf("engine metrics parse failed", map[string]string{
"address": o.addr,
"err": err.Error(),
})
}
o.recordFailure()
return
}
o.markDegraded(false)
o.recordSuccess(running, waiting, kvUsage)
}
// recordSuccess stores a fresh reading and wakes parked waiters.
func (o *EngineObserver) recordSuccess(running, waiting int64, kvUsage float64) {
o.mu.Lock()
o.state = EngineState{
Running: running,
Waiting: waiting,
KVUsage: kvUsage,
UpdatedAt: time.Now(),
Healthy: true,
}
o.failures = 0
close(o.tick)
o.tick = make(chan struct{})
o.mu.Unlock()
m := observability.Global()
m.SetVLLMRequestsRunning(running)
m.SetVLLMKVUsagePerc(kvUsage)
}
// recordFailure counts a failed poll and, on the third consecutive failure,
// marks the state unhealthy and counts the staleness transition once — not
// once per failed poll. Last known values are kept.
func (o *EngineObserver) recordFailure() {
o.mu.Lock()
o.failures++
stale := o.failures == engineStaleThreshold
if o.failures >= engineStaleThreshold {
o.state.Healthy = false
}
failures := o.failures
o.mu.Unlock()
if stale {
observability.Global().IncrEngineStateStale()
logging.Warnf("engine metrics stale; admission is failing open", map[string]string{
"address": o.addr,
"failures": strconv.Itoa(failures),
})
}
}
// ParseEngineMetrics extracts the running, waiting and KV-cache-usage gauges
// from a vLLM Prometheus exposition body. Comment lines, unrelated series and
// non-finite samples are ignored; values are summed across repeated label sets.
// It returns an error naming any of the three series that is absent.
func ParseEngineMetrics(r io.Reader) (running int64, waiting int64, kvUsage float64, err error) {
scanner := bufio.NewScanner(r)
// The real exposition body is large; give the scanner room per line.
scanner.Buffer(make([]byte, 0, 64*1024), 1024*1024)
var hasRunning, hasWaiting, hasKV bool
for scanner.Scan() {
line := strings.TrimSpace(scanner.Text())
if line == "" || strings.HasPrefix(line, "#") {
continue
}
name, value, ok := splitMetricSample(line)
if !ok {
continue
}
switch name {
case metricRequestsRunning, metricRequestsWaiting, metricKVCacheUsage:
default:
continue
}
f, err := strconv.ParseFloat(value, 64)
if err != nil || math.IsNaN(f) || math.IsInf(f, 0) {
continue
}
switch name {
case metricRequestsRunning:
running += int64(f)
hasRunning = true
case metricRequestsWaiting:
waiting += int64(f)
hasWaiting = true
case metricKVCacheUsage:
kvUsage += f
hasKV = true
}
}
if err := scanner.Err(); err != nil {
return 0, 0, 0, fmt.Errorf("read engine metrics: %w", err)
}
var missing []string
if !hasRunning {
missing = append(missing, metricRequestsRunning)
}
if !hasWaiting {
missing = append(missing, metricRequestsWaiting)
}
if !hasKV {
missing = append(missing, metricKVCacheUsage)
}
if len(missing) > 0 {
return 0, 0, 0, fmt.Errorf("engine metrics missing series: %s", strings.Join(missing, ", "))
}
return running, waiting, kvUsage, nil
}
// splitMetricSample splits one Prometheus exposition line into its metric name
// and its first sample value. Lines without a value are reported as not ok.
func splitMetricSample(line string) (name, value string, ok bool) {
end := strings.IndexAny(line, "{ \t")
if end <= 0 {
return "", "", false
}
name = line[:end]
rest := line[end:]
if rest[0] == '{' {
closeBrace := strings.LastIndex(rest, "}")
if closeBrace < 0 {
return "", "", false
}
rest = rest[closeBrace+1:]
}
fields := strings.Fields(rest)
if len(fields) == 0 {
return "", "", false
}
return name, fields[0], true
}
+68 -96
View File
@@ -3,7 +3,6 @@ package proxy
import (
"bufio"
"fmt"
"io"
"net"
"net/http"
"strings"
@@ -12,87 +11,37 @@ import (
"forgejo.riotpiao.com/rock/homelab-frontend/internal/observability"
)
// LLMMetricsCapture wraps a response writer to capture TTFT and ITL metrics
type LLMMetricsCapture struct {
writer io.WriteCloser
model string
metrics *observability.Metrics
firstTokenTime time.Time
lastTokenTime time.Time
requestStartTime time.Time
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
// ResponseWriterWrapper wraps an http.ResponseWriter to capture per-request LLM
// latency metrics (TTFT, ITL, generated tokens) while a streaming response is
// copied from the upstream by httputil.ReverseProxy.
//
// Only server-sent-event responses are measured. A non-streaming completion is
// written in a single Write, so recording it would emit one "TTFT" sample equal
// to the whole request latency and no ITL samples at all — numbers that look
// plausible but mean nothing.
type ResponseWriterWrapper struct {
writer http.ResponseWriter
statusCode int
metrics *observability.Metrics
model string
startTime time.Time
firstByteTime time.Time
lastWriteTime time.Time
ttftRecorded bool
writer http.ResponseWriter
statusCode int
metrics *observability.Metrics
model string
startTime time.Time
firstByteTime time.Time
lastWriteTime time.Time
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 {
return &ResponseWriterWrapper{
writer: w,
model: model,
metrics: metrics,
startTime: startTime,
statusCode: 200,
writer: w,
model: model,
metrics: metrics,
startTime: startTime,
statusCode: 200,
}
}
@@ -101,29 +50,48 @@ func (w *ResponseWriterWrapper) Header() http.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
func (w *ResponseWriterWrapper) Write(b []byte) (int, error) {
// Record TTFT on first write
if !w.ttftRecorded && len(b) > 0 {
if len(b) > 0 && w.isStreaming() {
now := time.Now()
ttft := now.Sub(w.startTime).Milliseconds()
w.metrics.RecordTTFT(w.model, ttft)
w.ttftRecorded = true
w.firstByteTime = now
w.lastWriteTime = now
}
// Record ITL for subsequent writes (for streaming)
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)
}
if !w.ttftRecorded {
w.metrics.RecordTTFT(w.model, now.Sub(w.startTime).Milliseconds())
w.ttftRecorded = true
w.firstByteTime = 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.
w.metrics.RecordITL(w.model, itl)
}
w.lastWriteTime = now
if n := countSSETokens(b); n > 0 {
w.metrics.RecordTokenCount(w.model, n)
}
}
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")
}
// 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
}
+171
View File
@@ -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")
}
}
+263 -27
View File
@@ -2,6 +2,7 @@
package proxy
import (
"bytes"
"context"
"encoding/json"
"errors"
@@ -15,6 +16,7 @@ import (
"sort"
"strconv"
"strings"
"sync"
"syscall"
"time"
@@ -46,6 +48,12 @@ type Handler struct {
schedulers map[string]*Scheduler
// CAG: delay applied to warmup requests before queue entry (#45)
warmupDelay time.Duration
// engineObservers polls each model backend's /metrics for admission (#64)
engineObservers map[string]*EngineObserver
// maxBodyBytes caps synthesized model-route request bodies (MAX_BODY_BYTES)
maxBodyBytes int64
// transportsMu guards h.transports for goroutine-safe reads/writes.
transportsMu sync.RWMutex
}
// Route represents a reverse proxy route.
@@ -119,18 +127,80 @@ func New(cfg *config.Config) *Handler {
maxSlots = n
}
}
maxQueue := 20
maxQueue := 8
if v := os.Getenv("VLLM_MAX_QUEUE_DEPTH"); v != "" {
if n, err := strconv.Atoi(v); err == nil && n > 0 {
maxQueue = n
}
}
// Issue #64: engine-wide admission thresholds. MaxInFlight defaults to the
// same value as VLLM_MAX_NUM_SEQS; the observers below read the engine's
// own gauges, which every replica shares.
admissionMaxInFlight := maxSlots
if v := os.Getenv("ADMISSION_MAX_INFLIGHT"); v != "" {
if n, err := strconv.Atoi(v); err == nil && n > 0 {
admissionMaxInFlight = n
}
}
admissionMaxKVUsage := 0.90
if v := os.Getenv("ADMISSION_MAX_KV_USAGE"); v != "" {
if f, err := strconv.ParseFloat(v, 64); err == nil && f > 0 {
admissionMaxKVUsage = f
}
}
maxParkSeconds := 120
if v := os.Getenv("ADMISSION_MAX_PARK_SECONDS"); v != "" {
if n, err := strconv.Atoi(v); err == nil && n > 0 {
maxParkSeconds = n
}
}
h.maxBodyBytes = 1048576
if v := os.Getenv("MAX_BODY_BYTES"); v != "" {
if n, err := strconv.ParseInt(v, 10, 64); err == nil && n > 0 {
h.maxBodyBytes = n
}
}
limits := AdmissionLimits{
MaxInFlight: admissionMaxInFlight,
MaxKVUsage: admissionMaxKVUsage,
MaxPark: time.Duration(maxParkSeconds) * time.Second,
}
// Issue #64: only backends that actually publish the vLLM scheduler gauges
// may be observed. `reasoning` is the only vLLM engine here; ornith and
// qwen-cpu are Ollama and embeddings/rerank are TEI, none of which export
// vllm:* series. Observing them would trip the startup self-check and hold
// observer_degraded at 1 forever. Unset (the default) observes nothing and
// leaves admission exactly as it was before this change.
observeModels := make(map[string]bool)
for _, name := range strings.Split(os.Getenv("ADMISSION_OBSERVE_MODELS"), ",") {
if name = strings.TrimSpace(name); name != "" {
observeModels[name] = true
}
}
h.engineObservers = make(map[string]*EngineObserver)
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)
if len(addrs) == 0 {
continue
}
if observeModels[modelName] {
for _, addr := range addrs {
if _, ok := h.engineObservers[addr]; !ok {
h.engineObservers[addr] = NewEngineObserver(addr, enginePollInterval)
}
}
}
h.schedulers[modelName] = NewScheduler(addrs, maxSlots, maxQueue, h.engineObservers, limits)
}
// The handler has no request-path lifecycle to hang cancellation on, so the
// observers poll for the lifetime of the process.
for _, o := range h.engineObservers {
o.Start(context.Background())
}
// Pre-populate transports for all scheduler backends at startup (single-threaded).
// getOrCreateTransport is not goroutine-safe; pre-populating avoids concurrent writes.
@@ -148,7 +218,7 @@ func New(cfg *config.Config) *Handler {
// Create a transport per unique upstream address for connection reuse
transport := h.getOrCreateTransport(route.Upstream.Address, &route.Upstream)
upstreamURL, _ := url.Parse("http://" + route.Upstream.Address)
upstreamURL := &url.URL{Scheme: "http", Host: route.Upstream.Address}
r := &Route{
Name: name,
@@ -167,7 +237,16 @@ func New(cfg *config.Config) *Handler {
// getOrCreateTransport returns a shared http.Transport for the given upstream address.
// This ensures connections are pooled and reused across requests to the same upstream.
func (h *Handler) getOrCreateTransport(addr string, up *config.Upstream) *http.Transport {
h.transportsMu.RLock()
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
}
@@ -216,7 +295,7 @@ func (h *Handler) getOrCreateTransport(addr string, up *config.Upstream) *http.T
}
h.transports[addr] = transport
return transport
return transport // written under transportsMu write lock
}
// directorFunc modifies the request to be sent to the upstream.
@@ -302,6 +381,10 @@ func writeProblemDetail(w http.ResponseWriter, status int, problemType, title, d
// ServeHTTP implements http.Handler.
func (h *Handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
// Captured before any work so TTFT measures from the moment the gateway
// accepted the request, not from when the upstream response is copied.
requestStart := time.Now()
if h.handleAuthEndpoint(w, r) {
return
}
@@ -346,44 +429,137 @@ func (h *Handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
// For streaming responses (SSE, chunked), disable buffering. Issue #33.
w.Header().Set("X-Accel-Buffering", "no")
proxy.ServeHTTP(w, r)
// 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, and records
// the metric. It is a no-op when the header is absent or not "true".
func (h *Handler) applyWarmup(r *http.Request) {
// 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
return true
}
r.Header.Del("X-Poimen-Warmup")
r.Header.Set("X-Priority", "low")
time.Sleep(h.warmupDelay)
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) {
// to Release (caller must defer), the (possibly patched) route, whether SSE
// headers were already written while parking, and ok=true.
// If the slot cannot be acquired it writes a 503 and returns ok=false.
func (h *Handler) acquireScheduledBackend(ctx context.Context, w http.ResponseWriter, r *http.Request, sched *Scheduler, route *Route, cacheKey string) (*Backend, *Route, bool, bool) {
backend := sched.Select(cacheKey)
h.applyWarmup(r)
if backend == nil {
return nil, route, true
if !h.applyWarmup(ctx, r) {
return nil, nil, false, false
}
if err := backend.Acquire(ctx); err != nil {
if errors.Is(err, ErrQueueFull) {
if backend == nil {
return nil, route, false, true
}
// Decision Q22: a request that actually parks on the engine gate and wants
// an SSE stream gets its headers pre-written plus a 15 s keep-alive, so no
// intermediary times the connection out before the engine admits it.
var (
prewritten bool
hbStop chan struct{}
hbDone chan struct{}
)
onPark := func() {
if !requestWantsStream(r) {
return
}
flusher, ok := w.(http.Flusher)
if !ok {
return
}
prewritten = true
hbStop = make(chan struct{})
hbDone = make(chan struct{})
w.Header().Set("Content-Type", "text/event-stream")
w.Header().Set("Cache-Control", "no-cache")
w.Header().Set("Connection", "keep-alive")
w.Header().Set("X-Accel-Buffering", "no")
w.WriteHeader(http.StatusOK)
fmt.Fprint(w, ": keep-alive\n\n")
flusher.Flush()
go func() {
defer close(hbDone)
ticker := time.NewTicker(15 * time.Second)
defer ticker.Stop()
for {
select {
case <-hbStop:
return
case <-ctx.Done():
return
case <-ticker.C:
fmt.Fprint(w, ": keep-alive\n\n")
flusher.Flush()
}
}
}()
}
err := backend.acquire(ctx, onPark)
if prewritten {
close(hbStop)
<-hbDone
}
if err != nil {
// Once the SSE status line is spent, report the rejection as an event
// instead of a second WriteHeader.
reject := func(body string) {
if prewritten {
fmt.Fprintf(w, "event: error\ndata: {\"error\":%q}\n\n", body)
if f, ok := w.(http.Flusher); ok {
f.Flush()
}
return
}
w.WriteHeader(http.StatusServiceUnavailable)
fmt.Fprintf(w, "%s", body)
}
if errors.Is(err, ErrAdmissionTimeout) {
observability.Global().IncrAdmissionRejected("park_timeout")
logging.Warnf("admission timeout", map[string]string{
"backend": backend.Address,
"path": r.URL.Path,
"model": route.ModelName,
})
if !prewritten {
w.Header().Set("Retry-After", "5")
}
reject("admission timeout")
} else if errors.Is(err, ErrQueueFull) {
observability.Global().IncrQueueFull()
observability.Global().IncrAdmissionRejected("queue_full")
logging.Warnf("backend queue full", map[string]string{
"backend": backend.Address,
"path": r.URL.Path,
"model": route.ModelName,
})
w.WriteHeader(http.StatusServiceUnavailable)
fmt.Fprintf(w, "queue full")
if !prewritten {
w.Header().Set("Retry-After", "5")
}
reject("queue full")
} else {
logging.Warnf("backend acquire failed", map[string]string{
"backend": backend.Address,
@@ -391,15 +567,14 @@ func (h *Handler) acquireScheduledBackend(ctx context.Context, w http.ResponseWr
"model": route.ModelName,
"err": err.Error(),
})
w.WriteHeader(http.StatusServiceUnavailable)
fmt.Fprintf(w, "upstream busy")
reject("upstream busy")
}
return nil, nil, false
return nil, nil, prewritten, false
}
if backend.Address != route.Upstream.Address {
patchedCfg := *route.Upstream
patchedCfg.Address = backend.Address
patchedURL, _ := url.Parse("http://" + backend.Address)
patchedURL := &url.URL{Scheme: "http", Host: backend.Address}
route = &Route{
Name: route.Name,
ModelName: route.ModelName,
@@ -410,7 +585,27 @@ func (h *Handler) acquireScheduledBackend(ctx context.Context, w http.ResponseWr
},
}
}
return backend, route, true
return backend, route, prewritten, true
}
// requestWantsStream reports whether the JSON request body asks for an SSE
// stream ("stream": true). The body is restored for the upstream.
func requestWantsStream(r *http.Request) bool {
if r.Body == nil {
return false
}
body, err := io.ReadAll(r.Body)
if err != nil {
return false
}
r.Body = io.NopCloser(bytes.NewReader(body))
var probe struct {
Stream bool `json:"stream"`
}
if err := json.Unmarshal(body, &probe); err != nil {
return false
}
return probe.Stream
}
// handleAuthEndpoint routes /auth/* requests and returns true if handled.
@@ -573,9 +768,10 @@ func (h *Handler) enforceBodyLimit(w http.ResponseWriter, r *http.Request, route
// Returns (nil, nil, false) when the backend slot cannot be acquired (response written).
func (h *Handler) setupCAGProxy(ctx context.Context, w http.ResponseWriter, r *http.Request, route *Route, sched *Scheduler, cacheKey string) (*httputil.ReverseProxy, *Backend, bool) {
var selectedBackend *Backend
var headersPrewritten bool
if sched != nil {
var ok bool
selectedBackend, route, ok = h.acquireScheduledBackend(ctx, w, r, sched, route, cacheKey)
selectedBackend, route, headersPrewritten, ok = h.acquireScheduledBackend(ctx, w, r, sched, route, cacheKey)
if !ok {
return nil, nil, false
}
@@ -590,12 +786,52 @@ func (h *Handler) setupCAGProxy(ctx context.Context, w http.ResponseWriter, r *h
capturedBackend := selectedBackend
capturedRoute := route
if headersPrewritten {
// The 200 + text/event-stream status is already on the wire, so an
// upstream status can no longer be forwarded: httputil.ReverseProxy
// calls WriteHeader again and net/http keeps the first one. Frame a
// failure as an SSE event rather than letting a raw error body land
// inside the stream, which a streaming client reads as an empty
// completion and reports as success.
proxy.ModifyResponse = func(resp *http.Response) error {
if resp.StatusCode < http.StatusBadRequest {
return nil
}
body, _ := io.ReadAll(io.LimitReader(resp.Body, 4096))
resp.Body.Close()
logging.Warnf("upstream error after parked stream started", map[string]string{
"upstream": capturedRoute.Upstream.Address,
"status": strconv.Itoa(resp.StatusCode),
})
resp.Body = io.NopCloser(strings.NewReader(fmt.Sprintf(
"event: error\ndata: {\"upstream_status\":%d,\"error\":%q}\n\n",
resp.StatusCode, strings.TrimSpace(string(body)))))
resp.ContentLength = -1
resp.Header.Del("Content-Length")
resp.Header.Set("Content-Type", "text/event-stream")
return nil
}
}
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)
}
}
if headersPrewritten {
// SSE headers and a 200 are already on the wire: a 502 here would be
// an ignored WriteHeader and would tear the event stream. Report the
// failure as an event instead.
fmt.Fprint(w, "event: error\ndata: {\"error\":\"upstream error\"}\n\n")
if f, ok := w.(http.Flusher); ok {
f.Flush()
}
logging.Errorf("upstream error", err, map[string]string{
"upstream": capturedRoute.Upstream.Address,
"path": r.URL.Path,
})
return
}
w.WriteHeader(http.StatusBadGateway)
fmt.Fprintf(w, "upstream error")
logging.Errorf("upstream error", err, map[string]string{
+112 -12
View File
@@ -69,9 +69,10 @@ func (h *Handler) routeByModel(r *http.Request, path string) (*Route, error) {
// Restore the body so it can be read again by the upstream
r.Body = io.NopCloser(bytes.NewReader(bodyBytes))
// Parse the JSON to find the model field
var payload map[string]interface{}
if err := json.Unmarshal(bodyBytes, &payload); err != nil {
// Extract only the top-level "model" value with a token walk: the common
// path must not materialise the whole body (messages and all) as a map.
modelRaw, hasModel, err := extractTopLevelModel(bodyBytes)
if err != nil {
return nil, &modelValidationError{
Kind: "invalid_json",
Message: "request body is not valid JSON",
@@ -79,7 +80,6 @@ func (h *Handler) routeByModel(r *http.Request, path string) (*Route, error) {
}
// Extract the model name
modelVal, hasModel := payload["model"]
if !hasModel {
return nil, &modelValidationError{
Kind: "missing_model",
@@ -88,7 +88,7 @@ func (h *Handler) routeByModel(r *http.Request, path string) (*Route, error) {
}
// Handle null model
if modelVal == nil {
if bytes.Equal(bytes.TrimSpace(modelRaw), []byte("null")) {
return nil, &modelValidationError{
Kind: "null_model",
Message: "'model' field is null",
@@ -96,8 +96,8 @@ func (h *Handler) routeByModel(r *http.Request, path string) (*Route, error) {
}
// Extract as string
modelName, ok := modelVal.(string)
if !ok {
var modelName string
if err := json.Unmarshal(modelRaw, &modelName); err != nil {
return nil, &modelValidationError{
Kind: "missing_model",
Message: "'model' field must be a string",
@@ -122,10 +122,26 @@ func (h *Handler) routeByModel(r *http.Request, path string) (*Route, error) {
}
}
// If upstream expects a different model name, rewrite the body
// If upstream expects a different model name, rewrite the body. This path is
// uncommon, so it may decode the body as raw fields; the common path above
// does not build a full map.
if modelUpstream.UpstreamModel != "" && modelUpstream.UpstreamModel != modelName {
payload["model"] = modelUpstream.UpstreamModel
newBody, err := json.Marshal(payload)
var fields map[string]json.RawMessage
if err := json.Unmarshal(bodyBytes, &fields); err != nil {
return nil, &modelValidationError{
Kind: "invalid_request",
Message: fmt.Sprintf("failed to rewrite model name: %v", err),
}
}
newModel, err := json.Marshal(modelUpstream.UpstreamModel)
if err != nil {
return nil, &modelValidationError{
Kind: "invalid_request",
Message: fmt.Sprintf("failed to rewrite model name: %v", err),
}
}
fields["model"] = newModel
newBody, err := json.Marshal(fields)
if err != nil {
return nil, &modelValidationError{
Kind: "invalid_request",
@@ -151,11 +167,11 @@ func (h *Handler) routeByModel(r *http.Request, path string) (*Route, error) {
ConnectTimeout: h.defaultConnectTimeout,
ReadTimeout: h.defaultReadTimeout,
WriteTimeout: h.defaultWriteTimeout,
MaxBodySize: h.defaultMaxBodySize,
MaxBodySize: h.maxBodyBytes,
AuthRequired: false,
}
targetURL, _ := url.Parse("http://" + modelUpstream.Address)
targetURL := &url.URL{Scheme: "http", Host: modelUpstream.Address}
route := &Route{
Name: "v1-chat-" + modelName,
@@ -169,3 +185,87 @@ func (h *Handler) routeByModel(r *http.Request, path string) (*Route, error) {
return route, nil
}
// extractTopLevelModel walks a JSON object body with a streaming decoder and
// returns the raw value of the top-level "model" field. It never builds the
// full document as a map, so a large "messages" array is skipped by token.
// hasModel is false when the key is absent or the document is a bare null.
func extractTopLevelModel(body []byte) (model json.RawMessage, hasModel bool, err error) {
dec := json.NewDecoder(bytes.NewReader(body))
tok, err := dec.Token()
if err != nil {
return nil, false, err
}
if tok == nil {
// A bare `null` body decoded to a nil map before, i.e. no model field.
if _, err := dec.Token(); err != io.EOF {
if err == nil {
return nil, false, fmt.Errorf("unexpected data after JSON null")
}
return nil, false, err
}
return nil, false, nil
}
if delim, ok := tok.(json.Delim); !ok || delim != '{' {
return nil, false, fmt.Errorf("request body is not a JSON object")
}
for dec.More() {
keyTok, err := dec.Token()
if err != nil {
return nil, false, err
}
key, _ := keyTok.(string)
if key == "model" {
if err := dec.Decode(&model); err != nil {
return nil, false, err
}
hasModel = true
continue
}
if err := skipJSONValue(dec); err != nil {
return nil, false, err
}
}
// Consume the closing brace and reject trailing data, preserving the
// invalid_json answer the old whole-body decode produced.
if _, err := dec.Token(); err != nil {
return nil, false, err
}
if _, err := dec.Token(); err != io.EOF {
if err == nil {
return nil, false, fmt.Errorf("unexpected data after JSON object")
}
return nil, false, err
}
return model, hasModel, nil
}
// skipJSONValue advances dec past the next JSON value without materialising it.
func skipJSONValue(dec *json.Decoder) error {
tok, err := dec.Token()
if err != nil {
return err
}
delim, ok := tok.(json.Delim)
if !ok {
return nil // scalar
}
if delim != '{' && delim != '[' {
return fmt.Errorf("unexpected delimiter %v", delim)
}
depth := 1
for depth > 0 {
tok, err := dec.Token()
if err != nil {
return err
}
if d, ok := tok.(json.Delim); ok {
if d == '{' || d == '[' {
depth++
} else {
depth--
}
}
}
return nil
}
+71 -4
View File
@@ -8,20 +8,41 @@ import (
"sync"
"sync/atomic"
"time"
"forgejo.riotpiao.com/rock/homelab-frontend/internal/observability"
)
// ErrQueueFull is returned when the in-memory queue is at capacity.
var ErrQueueFull = errors.New("backend queue full")
// ErrAdmissionTimeout is returned when a request parked on engine admission
// exceeds the configured maximum park duration.
var ErrAdmissionTimeout = errors.New("admission timeout")
// AdmissionLimits are the engine-wide admission thresholds applied to every
// backend of a scheduler (#64).
type AdmissionLimits struct {
MaxInFlight int
MaxKVUsage float64
MaxPark time.Duration
}
// 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)
//
// The engine gate is applied before the queue: admission reads the engine's
// own published gauges via observer, which is global across replicas.
type Backend struct {
Address string
queue chan struct{}
slots chan struct{}
Address string
queue chan struct{}
slots chan struct{}
observer *EngineObserver
maxInFlight int
maxKVUsage float64
maxPark time.Duration
}
// newBackend creates a Backend with given queue depth and concurrency limit.
@@ -40,7 +61,48 @@ func newBackend(addr string, maxSlots, maxQueue int) *Backend {
// 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.
// Returns ErrAdmissionTimeout if the engine gate parks the request past maxPark.
func (b *Backend) Acquire(ctx context.Context) error {
return b.acquire(ctx, nil)
}
// acquire is Acquire with an optional onPark callback. onPark is invoked once,
// on the calling goroutine, the first time the request has to wait for engine
// admission; the parked streaming path uses it to start an SSE heartbeat.
func (b *Backend) acquire(ctx context.Context, onPark func()) error {
// Engine gate (#64): admit on the engine's own gauges, not on this pod's
// counters. Failing OPEN on !Healthy is deliberate — a metrics outage must
// not become an inference outage; the three-replica gate is only a fairness
// improvement on top of the engine's own queueing.
if b.observer != nil {
deadline := time.Now().Add(b.maxPark)
timer := time.NewTimer(time.Until(deadline))
defer timer.Stop()
parked := false
for {
// State and tick channel are read together under one lock: taking
// them separately can miss a poll that lands between the two calls.
st, tick := b.observer.Snapshot()
if !st.Healthy || (st.InFlight() < int64(b.maxInFlight) && st.KVUsage < b.maxKVUsage) {
break
}
if !parked {
observability.Global().IncrAdmissionParked()
parked = true
if onPark != nil {
onPark()
}
}
select {
case <-tick:
case <-ctx.Done():
return ctx.Err()
case <-timer.C:
return ErrAdmissionTimeout
}
}
}
// Non-blocking queue entry: reject immediately when queue is full.
select {
case b.queue <- struct{}{}:
@@ -76,10 +138,15 @@ type Scheduler struct {
// 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 {
// observers holds one engine observer per backend address (may be nil).
func NewScheduler(addresses []string, maxSlots, maxQueue int, observers map[string]*EngineObserver, limits AdmissionLimits) *Scheduler {
backends := make([]*Backend, len(addresses))
for i, addr := range addresses {
backends[i] = newBackend(addr, maxSlots, maxQueue)
backends[i].observer = observers[addr]
backends[i].maxInFlight = limits.MaxInFlight
backends[i].maxKVUsage = limits.MaxKVUsage
backends[i].maxPark = limits.MaxPark
}
return &Scheduler{
backends: backends,
+3 -3
View File
@@ -67,7 +67,7 @@ func TestBackendAcquireReturnsErrQueueFullWhenQueueFull(t *testing.T) {
// 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)
s := NewScheduler(addrs, 10, 10, nil, AdmissionLimits{})
first := s.Select("same-cache-key")
if first == nil {
@@ -85,7 +85,7 @@ func TestSchedulerSelectRendezvousHashDeterministic(t *testing.T) {
// 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)
s := NewScheduler(addrs, 10, 10, nil, AdmissionLimits{})
counts := make(map[string]int)
for i := 0; i < 9; i++ {
@@ -160,7 +160,7 @@ func TestBackendAcquireCancelledContext(t *testing.T) {
// 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 := NewScheduler(addrs, 10, 10, nil, AdmissionLimits{})
s.MarkFailed("a:80")
+3
View File
@@ -28,6 +28,9 @@ func AdapterFactory(serviceName string) interface{} {
return &IAMAdapter{}
case "memory":
return &MemoryAdapter{}
case "sessions":
// code-registered; see NewSessionAdapter()
return nil
default:
return nil
}
+8 -5
View File
@@ -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",
Methods: []Method{
{
Verb: "GET",
UpstreamPath: "/v1/poimen/executions/{id}",
UpstreamPath: "/v1/poimen/executions/detail",
Auth: &Auth{Required: true, Capability: "workflow:read"},
},
{
Verb: "DELETE",
UpstreamPath: "/v1/poimen/executions/{id}",
UpstreamPath: "/v1/poimen/executions/detail",
Auth: &Auth{Required: true, Capability: "workflow:execute"},
},
},
@@ -72,7 +75,7 @@ func GetPoimenSpec() *Spec {
Name: "execution-history",
Methods: []Method{{
Verb: "GET",
UpstreamPath: "/v1/poimen/executions/{id}/history",
UpstreamPath: "/v1/poimen/executions/history",
Auth: &Auth{Required: true, Capability: "workflow:read"},
}},
},
@@ -80,7 +83,7 @@ func GetPoimenSpec() *Spec {
Name: "execution-stream",
Methods: []Method{{
Verb: "GET",
UpstreamPath: "/v1/poimen/executions/{id}/stream",
UpstreamPath: "/v1/poimen/executions/stream",
Auth: &Auth{Required: true, Capability: "workflow:read"},
}},
},
@@ -88,7 +91,7 @@ func GetPoimenSpec() *Spec {
Name: "execution-signal",
Methods: []Method{{
Verb: "POST",
UpstreamPath: "/v1/poimen/executions/{id}/signal",
UpstreamPath: "/v1/poimen/executions/signal",
Auth: &Auth{Required: true, Capability: "workflow:execute"},
}},
},
+123
View File
@@ -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"},
},
},
},
},
}
}
+97
View File
@@ -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
}
+634
View File
@@ -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
}
+594
View File
@@ -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)
}
}
+492
View File
@@ -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
}
+271
View File
@@ -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
View File
@@ -4,7 +4,6 @@ import (
"encoding/json"
"fmt"
"net/http"
"strings"
"time"
)
@@ -33,26 +32,29 @@ func (p *PoimenHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
p.handleExecute(w, r)
case path == "/v1/poimen/executions" && r.Method == http.MethodGet:
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)
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)
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)
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)
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)
default:
writeJSON(w, http.StatusNotFound, map[string]string{"error": "endpoint not found"})
}
}
// extractWorkflowID extracts the workflow ID from paths like /v1/poimen/executions/{id}/...
func extractWorkflowID(path string) string {
parts := strings.Split(strings.TrimPrefix(path, "/v1/poimen/executions/"), "/")
if len(parts) > 0 && parts[0] != "" {
return parts[0]
// getWorkflowID reads the workflow ID from query param (preferred) or path.
// The ServiceAdapter sets X-Upstream-Path to a template like
// /v1/poimen/executions/{id}/stream — the {id} is literal, not substituted.
// So workflow_id must come from a query parameter.
func getWorkflowID(r *http.Request) string {
// Query param is the primary source
if id := r.URL.Query().Get("workflow_id"); id != "" {
return id
}
return ""
}
@@ -176,13 +178,9 @@ func (p *PoimenHandler) handleDescribe(w http.ResponseWriter, r *http.Request) {
return
}
path := r.Header.Get("X-Upstream-Path")
if path == "" {
path = r.URL.Path
}
workflowID := extractWorkflowID(path)
workflowID := getWorkflowID(r)
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
}
@@ -209,13 +207,9 @@ func (p *PoimenHandler) handleHistory(w http.ResponseWriter, r *http.Request) {
return
}
path := r.Header.Get("X-Upstream-Path")
if path == "" {
path = r.URL.Path
}
workflowID := extractWorkflowID(path)
workflowID := getWorkflowID(r)
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
}
@@ -248,15 +242,9 @@ func (p *PoimenHandler) handleSignal(w http.ResponseWriter, r *http.Request) {
return
}
path := r.Header.Get("X-Upstream-Path")
if path == "" {
path = r.URL.Path
}
// Remove /signal suffix to get workflow ID
cleanPath := strings.TrimSuffix(path, "/signal")
workflowID := extractWorkflowID(cleanPath)
workflowID := getWorkflowID(r)
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
}
@@ -290,13 +278,9 @@ func (p *PoimenHandler) handleTerminate(w http.ResponseWriter, r *http.Request)
return
}
path := r.Header.Get("X-Upstream-Path")
if path == "" {
path = r.URL.Path
}
workflowID := extractWorkflowID(path)
workflowID := getWorkflowID(r)
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
}
@@ -317,15 +301,9 @@ func (p *PoimenHandler) handleTerminate(w http.ResponseWriter, r *http.Request)
// --- Stream (delegate to existing SSE handler) ---
func (p *PoimenHandler) handleStream(w http.ResponseWriter, r *http.Request) {
path := r.Header.Get("X-Upstream-Path")
if path == "" {
path = r.URL.Path
}
// Remove /stream suffix
cleanPath := strings.TrimSuffix(path, "/stream")
workflowID := extractWorkflowID(cleanPath)
workflowID := getWorkflowID(r)
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
}
+15
View File
@@ -36,6 +36,21 @@ func (rw *responseWriter) Write(b []byte) (int, error) {
return n, err
}
// Flush forwards to the wrapped writer so that streaming handlers can still
// detect http.Flusher through this wrapper. Without it a bare
// `w.(http.Flusher)` assertion fails for every request that passes through
// this middleware, which silently disables SSE flushing and keep-alives.
func (rw *responseWriter) Flush() {
if f, ok := rw.ResponseWriter.(http.Flusher); ok {
f.Flush()
}
}
// Unwrap exposes the wrapped writer to http.ResponseController.
func (rw *responseWriter) Unwrap() http.ResponseWriter {
return rw.ResponseWriter
}
// Middleware returns an HTTP middleware that adds tracing to requests.
func Middleware(next http.Handler) http.Handler {
tracer := otel.Tracer(tracerName)
+93 -9
View File
@@ -11,6 +11,7 @@ import (
"net/http"
"os"
"strings"
"time"
"forgejo.riotpiao.com/rock/homelab-frontend/internal/notification"
)
@@ -18,8 +19,9 @@ import (
// ForgejoHandler receives Forgejo webhook payloads and forwards them to Gotify.
// Forgejo sends a Gitea-compatible JSON payload with X-Gitea-Signature-256 header.
type ForgejoHandler struct {
secret string
gotify *notification.GotifyClient
secret string
gotify *notification.GotifyClient
promoter *ImagePromoter
}
// NewForgejoHandler creates a handler from environment variables.
@@ -35,8 +37,9 @@ func NewForgejoHandler() *ForgejoHandler {
}
return &ForgejoHandler{
secret: os.Getenv("FORGEJO_WEBHOOK_SECRET"),
gotify: client,
secret: os.Getenv("FORGEJO_WEBHOOK_SECRET"),
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 {
log.Printf("forgejo webhook received but Gotify not configured (GOTIFY_URL/GOTIFY_APP_TOKEN missing)")
w.WriteHeader(http.StatusOK)
return
}
event := r.Header.Get("X-Gitea-Event")
if event == "" {
event = r.Header.Get("X-GitHub-Event")
}
title, message, priority := h.formatMessage(event, body)
if title == "" {
// Unhandled event type — ack and ignore
@@ -99,6 +108,81 @@ func (h *ForgejoHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
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>
func (h *ForgejoHandler) verifySignature(body []byte, sig string) bool {
sig = strings.TrimPrefix(sig, "sha256=")
+424
View File
@@ -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
}
+408
View File
@@ -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)
}
File diff suppressed because one or more lines are too long
+30 -3
View File
@@ -21,9 +21,10 @@ spec:
labels:
app: api-gateway
component: gateway
# Required by the llm-serving-default-deny NetworkPolicy, which admits
# only pods labelled llm-client=true (from any namespace) on port 8080.
# Without it every upstream dial times out and dispatch returns 502.
# The llm-serving-ingress NetworkPolicy in the llm-serving namespace now
# selects by namespace (kubernetes.io/metadata.name: api), not by this
# label. Kept for older NetworkPolicy revisions that still admit on
# llm-client=true.
llm-client: "true"
annotations:
reloader.stakater.com/auto: "true"
@@ -67,6 +68,24 @@ spec:
value: "5m"
- name: LOG_LEVEL
value: "info"
- name: ADMISSION_MAX_PARK_SECONDS
value: "120"
# Must track the engine's --max-num-seqs (homelab: reasoning.yaml).
# The gauge this compares against is engine-wide, so this number is a
# ceiling on concurrent requests across all gateway replicas, not a
# per-pod slot count. Drift here is silent; startup logs the value.
- name: ADMISSION_MAX_INFLIGHT
value: "4"
# Only backends that publish the vLLM scheduler gauges may be observed.
# `reasoning` is the only vLLM engine; ornith/qwen-cpu are Ollama and
# embeddings/rerank are TEI, so observing them would hold
# observer_degraded at 1 forever.
- name: ADMISSION_OBSERVE_MODELS
value: "reasoning"
- name: VLLM_MAX_QUEUE_DEPTH
value: "8"
- name: MAX_BODY_BYTES
value: "1048576"
# OpenTelemetry tracing configuration
- name: OTEL_EXPORTER_OTLP_ENDPOINT
value: "otel-collector.tracing.svc.cluster.local:4317"
@@ -95,6 +114,14 @@ spec:
name: gotify-webhook-secret
key: forgejo-webhook-secret
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:
- name: config
mountPath: /etc/gateway
+1
View File
@@ -9,6 +9,7 @@ resources:
- deployment.yaml
- network-policy.yaml
- gateway-config-secret.yaml
- rbac.yaml
# 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.
+44 -3
View File
@@ -21,11 +21,17 @@ spec:
ports:
- protocol: TCP
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:
- namespaceSelector:
matchLabels:
name: monitoring
kubernetes.io/metadata.name: monitoring
ports:
- protocol: TCP
port: 8080
@@ -122,7 +128,7 @@ spec:
ports:
- protocol: TCP
port: 7233
# Allow to Poimen (memory/semantic search)
# Allow to Poimen (memory:8080 + graphql:8090 + redis:6379 + fleet-controller:8788)
- to:
- namespaceSelector:
matchLabels:
@@ -130,6 +136,20 @@ spec:
ports:
- protocol: TCP
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)
# Service `minio` listens on port 80 (targetPort 9000).
# Headless `minio-cluster-hl` is 9000. Allow both.
@@ -142,3 +162,24 @@ spec:
port: 80
- protocol: TCP
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
+62
View File
@@ -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