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
22 changed files with 2303 additions and 183 deletions
+58
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
@@ -177,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")
}
}
+236 -21
View File
@@ -2,6 +2,7 @@
package proxy
import (
"bytes"
"context"
"encoding/json"
"errors"
@@ -47,6 +48,10 @@ 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
}
@@ -122,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.
@@ -151,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,
@@ -314,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
}
@@ -358,7 +429,14 @@ 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)
}
@@ -384,27 +462,104 @@ func (h *Handler) applyWarmup(ctx context.Context, r *http.Request) bool {
// 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)
if !h.applyWarmup(ctx, r) {
return nil, nil, false
return nil, nil, false, false
}
if backend == nil {
return nil, route, true
return nil, route, false, true
}
if err := backend.Acquire(ctx); err != nil {
if errors.Is(err, ErrQueueFull) {
// 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.Header().Set("Retry-After", "5")
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,
@@ -412,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,
@@ -431,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.
@@ -594,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
}
@@ -611,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")
@@ -102,6 +102,22 @@ func GetSessionSpec() *Spec {
},
},
},
// 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
}
+126 -3
View File
@@ -12,9 +12,11 @@ package sessions
import (
"bytes"
"context"
"crypto/sha256"
"encoding/json"
"fmt"
"log"
"net/http"
"os"
"time"
@@ -30,9 +32,13 @@ const (
// 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.
@@ -42,10 +48,34 @@ func NewHandler() *Handler {
if gql == "" {
gql = defaultGraphQLURL
}
return &Handler{
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.
@@ -66,6 +96,8 @@ func (h *Handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
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"})
}
@@ -180,9 +212,12 @@ func (h *Handler) handleList(w http.ResponseWriter, r *http.Request) {
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!){
hybridSearch(query:$q,limit:$limit,entityType:SESSION){
entities{id name summary tCreated}
entities(query:$q,limit:$limit,entityType:SESSION){
entities{id name summary projectId tCreated}
searchMode
}
}`
@@ -495,6 +530,94 @@ func (h *Handler) gql(r *http.Request, query string, variables map[string]any) (
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.
+6 -1
View File
@@ -18,7 +18,8 @@ import (
// "query" field and records captured request bodies.
//
// - "upsertEntity" → mutation response
// - "hybridSearch" → search 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()
@@ -33,7 +34,11 @@ func newGQLServer(t *testing.T, captured *[]byte) *httptest.Server {
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:
+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")
}
}
+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)
+4 -1
View File
@@ -260,9 +260,12 @@ func TestPromote_KustomizeSuccess(t *testing.T) {
case http.MethodPatch:
atomic.AddInt32(&patchCalled, 1)
raw, _ := readAll(r.Body)
if !strings.Contains(string(raw), "/spec/sources/0/kustomize/images") {
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)
}
File diff suppressed because one or more lines are too long
+22 -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"
+15 -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,12 @@ 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:
+34
View File
@@ -26,3 +26,37 @@ 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