diff --git a/gateway/internal/adminapi/budgets_test.go b/gateway/internal/adminapi/budgets_test.go index 93e3d13e0..6257f52a6 100644 --- a/gateway/internal/adminapi/budgets_test.go +++ b/gateway/internal/adminapi/budgets_test.go @@ -146,12 +146,11 @@ func TestBudget_HourlyWindow(t *testing.T) { func TestBudget_404OnDeepPath(t *testing.T) { srv, _ := newBudgetTestServer(t, nil) - // phase-9 will own /:name/state and /:name/kill; phase-8 only - // owns /:name/budget. Anything else under /_plugin/agents/ - // should 404. + // /:name/{budget,state,kill,catalog,tools,skills,evals} are + // owned; anything else under /_plugin/agents/ should 404. for _, p := range []string{ "/_plugin/agents/coder", - "/_plugin/agents/coder/state", + "/_plugin/agents/coder/other", "/_plugin/agents/coder/budget/extra", "/_plugin/agents//budget", } { diff --git a/gateway/internal/adminapi/hotstate.go b/gateway/internal/adminapi/hotstate.go new file mode 100644 index 000000000..8fa646bb9 --- /dev/null +++ b/gateway/internal/adminapi/hotstate.go @@ -0,0 +1,172 @@ +package adminapi + +import ( + "errors" + "net/http" + "time" + + "github.com/stakwork/stakgraph/gateway/internal/auth" + "github.com/stakwork/stakgraph/gateway/internal/pluginlog" +) + +// Phase-6 hot-state routes: kill switches and Redis snapshots for +// runs and agents. Thin wrappers over auth.KillRun / auth.GetRunState +// et al. — every decision about keys, TTLs and what "killed" means +// lives in gateway/internal/auth/kill.go; this file only parses the +// URL, picks by method, and shapes the JSON. +// +// POST /_plugin/runs/:id/kill → 200 KillRunResponse +// DELETE /_plugin/runs/:id/kill → 204 +// GET /_plugin/runs/:id/state → 200 RunStateResponse +// POST /_plugin/agents/:name/kill → 200 KillAgentResponse +// DELETE /_plugin/agents/:name/kill → 204 +// GET /_plugin/agents/:name/state → 200 AgentStateResponse (?window=1d) +// +// All cookie-or-bearer: the operator dashboard drives these from a +// session cookie (with the CSRF header on mutations); Hive uses the +// bearer. Redis unconfigured ⇒ 503, same as the session store. + +// KillRunResponse is the wire shape for POST /_plugin/runs/:id/kill. +type KillRunResponse struct { + RunID string `json:"run_id"` + KilledAt string `json:"killed_at"` // RFC3339 UTC +} + +// KillAgentResponse is the wire shape for POST /_plugin/agents/:name/kill. +type KillAgentResponse struct { + AgentName string `json:"agent_name"` + KilledAt string `json:"killed_at"` // RFC3339 UTC +} + +// RunStateResponse is the wire shape for GET /_plugin/runs/:id/state — +// the run's live phase-6 accumulators. A run that has never made a +// call reads as all-zero with ttl_seconds = -2 (no key), not 404. +type RunStateResponse struct { + RunID string `json:"run_id"` + CostUSD float64 `json:"cost_usd"` + Steps int64 `json:"steps"` + Tools []string `json:"tools"` // last 10 tool names, most recent first + Killed bool `json:"killed"` + // TTLSeconds is the remaining lifetime of the cost accumulator: + // -2 when the run has no state yet, -1 when it has no expiry. + TTLSeconds int64 `json:"ttl_seconds"` +} + +// AgentStateResponse is the wire shape for GET /_plugin/agents/:name/state. +type AgentStateResponse struct { + AgentName string `json:"agent_name"` + Window string `json:"window"` + BucketKey string `json:"bucket_key"` + CurrentSpendUSD float64 `json:"current_spend_usd"` + // ConfiguredCapUSD is null when the agent has no agent_budgets + // entry; then `window` is informational (?window= or "1d"). + ConfiguredCapUSD *float64 `json:"configured_cap_usd"` + Killed bool `json:"killed"` +} + +type hotStateHandlers struct{} + +func newHotStateHandlers() *hotStateHandlers { return &hotStateHandlers{} } + +func (h *hotStateHandlers) runKill(w http.ResponseWriter, r *http.Request, runID string) { + switch r.Method { + case http.MethodPost: + if err := auth.KillRun(r.Context(), runID); err != nil { + writeHotStateErr(w, err, "runs.kill") + return + } + pluginlog.Logf("adminapi: run killed run_id=%s", runID) + writeJSON(w, http.StatusOK, KillRunResponse{ + RunID: runID, + KilledAt: time.Now().UTC().Format(time.RFC3339), + }) + case http.MethodDelete: + if err := auth.UnkillRun(r.Context(), runID); err != nil { + writeHotStateErr(w, err, "runs.unkill") + return + } + pluginlog.Logf("adminapi: run unkilled run_id=%s", runID) + w.WriteHeader(http.StatusNoContent) + default: + methodNotAllowed(w, http.MethodPost, http.MethodDelete) + } +} + +func (h *hotStateHandlers) runState(w http.ResponseWriter, r *http.Request, runID string) { + if r.Method != http.MethodGet { + methodNotAllowed(w, http.MethodGet) + return + } + st, err := auth.GetRunState(r.Context(), runID) + if err != nil { + writeHotStateErr(w, err, "runs.state") + return + } + writeJSON(w, http.StatusOK, RunStateResponse{ + RunID: st.RunID, + CostUSD: st.CostUSD, + Steps: st.Steps, + Tools: st.Tools, + Killed: st.Killed, + TTLSeconds: st.TTLSeconds, + }) +} + +func (h *hotStateHandlers) agentKill(w http.ResponseWriter, r *http.Request, name string) { + switch r.Method { + case http.MethodPost: + if err := auth.KillAgent(r.Context(), name); err != nil { + writeHotStateErr(w, err, "agents.kill") + return + } + pluginlog.Logf("adminapi: agent killed agent=%s", name) + writeJSON(w, http.StatusOK, KillAgentResponse{ + AgentName: name, + KilledAt: time.Now().UTC().Format(time.RFC3339), + }) + case http.MethodDelete: + if err := auth.UnkillAgent(r.Context(), name); err != nil { + writeHotStateErr(w, err, "agents.unkill") + return + } + pluginlog.Logf("adminapi: agent unkilled agent=%s", name) + w.WriteHeader(http.StatusNoContent) + default: + methodNotAllowed(w, http.MethodPost, http.MethodDelete) + } +} + +func (h *hotStateHandlers) agentState(w http.ResponseWriter, r *http.Request, name string) { + if r.Method != http.MethodGet { + methodNotAllowed(w, http.MethodGet) + return + } + st, err := auth.GetAgentState(r.Context(), name, r.URL.Query().Get("window"), time.Now().UTC()) + if err != nil { + writeHotStateErr(w, err, "agents.state") + return + } + writeJSON(w, http.StatusOK, AgentStateResponse{ + AgentName: st.AgentName, + Window: st.Window, + BucketKey: st.BucketKey, + CurrentSpendUSD: st.CurrentSpendUSD, + ConfiguredCapUSD: st.ConfiguredCapUSD, + Killed: st.Killed, + }) +} + +// writeHotStateErr maps auth-package errors onto HTTP: Redis +// unconfigured ⇒ 503 (operator can retry once the link is up), +// anything else ⇒ 400. The auth helpers only return validation +// errors and Redis errors; a Redis I/O failure also lands on 400 +// here rather than 500 because the message is safe to show and +// retrying is the right move either way. +func writeHotStateErr(w http.ResponseWriter, err error, op string) { + if errors.Is(err, auth.ErrRedisUnavailable) { + http.Error(w, "redis not configured", http.StatusServiceUnavailable) + return + } + pluginlog.Warnf("adminapi: %s: %v", op, err) + http.Error(w, err.Error(), http.StatusBadRequest) +} diff --git a/gateway/internal/adminapi/hotstate_test.go b/gateway/internal/adminapi/hotstate_test.go new file mode 100644 index 000000000..254e69a6c --- /dev/null +++ b/gateway/internal/adminapi/hotstate_test.go @@ -0,0 +1,366 @@ +package adminapi + +import ( + "encoding/json" + "io" + "net/http" + "net/http/httptest" + "strings" + "testing" + "time" + + "github.com/stakwork/stakgraph/gateway/internal/auth" + "github.com/stakwork/stakgraph/gateway/internal/redisclient" +) + +// bearerDo issues an arbitrary-method request with the provisioning +// bearer. Body is optional JSON. +func bearerDo(t *testing.T, srv *httptest.Server, method, path, body string) *http.Response { + t.Helper() + var rd io.Reader + if body != "" { + rd = strings.NewReader(body) + } + req, _ := http.NewRequest(method, srv.URL+path, rd) + req.Header.Set("Authorization", "Bearer "+testToken) + if body != "" { + req.Header.Set("Content-Type", "application/json") + } + resp, err := srv.Client().Do(req) + if err != nil { + t.Fatal(err) + } + return resp +} + +func decodeBody(t *testing.T, resp *http.Response, into any) { + t.Helper() + defer resp.Body.Close() + if err := json.NewDecoder(resp.Body).Decode(into); err != nil { + t.Fatalf("decode: %v", err) + } +} + +func TestRunKill_RoundTrip(t *testing.T) { + srv, mr := newBudgetTestServer(t, nil) + + resp := bearerDo(t, srv, http.MethodPost, "/_plugin/runs/r_1/kill", "") + if resp.StatusCode != http.StatusOK { + t.Fatalf("kill: want 200, got %d", resp.StatusCode) + } + var kr KillRunResponse + decodeBody(t, resp, &kr) + if kr.RunID != "r_1" || kr.KilledAt == "" { + t.Fatalf("kill response: %+v", kr) + } + if !mr.Exists("bifrost:kill:r_1") { + t.Fatal("kill key not written") + } + if ttl := mr.TTL("bifrost:kill:r_1"); ttl != time.Hour { + t.Fatalf("ttl: %v", ttl) + } + + resp = bearerDo(t, srv, http.MethodGet, "/_plugin/runs/r_1/state", "") + if resp.StatusCode != http.StatusOK { + t.Fatalf("state: want 200, got %d", resp.StatusCode) + } + var st RunStateResponse + decodeBody(t, resp, &st) + if !st.Killed || st.RunID != "r_1" { + t.Fatalf("state after kill: %+v", st) + } + + resp = bearerDo(t, srv, http.MethodDelete, "/_plugin/runs/r_1/kill", "") + resp.Body.Close() + if resp.StatusCode != http.StatusNoContent { + t.Fatalf("unkill: want 204, got %d", resp.StatusCode) + } + if mr.Exists("bifrost:kill:r_1") { + t.Fatal("kill key not deleted") + } +} + +func TestRunState_SeededAccumulators(t *testing.T) { + srv, mr := newBudgetTestServer(t, nil) + mr.HSet("bifrost:cost:run:r_2", "total", "0.75") + mr.HSet("bifrost:steps:run:r_2", "total", "3") + mr.Lpush("bifrost:tools:run:r_2", "bash") + mr.SetTTL("bifrost:cost:run:r_2", 2*time.Hour) + + resp := bearerDo(t, srv, http.MethodGet, "/_plugin/runs/r_2/state", "") + if resp.StatusCode != http.StatusOK { + t.Fatalf("want 200, got %d", resp.StatusCode) + } + var st RunStateResponse + decodeBody(t, resp, &st) + if st.CostUSD != 0.75 || st.Steps != 3 || st.Killed { + t.Fatalf("state: %+v", st) + } + if len(st.Tools) != 1 || st.Tools[0] != "bash" { + t.Fatalf("tools: %v", st.Tools) + } + if st.TTLSeconds != 7200 { + t.Fatalf("ttl: %d", st.TTLSeconds) + } +} + +func TestRunState_EmptyRunIsZeroNot404(t *testing.T) { + srv, _ := newBudgetTestServer(t, nil) + resp := bearerDo(t, srv, http.MethodGet, "/_plugin/runs/r_never/state", "") + if resp.StatusCode != http.StatusOK { + t.Fatalf("want 200, got %d", resp.StatusCode) + } + var st RunStateResponse + decodeBody(t, resp, &st) + if st.Tools == nil { + t.Fatal("tools must serialize as [] not null") + } + if st.TTLSeconds != -2 { + t.Fatalf("ttl sentinel: %d", st.TTLSeconds) + } +} + +func TestAgentKill_RoundTrip(t *testing.T) { + srv, mr := newBudgetTestServer(t, map[string]auth.AgentBudget{ + "coder": {CapUSD: 5, Window: "1d"}, + }) + today := time.Now().UTC().Format("2006-01-02") + mr.HSet("bifrost:cost:agent:coder:"+today, "total", "1.5") + + resp := bearerDo(t, srv, http.MethodPost, "/_plugin/agents/coder/kill", "") + if resp.StatusCode != http.StatusOK { + t.Fatalf("kill: want 200, got %d", resp.StatusCode) + } + var kr KillAgentResponse + decodeBody(t, resp, &kr) + if kr.AgentName != "coder" { + t.Fatalf("kill response: %+v", kr) + } + if ttl := mr.TTL("bifrost:kill:agent:coder"); ttl != 24*time.Hour { + t.Fatalf("ttl: %v", ttl) + } + + // ?window=1h is ignored because the configured cap pins 1d. + resp = bearerDo(t, srv, http.MethodGet, "/_plugin/agents/coder/state?window=1h", "") + if resp.StatusCode != http.StatusOK { + t.Fatalf("state: want 200, got %d", resp.StatusCode) + } + var st AgentStateResponse + decodeBody(t, resp, &st) + if !st.Killed || st.Window != "1d" || st.BucketKey != today || st.CurrentSpendUSD != 1.5 { + t.Fatalf("state: %+v", st) + } + if st.ConfiguredCapUSD == nil || *st.ConfiguredCapUSD != 5 { + t.Fatalf("cap: %+v", st.ConfiguredCapUSD) + } + + resp = bearerDo(t, srv, http.MethodDelete, "/_plugin/agents/coder/kill", "") + resp.Body.Close() + if resp.StatusCode != http.StatusNoContent { + t.Fatalf("unkill: want 204, got %d", resp.StatusCode) + } + if mr.Exists("bifrost:kill:agent:coder") { + t.Fatal("kill key not deleted") + } +} + +func TestAgentState_NoCap_WindowParam(t *testing.T) { + srv, _ := newBudgetTestServer(t, nil) + resp := bearerDo(t, srv, http.MethodGet, "/_plugin/agents/free/state?window=1h", "") + if resp.StatusCode != http.StatusOK { + t.Fatalf("want 200, got %d", resp.StatusCode) + } + var st AgentStateResponse + decodeBody(t, resp, &st) + if st.Window != "1h" || st.ConfiguredCapUSD != nil || st.Killed { + t.Fatalf("state: %+v", st) + } + + resp = bearerDo(t, srv, http.MethodGet, "/_plugin/agents/free/state?window=bogus", "") + resp.Body.Close() + if resp.StatusCode != http.StatusBadRequest { + t.Fatalf("bogus window: want 400, got %d", resp.StatusCode) + } +} + +func TestHotState_MethodsAndValidation(t *testing.T) { + srv, _ := newBudgetTestServer(t, nil) + cases := []struct { + method, path string + want int + }{ + {http.MethodGet, "/_plugin/runs/r_1/kill", http.StatusMethodNotAllowed}, + {http.MethodPost, "/_plugin/runs/r_1/state", http.StatusMethodNotAllowed}, + {http.MethodGet, "/_plugin/agents/coder/kill", http.StatusMethodNotAllowed}, + {http.MethodPost, "/_plugin/agents/coder/state", http.StatusMethodNotAllowed}, + {http.MethodPost, "/_plugin/runs/has%20space/kill", http.StatusBadRequest}, + {http.MethodPost, "/_plugin/runs//kill", http.StatusNotFound}, + } + for _, c := range cases { + resp := bearerDo(t, srv, c.method, c.path, "") + resp.Body.Close() + if resp.StatusCode != c.want { + t.Errorf("%s %s: want %d, got %d", c.method, c.path, c.want, resp.StatusCode) + } + } +} + +func TestHotState_RedisUnconfigured503(t *testing.T) { + redisclient.SetClientForTest(nil) + mux := http.NewServeMux() + registerRoutes(mux, routeDeps{ + adminUser: "admin", adminPass: "hunter2", provisioningToken: testToken, + }) + srv := httptest.NewServer(mux) + t.Cleanup(srv.Close) + + for _, c := range []struct{ method, path string }{ + {http.MethodPost, "/_plugin/runs/r_1/kill"}, + {http.MethodGet, "/_plugin/runs/r_1/state"}, + {http.MethodPost, "/_plugin/agents/coder/kill"}, + {http.MethodGet, "/_plugin/agents/coder/state"}, + {http.MethodPost, "/_plugin/revoke/nonce/aaaa000000000000000000000000aaaa"}, + } { + resp := bearerDo(t, srv, c.method, c.path, "") + resp.Body.Close() + if resp.StatusCode != http.StatusServiceUnavailable { + t.Errorf("%s %s: want 503, got %d", c.method, c.path, resp.StatusCode) + } + } +} + +func TestHotState_RequiresAuth(t *testing.T) { + srv, _ := newBudgetTestServer(t, nil) + for _, c := range []struct{ method, path string }{ + {http.MethodPost, "/_plugin/runs/r_1/kill"}, + {http.MethodGet, "/_plugin/runs/r_1/state"}, + {http.MethodPost, "/_plugin/agents/coder/kill"}, + {http.MethodPost, "/_plugin/revoke/nonce/aaaa000000000000000000000000aaaa"}, + } { + req, _ := http.NewRequest(c.method, srv.URL+c.path, nil) + resp, err := srv.Client().Do(req) + if err != nil { + t.Fatal(err) + } + resp.Body.Close() + if resp.StatusCode != http.StatusUnauthorized { + t.Errorf("%s %s: want 401, got %d", c.method, c.path, resp.StatusCode) + } + } +} + +func TestRevokeNonce_RoundTrip(t *testing.T) { + srv, mr := newBudgetTestServer(t, nil) + const nonce = "bbbb000000000000000000000000bbbb" + exp := time.Now().UTC().Add(3 * time.Hour).Format(time.RFC3339) + + resp := bearerDo(t, srv, http.MethodPost, "/_plugin/revoke/nonce/"+nonce, `{"exp":"`+exp+`"}`) + if resp.StatusCode != http.StatusOK { + t.Fatalf("revoke: want 200, got %d", resp.StatusCode) + } + var rr RevokeNonceResponse + decodeBody(t, resp, &rr) + if rr.Nonce != nonce || rr.ExpiresAt == "" { + t.Fatalf("response: %+v", rr) + } + if !mr.Exists("bifrost:revoke:" + nonce) { + t.Fatal("tombstone not written") + } + // clamp(exp - now + 1h) ≈ 4h. + if ttl := mr.TTL("bifrost:revoke:" + nonce); ttl < 3*time.Hour+59*time.Minute || ttl > 4*time.Hour+time.Minute { + t.Fatalf("ttl: %v", ttl) + } + + resp = bearerDo(t, srv, http.MethodDelete, "/_plugin/revoke/nonce/"+nonce, "") + resp.Body.Close() + if resp.StatusCode != http.StatusNoContent { + t.Fatalf("unrevoke: want 204, got %d", resp.StatusCode) + } + if mr.Exists("bifrost:revoke:" + nonce) { + t.Fatal("tombstone not deleted") + } +} + +func TestRevokeNonce_DefaultsToCeiling_And_Validates(t *testing.T) { + srv, mr := newBudgetTestServer(t, nil) + const nonce = "cccc000000000000000000000000cccc" + + resp := bearerDo(t, srv, http.MethodPost, "/_plugin/revoke/nonce/"+nonce, "") + resp.Body.Close() + if resp.StatusCode != http.StatusOK { + t.Fatalf("no body: want 200, got %d", resp.StatusCode) + } + if ttl := mr.TTL("bifrost:revoke:" + nonce); ttl != 7*24*time.Hour { + t.Fatalf("default ttl should be the 7d ceiling, got %v", ttl) + } + + resp = bearerDo(t, srv, http.MethodPost, "/_plugin/revoke/nonce/not-hex", "") + resp.Body.Close() + if resp.StatusCode != http.StatusBadRequest { + t.Fatalf("bad nonce: want 400, got %d", resp.StatusCode) + } + resp = bearerDo(t, srv, http.MethodPost, "/_plugin/revoke/nonce/"+nonce, `{"exp":"yesterday"}`) + resp.Body.Close() + if resp.StatusCode != http.StatusBadRequest { + t.Fatalf("bad exp: want 400, got %d", resp.StatusCode) + } +} + +func TestRevokeUser_RoundTrip(t *testing.T) { + srv, mr := newBudgetTestServer(t, nil) + + resp := bearerDo(t, srv, http.MethodGet, "/_plugin/revoke/user/u_bob", "") + resp.Body.Close() + if resp.StatusCode != http.StatusNotFound { + t.Fatalf("unset cutoff: want 404, got %d", resp.StatusCode) + } + + resp = bearerDo(t, srv, http.MethodPut, "/_plugin/revoke/user/u_bob", `{"before":"2026-09-01T00:00:00Z"}`) + if resp.StatusCode != http.StatusOK { + t.Fatalf("set: want 200, got %d", resp.StatusCode) + } + var ur RevokeUserResponse + decodeBody(t, resp, &ur) + if ur.UserID != "u_bob" || ur.Before != "2026-09-01T00:00:00Z" { + t.Fatalf("response: %+v", ur) + } + if got, _ := mr.Get("bifrost:revoke_user_before:u_bob"); got != "2026-09-01T00:00:00Z" { + t.Fatalf("stored cutoff: %q", got) + } + + resp = bearerDo(t, srv, http.MethodGet, "/_plugin/revoke/user/u_bob", "") + if resp.StatusCode != http.StatusOK { + t.Fatalf("get: want 200, got %d", resp.StatusCode) + } + decodeBody(t, resp, &ur) + if ur.Before != "2026-09-01T00:00:00Z" { + t.Fatalf("get response: %+v", ur) + } + + resp = bearerDo(t, srv, http.MethodDelete, "/_plugin/revoke/user/u_bob", "") + resp.Body.Close() + if resp.StatusCode != http.StatusNoContent { + t.Fatalf("clear: want 204, got %d", resp.StatusCode) + } + if mr.Exists("bifrost:revoke_user_before:u_bob") { + t.Fatal("cutoff not cleared") + } +} + +func TestRevokeUser_DefaultsToNow(t *testing.T) { + srv, mr := newBudgetTestServer(t, nil) + before := time.Now().UTC() + resp := bearerDo(t, srv, http.MethodPut, "/_plugin/revoke/user/u_now", "") + resp.Body.Close() + if resp.StatusCode != http.StatusOK { + t.Fatalf("want 200, got %d", resp.StatusCode) + } + raw, _ := mr.Get("bifrost:revoke_user_before:u_now") + got, err := time.Parse(time.RFC3339, raw) + if err != nil { + t.Fatalf("stored cutoff not RFC3339: %q", raw) + } + if got.Before(before.Truncate(time.Second)) || got.After(time.Now().Add(time.Minute)) { + t.Fatalf("cutoff %v not ≈ now", got) + } +} diff --git a/gateway/internal/adminapi/observability.go b/gateway/internal/adminapi/observability.go index 8ab42ef41..6f8ebb894 100644 --- a/gateway/internal/adminapi/observability.go +++ b/gateway/internal/adminapi/observability.go @@ -628,8 +628,9 @@ func dimensionValue(l logstoreLog, dim string) string { // /_plugin/runs/{run_id} → runDetail (list) // /_plugin/runs/{run_id}/calls/{call_id} → runCallDetail (body) // -// Phase 6 will add /:id/state and /:id/kill under the same prefix; -// any other shape returns 404. We don't pull in a router library +// /:id/state and /:id/kill live under the same prefix but are +// routed to the phase-6 hot-state handlers before this one runs +// (see runsSubtree in server.go); any other shape returns 404. We don't pull in a router library // since the dispatch fits in a switch. func (h *observabilityHandlers) runDetail(w http.ResponseWriter, r *http.Request) { if r.Method != http.MethodGet { @@ -651,8 +652,9 @@ func (h *observabilityHandlers) runDetail(w http.ResponseWriter, r *http.Request case len(parts) == 3 && parts[1] == "calls" && parts[2] != "": h.runCallDetail(w, r, parts[0], parts[2]) default: - // Anything else (e.g. /:id/state, /:id/kill, trailing slash, - // 4-segment paths) is phase-6 territory or malformed; 404. + // /:id/state and /:id/kill are peeled off by the runs + // dispatcher in server.go before we get here; anything + // else (trailing slash, 4-segment paths) is malformed; 404. http.NotFound(w, r) } } diff --git a/gateway/internal/adminapi/observability_test.go b/gateway/internal/adminapi/observability_test.go index 19a6a8e81..155324d8a 100644 --- a/gateway/internal/adminapi/observability_test.go +++ b/gateway/internal/adminapi/observability_test.go @@ -472,10 +472,12 @@ func TestRunDetail_FiltersByRunID(t *testing.T) { func TestRunDetail_404OnSubpath(t *testing.T) { bf := newFakeBifrost(t, nil) srv := newObservabilityTestServer(t, bf) - resp := bearerGet(t, srv, "/_plugin/runs/r1/state") + // /:id/state and /:id/kill are phase-6 hot state (hotstate_test); + // any other subpath is malformed. + resp := bearerGet(t, srv, "/_plugin/runs/r1/other") defer resp.Body.Close() if resp.StatusCode != http.StatusNotFound { - t.Fatalf("want 404 on /:id/state (phase-6 territory), got %d", resp.StatusCode) + t.Fatalf("want 404 on unknown subpath, got %d", resp.StatusCode) } } diff --git a/gateway/internal/adminapi/revoke.go b/gateway/internal/adminapi/revoke.go new file mode 100644 index 000000000..000436216 --- /dev/null +++ b/gateway/internal/adminapi/revoke.go @@ -0,0 +1,165 @@ +package adminapi + +import ( + "net/http" + "strings" + "time" + + "github.com/stakwork/stakgraph/gateway/internal/auth" + "github.com/stakwork/stakgraph/gateway/internal/pluginlog" +) + +// Revocation routes — the HTTP face of auth/admin.go's helpers. +// Bearer-only: revocation is issuer territory (Hive's +// /macaroons/revoke fans out here), never a dashboard click. +// +// POST /_plugin/revoke/nonce/:nonce body {exp?} → 200 RevokeNonceResponse +// DELETE /_plugin/revoke/nonce/:nonce → 204 +// PUT /_plugin/revoke/user/:user_id body {before?} → 200 RevokeUserResponse +// GET /_plugin/revoke/user/:user_id → 200 RevokeUserResponse | 404 +// DELETE /_plugin/revoke/user/:user_id → 204 +// +// `exp` is the RFC3339 expiry of the layer whose nonce is being +// revoked; the tombstone's TTL is derived from it (phase-6: "TTL = +// layer.exp"). Omitted ⇒ the 7d ceiling, which is always safe +// (macaroon layers never outlive it). `before` defaults to now. + +const revokePrefixPath = "/_plugin/revoke/" + +// RevokeNonceRequest is the body for POST /_plugin/revoke/nonce/:nonce. +type RevokeNonceRequest struct { + Exp string `json:"exp,omitempty"` // RFC3339; layer expiry +} + +// RevokeNonceResponse is the wire shape for POST /_plugin/revoke/nonce/:nonce. +type RevokeNonceResponse struct { + Nonce string `json:"nonce"` + ExpiresAt string `json:"expires_at"` // RFC3339 UTC; when the tombstone lapses +} + +// RevokeUserRequest is the body for PUT /_plugin/revoke/user/:user_id. +type RevokeUserRequest struct { + Before string `json:"before,omitempty"` // RFC3339; defaults to now +} + +// RevokeUserResponse is the wire shape for PUT/GET /_plugin/revoke/user/:user_id. +type RevokeUserResponse struct { + UserID string `json:"user_id"` + Before string `json:"before"` // RFC3339 UTC +} + +type revokeHandlers struct{} + +func newRevokeHandlers() *revokeHandlers { return &revokeHandlers{} } + +// dispatch routes /_plugin/revoke/{nonce,user}/. +func (h *revokeHandlers) dispatch(w http.ResponseWriter, r *http.Request) { + parts := strings.Split(strings.TrimPrefix(r.URL.Path, revokePrefixPath), "/") + if len(parts) != 2 || parts[1] == "" { + http.NotFound(w, r) + return + } + switch parts[0] { + case "nonce": + h.nonce(w, r, parts[1]) + case "user": + h.user(w, r, parts[1]) + default: + http.NotFound(w, r) + } +} + +func (h *revokeHandlers) nonce(w http.ResponseWriter, r *http.Request, nonce string) { + switch r.Method { + case http.MethodPost: + var body RevokeNonceRequest + if r.ContentLength != 0 { + if err := decodeJSON(r, &body); err != nil { + http.Error(w, "bad json: "+err.Error(), http.StatusBadRequest) + return + } + } + now := time.Now().UTC() + // Default to the ceiling: runKeyTTL clamps anything past 7d. + exp := now.Add(365 * 24 * time.Hour) + if body.Exp != "" { + t, err := time.Parse(time.RFC3339, body.Exp) + if err != nil { + http.Error(w, "exp must be RFC3339", http.StatusBadRequest) + return + } + exp = t + } + ttl := auth.RevocationTTL(exp, now) + if err := auth.RevokeNonce(r.Context(), nonce, ttl); err != nil { + writeHotStateErr(w, err, "revoke.nonce") + return + } + pluginlog.Logf("adminapi: nonce revoked ttl=%s", ttl) + writeJSON(w, http.StatusOK, RevokeNonceResponse{ + Nonce: nonce, + ExpiresAt: now.Add(ttl).Format(time.RFC3339), + }) + case http.MethodDelete: + if err := auth.UnrevokeNonce(r.Context(), nonce); err != nil { + writeHotStateErr(w, err, "revoke.unrevoke_nonce") + return + } + w.WriteHeader(http.StatusNoContent) + default: + methodNotAllowed(w, http.MethodPost, http.MethodDelete) + } +} + +func (h *revokeHandlers) user(w http.ResponseWriter, r *http.Request, userID string) { + switch r.Method { + case http.MethodPut: + var body RevokeUserRequest + if r.ContentLength != 0 { + if err := decodeJSON(r, &body); err != nil { + http.Error(w, "bad json: "+err.Error(), http.StatusBadRequest) + return + } + } + before := time.Now().UTC() + if body.Before != "" { + t, err := time.Parse(time.RFC3339, body.Before) + if err != nil { + http.Error(w, "before must be RFC3339", http.StatusBadRequest) + return + } + before = t.UTC() + } + if err := auth.SetUserRevokeCutoff(r.Context(), userID, before); err != nil { + writeHotStateErr(w, err, "revoke.user") + return + } + pluginlog.Logf("adminapi: user revoke cutoff set user=%s before=%s", userID, before.Format(time.RFC3339)) + writeJSON(w, http.StatusOK, RevokeUserResponse{ + UserID: userID, + Before: before.Format(time.RFC3339), + }) + case http.MethodGet: + cutoff, ok, err := auth.GetUserRevokeCutoff(r.Context(), userID) + if err != nil { + writeHotStateErr(w, err, "revoke.get_user") + return + } + if !ok { + http.NotFound(w, r) + return + } + writeJSON(w, http.StatusOK, RevokeUserResponse{ + UserID: userID, + Before: cutoff.UTC().Format(time.RFC3339), + }) + case http.MethodDelete: + if err := auth.ClearUserRevokeCutoff(r.Context(), userID); err != nil { + writeHotStateErr(w, err, "revoke.clear_user") + return + } + w.WriteHeader(http.StatusNoContent) + default: + methodNotAllowed(w, http.MethodPut, http.MethodGet, http.MethodDelete) + } +} diff --git a/gateway/internal/adminapi/server.go b/gateway/internal/adminapi/server.go index 95551d54c..89d022399 100644 --- a/gateway/internal/adminapi/server.go +++ b/gateway/internal/adminapi/server.go @@ -189,10 +189,12 @@ func methodMuxedAuth( // Auth model // ---------- // - `/_plugin/health`, `/_plugin/login`: anonymous. -// - `/_plugin/admin-credentials`, `/_plugin/trust/*`: bearer only -// (Hive's machine-to-plugin path; cookies are not honoured). -// - Everything else (observability, /me, /logout): cookie OR -// bearer, with cookie tried first. +// - `/_plugin/admin-credentials`, `/_plugin/trust/*`, +// `/_plugin/revoke/*`: bearer only (Hive's machine-to-plugin +// path; cookies are not honoured). +// - Everything else (observability, kill/state, /me, /logout): +// cookie OR bearer, with cookie tried first. Cookie-authed +// mutations (kill, unkill, toggles) need the CSRF header. // // The /_plugin/ui/* SPA is also cookie-or-bearer so curl with a // bearer can pull it for diagnostics, but browsers always reach it @@ -248,29 +250,61 @@ func registerRoutes(mux *http.ServeMux, deps routeDeps) { mux.HandleFunc("/_plugin/auth/ticket", bearer(ticketH.mint)) mux.HandleFunc("/_plugin/auth/redeem", ticketH.redeem) // anon; ticket IS the proof + // Phase-6 hot state: kill switches + Redis snapshots. Registered + // unconditionally (they need Redis, not the logstore) and + // dispatched from the shared /runs/ and /agents/ subtrees below. + hot := newHotStateHandlers() + // Cookie-or-bearer routes: phase-7 observability subset. + var obs *observabilityHandlers if deps.logstore != nil { - obs := newObservabilityHandlers(deps.logstore) + obs = newObservabilityHandlers(deps.logstore) mux.HandleFunc("/_plugin/spend/by-agent", cookieOrBearer(obs.spendByAgent)) mux.HandleFunc("/_plugin/spend/by-user", cookieOrBearer(obs.spendByUser)) mux.HandleFunc("/_plugin/spend/by-agent-user", cookieOrBearer(obs.spendByAgentUser)) mux.HandleFunc("/_plugin/histogram/cost", cookieOrBearer(obs.histogramCost)) - // /_plugin/runs/ takes a trailing path segment as run-id - mux.HandleFunc("/_plugin/runs/", cookieOrBearer(obs.runDetail)) // /_plugin/users/ takes a trailing path segment as user-id. // Phase-8 only exposes the rollup (KPIs + agents-used + // runs); phase-9 adds /:id/quota for spend-vs-cap. mux.HandleFunc("/_plugin/users/", cookieOrBearer(obs.userDetail)) } + // The `/_plugin/runs/` subtree: `/kill` and `/state` are + // phase-6 hot state; everything else (``, `/calls/`) + // is the phase-7 logs.db drill-down, which 404s when the logstore + // isn't configured. + runsSubtree := func(w http.ResponseWriter, r *http.Request) { + rest := strings.TrimPrefix(r.URL.Path, "/_plugin/runs/") + parts := strings.Split(rest, "/") + if len(parts) == 2 && parts[0] != "" { + switch parts[1] { + case "kill": + hot.runKill(w, r, parts[0]) + return + case "state": + hot.runState(w, r, parts[0]) + return + } + } + if obs == nil { + http.NotFound(w, r) + return + } + obs.runDetail(w, r) + } + mux.HandleFunc("/_plugin/runs/", cookieOrBearer(runsSubtree)) + + // Revocation admin (bearer-only; issuer territory). + rv := newRevokeHandlers() + mux.HandleFunc(revokePrefixPath, bearer(rv.dispatch)) + // Phase-8.5 per-agent budget view. Reads cap from plugin config // (auth.GetConfig().AgentBudgets) and current-bucket spend from // Redis (`bifrost:cost:agent::`) with a // fallback to summing `logs.db` rows when phase 6's PostHook // hasn't filled the Redis hash yet. Subtree routing on // `/_plugin/agents/`; the handler enforces the `/budget` - // shape and 404s on anything else (phase-9 `:name/state` and - // `:name/kill` live under the same prefix later). + // shape and 404s on anything else. bgt := newBudgetHandlers(deps.logstore) cat := newCatalogHandlers(deps.graph) hiveCb := newHiveCallbackHandlers() @@ -278,13 +312,22 @@ func registerRoutes(mux *http.ServeMux, deps routeDeps) { // The `/_plugin/agents/` subtree is shared: ServeMux allows only // one handler per pattern, so a small dispatcher fans the trailing - // `/` segment out to the budget (phase-8.5) or catalog - // (agent-catalog) read. Both are cookie-or-bearer reads. Anything - // else under the prefix 404s, preserving the budget handler's - // existing contract. + // `/` segment out to the budget (phase-8.5), catalog + // (agent-catalog), or hot-state (phase-6 kill/state) handler. All + // cookie-or-bearer. Anything else under the prefix 404s, + // preserving the budget handler's existing contract. agentsSubtree := func(w http.ResponseWriter, r *http.Request) { rest := strings.TrimPrefix(r.URL.Path, "/_plugin/agents/") parts := strings.Split(rest, "/") + // `/kill` (POST/DELETE) and `/state` (GET). + if len(parts) == 2 && parts[0] != "" && parts[1] == "kill" { + hot.agentKill(w, r, parts[0]) + return + } + if len(parts) == 2 && parts[0] != "" && parts[1] == "state" { + hot.agentState(w, r, parts[0]) + return + } // `/_plugin/agents/catalog` (single segment) is the catalog // list — every registry agent, traffic or not. Distinct from // `/catalog` (two segments) which is one agent's detail. diff --git a/gateway/internal/auth/admin.go b/gateway/internal/auth/admin.go index 5382fc752..e35e5be53 100644 --- a/gateway/internal/auth/admin.go +++ b/gateway/internal/auth/admin.go @@ -13,9 +13,10 @@ import ( ) // Admin-side helpers for managing revocation state. The HTTP routes -// that expose these (POST /_plugin/admin/revoke, etc.) live in -// gateway/internal/adminapi and call into this file. Phase-6 will -// grow the kill-switch and per-run/state endpoints alongside these. +// that expose these (`/_plugin/revoke/nonce/:nonce`, +// `/_plugin/revoke/user/:user_id`) live in gateway/internal/adminapi +// (revoke.go) and call into this file. The kill-switch and per-run / +// per-agent state primitives live next door in kill.go. // // These helpers exist now so a swarm operator can: // @@ -59,6 +60,15 @@ func RevokeNonce(ctx context.Context, nonce string, ttl time.Duration) error { return rdb.Set(octx, redisclient.Key(revokePrefix+nonce), "1", ttl).Err() } +// RevocationTTL computes the tombstone TTL for a nonce whose layer +// expires at `layerExp` — the same clamp(exp-now+1h, 1h, 7d) formula +// the accumulators use, so a revoke outlives the macaroon by the +// grace hour and never lingers past the 7d ceiling. Exported for the +// adminapi revoke handler. +func RevocationTTL(layerExp, now time.Time) time.Duration { + return runKeyTTL(layerExp, now) +} + // UnrevokeNonce removes a revocation tombstone. Mostly for operator // recovery from a mistakenly-pressed revoke button; revocations are // supposed to be permanent within the macaroon's lifetime. diff --git a/gateway/internal/auth/doc.go b/gateway/internal/auth/doc.go index 0eb305337..ebb84da6f 100644 --- a/gateway/internal/auth/doc.go +++ b/gateway/internal/auth/doc.go @@ -11,28 +11,33 @@ // - config.go enforce_macaroons flag (shadow → enforce rollout), // agent_budgets, model_pricing // - verifier.go Verify() — header extraction + trust lookup + pure verify -// - revocation.go CheckRevocations() — bifrost:revoke:* / revoke_user_before:* +// - revocation.go CheckRevocations() — phase-6 PIPELINE 1: +// bifrost:revoke:* / revoke_user_before:* (401) and the +// kill: (every chain layer) / kill:agent: switches (402) +// - kill.go KillRun/KillAgent + Unkill*, GetRunState/GetAgentState — +// the admin primitives behind /_plugin/{runs,agents}/:id/{kill,state} // - ttl.go clamp(exp-now+1h, 1h, 7d) shared by revocation + accumulators // - enforcement.go Evaluate() + ApplyToLLMPre() — hook glue // - accumulator.go ApplyToLLMPost() — phase-6 PostLLMHook pipeline: // cost:run / steps:run per chain layer, cost:ua envelope, // cost:agent windowed buckets, tools:run history // - pricing.go PriceCall() — model_pricing table → dollars -// - admin.go admin endpoints (revoke management — minimal scope) +// - admin.go revoke primitives behind /_plugin/revoke/* // // What's still out of scope (phase 6 read side) // --------------------------------------------- -// - Per-run cost/step cap walk in PreLLMHook (reads the -// accumulators this package now writes) +// - Per-run cost/step cap walk in PreLLMHook (PIPELINE 2 — reads +// the accumulators this package now writes) // - ua_budget / realm_budget / agent budget 402 rejections // - Tool-loop detection (reads tools:run) -// - Kill switches (kill:, kill:agent:) + admin routes +// - hard_ceiling / tool_loop config + the /_plugin/config/* overrides // // The write side landing first is deliberate: accumulators are // shadow-safe (they reject nothing), they light up the phase-8 // budget endpoint that currently falls back to logs.db, and the cap // walk needs weeks of real accumulated state to validate against -// before it starts rejecting. +// before it starts rejecting. Kill switches don't depend on that +// state, so they shipped ahead of the cap walk. // // Operational posture // ------------------- @@ -44,8 +49,10 @@ // - LOGS LOUDLY when a macaroon would have been rejected. // - Does NOT reject — the request continues to the provider. // -// With enforce_macaroons=true the failure path becomes 401/402 with -// a stable AdapterError.Code. Operators flip the flag per-swarm once +// With enforce_macaroons=true the failure path becomes 401 (bad, +// missing or revoked macaroon) or 402 (valid macaroon, but the run +// or agent was killed) with a stable AdapterError.Code. Operators +// flip the flag per-swarm once // the shadow-mode logs show no false positives — either in the // config.json plugin block or, without rebuilding the image, via the // BIFROST_PLUGIN_ENFORCE_MACAROONS env var (which wins when set; an diff --git a/gateway/internal/auth/enforcement.go b/gateway/internal/auth/enforcement.go index ebc3adb2b..6e6f379ab 100644 --- a/gateway/internal/auth/enforcement.go +++ b/gateway/internal/auth/enforcement.go @@ -82,9 +82,9 @@ func Evaluate(ctx context.Context, rawMacaroon string) Decision { // still tentatively-rejected until all post-verify checks pass. d.PostVerifyClaims = claims - // Note this is the only Redis touch in the phase-4 adapter; - // phase 6 adds cost/steps cap-walks and per-agent budget reads - // here. + // Phase-6 PIPELINE 1: revocations + kill switches in one Redis + // round-trip. The cost/steps cap walk (PIPELINE 2) is still to + // come and slots in after this. if revErr := CheckRevocations(ctx, claims); revErr != nil { d.Err = revErr return d @@ -235,7 +235,15 @@ func shortCircuitFromError(e *AdapterError) *schemas.LLMPluginShortCircuit { if status == 0 { status = 401 } + // 401s are verification failures (bad/missing/revoked + // macaroon); 402s are enforcement decisions against a valid + // macaroon (kill switches today, budget caps once the cap walk + // lands). Distinct Type so clients can tell "re-issue the + // macaroon" apart from "an operator or a cap stopped you". errType := "macaroon_verification_failed" + if status == 402 { + errType = "enforcement_rejected" + } return &schemas.LLMPluginShortCircuit{ Error: &schemas.BifrostError{ IsBifrostError: false, // it's an auth error, not a transport error diff --git a/gateway/internal/auth/kill.go b/gateway/internal/auth/kill.go new file mode 100644 index 000000000..7f94e0d91 --- /dev/null +++ b/gateway/internal/auth/kill.go @@ -0,0 +1,235 @@ +package auth + +import ( + "context" + "errors" + "fmt" + "strings" + "time" + + "github.com/redis/go-redis/v9" + + "github.com/stakwork/stakgraph/gateway/internal/duration" + "github.com/stakwork/stakgraph/gateway/internal/redisclient" +) + +// Kill-switch and hot-state helpers — the phase-6 admin primitives +// (`KillRun`, `UnkillRun`, `KillAgent`, `UnkillAgent`, `RunState`, +// `AgentState`). Pure Redis logic, no HTTP; the routes that expose +// them live in gateway/internal/adminapi (runs.go / agents.go). +// +// A kill does not invalidate the macaroon — it's an operator saying +// "stop this now". The hot path (CheckRevocations) turns a set kill +// key into a 402 `run_killed` / `agent_killed` on the next LLM call +// from the run (or any descendant) / from the agent. Both keys carry +// fixed TTLs per phase-6 "TTL policy": a kill is hot operational +// state, not configuration. To stop an agent permanently, set its +// agent_budget cap instead. + +const ( + // killRunTTL: if the run isn't dead within the hour, something + // else has gone wrong. + killRunTTL = 1 * time.Hour + // killAgentTTL: re-set if the kill needs to persist longer. + killAgentTTL = 24 * time.Hour +) + +// RunState is a snapshot of one run's phase-6 Redis accumulators. +type RunState struct { + RunID string + CostUSD float64 // HGET cost:run: total (0 when absent) + Steps int64 // HGET steps:run: total (0 when absent) + Tools []string // LRANGE tools:run: 0 9, most recent first + Killed bool // EXISTS kill: + TTLSeconds int64 // TTL cost:run:; -2 when the key is absent, -1 when no expiry +} + +// AgentState is a snapshot of one agent's current-bucket spend and +// kill state. ConfiguredCapUSD is nil when the agent has no entry in +// agent_budgets. +type AgentState struct { + AgentName string + Window string + BucketKey string + CurrentSpendUSD float64 + ConfiguredCapUSD *float64 + Killed bool +} + +// KillRun sets bifrost:kill: = "1" EX 1h. Idempotent; a +// re-kill refreshes the TTL. +func KillRun(ctx context.Context, runID string) error { + if err := validateKillID("run_id", runID); err != nil { + return err + } + rdb := redisclient.Client() + if rdb == nil { + return ErrRedisUnavailable + } + octx, cancel := context.WithTimeout(ctx, adminTimeout) + defer cancel() + return rdb.Set(octx, redisclient.Key(killRunPrefix+runID), "1", killRunTTL).Err() +} + +// UnkillRun deletes the per-run kill key. No-op if absent. +func UnkillRun(ctx context.Context, runID string) error { + if err := validateKillID("run_id", runID); err != nil { + return err + } + rdb := redisclient.Client() + if rdb == nil { + return ErrRedisUnavailable + } + octx, cancel := context.WithTimeout(ctx, adminTimeout) + defer cancel() + return rdb.Del(octx, redisclient.Key(killRunPrefix+runID)).Err() +} + +// KillAgent sets bifrost:kill:agent: = "1" EX 24h. +func KillAgent(ctx context.Context, name string) error { + if err := validateKillID("agent_name", name); err != nil { + return err + } + rdb := redisclient.Client() + if rdb == nil { + return ErrRedisUnavailable + } + octx, cancel := context.WithTimeout(ctx, adminTimeout) + defer cancel() + return rdb.Set(octx, redisclient.Key(killAgentPrefix+name), "1", killAgentTTL).Err() +} + +// UnkillAgent deletes the per-agent kill key. No-op if absent. +func UnkillAgent(ctx context.Context, name string) error { + if err := validateKillID("agent_name", name); err != nil { + return err + } + rdb := redisclient.Client() + if rdb == nil { + return ErrRedisUnavailable + } + octx, cancel := context.WithTimeout(ctx, adminTimeout) + defer cancel() + return rdb.Del(octx, redisclient.Key(killAgentPrefix+name)).Err() +} + +// GetRunState reads the run's accumulators + kill flag in one +// pipelined round-trip. Absent keys read as zero / empty / not +// killed — a run that has never made a call is a valid, empty state, +// not an error. +func GetRunState(ctx context.Context, runID string) (RunState, error) { + st := RunState{RunID: runID, Tools: []string{}} + if err := validateKillID("run_id", runID); err != nil { + return st, err + } + rdb := redisclient.Client() + if rdb == nil { + return st, ErrRedisUnavailable + } + octx, cancel := context.WithTimeout(ctx, adminTimeout) + defer cancel() + + costKey := redisclient.Key(costRunPrefix + runID) + pipe := rdb.Pipeline() + costCmd := pipe.HGet(octx, costKey, "total") + stepsCmd := pipe.HGet(octx, redisclient.Key(stepsRunPrefix+runID), "total") + toolsCmd := pipe.LRange(octx, redisclient.Key(toolsRunPrefix+runID), 0, toolHistoryLen-1) + killCmd := pipe.Exists(octx, redisclient.Key(killRunPrefix+runID)) + ttlCmd := pipe.TTL(octx, costKey) + if _, err := pipe.Exec(octx); err != nil && !errors.Is(err, redis.Nil) { + return st, fmt.Errorf("redis pipeline: %w", err) + } + + if v, err := costCmd.Float64(); err == nil { + st.CostUSD = v + } else if !errors.Is(err, redis.Nil) { + return st, fmt.Errorf("cost:run: %w", err) + } + if v, err := stepsCmd.Int64(); err == nil { + st.Steps = v + } else if !errors.Is(err, redis.Nil) { + return st, fmt.Errorf("steps:run: %w", err) + } + if v, err := toolsCmd.Result(); err == nil && v != nil { + st.Tools = v + } + if v, err := killCmd.Result(); err == nil { + st.Killed = v == 1 + } + // go-redis reports "no key" as -2ns and "no expiry" as -1ns; + // pass those sentinels through unchanged in seconds. + switch d, _ := ttlCmd.Result(); { + case d < 0: + st.TTLSeconds = int64(d) + default: + st.TTLSeconds = int64(d / time.Second) + } + return st, nil +} + +// GetAgentState reads the agent's current-bucket spend + kill flag. +// The window comes from agent_budgets when the agent has a +// configured cap (so the reported bucket is the one enforcement +// reads); otherwise `fallbackWindow` (the operator's ?window=, or +// "1d") picks which bucket to report — informational only, since no +// cap applies to it. +func GetAgentState(ctx context.Context, name, fallbackWindow string, now time.Time) (AgentState, error) { + st := AgentState{AgentName: name} + if err := validateKillID("agent_name", name); err != nil { + return st, err + } + rdb := redisclient.Client() + if rdb == nil { + return st, ErrRedisUnavailable + } + + window := fallbackWindow + if b, ok := GetConfig().AgentBudgets[name]; ok && b.CapUSD > 0 && b.Window != "" { + cap := b.CapUSD + st.ConfiguredCapUSD = &cap + window = b.Window + } + if window == "" { + window = "1d" + } + w, err := duration.Parse(window) + if err != nil { + return st, fmt.Errorf("window %q: %w", window, err) + } + st.Window = window + st.BucketKey = w.BucketKey(now) + + octx, cancel := context.WithTimeout(ctx, adminTimeout) + defer cancel() + pipe := rdb.Pipeline() + spendCmd := pipe.HGet(octx, redisclient.Key(costAgentPrefix+name+":"+st.BucketKey), "total") + killCmd := pipe.Exists(octx, redisclient.Key(killAgentPrefix+name)) + if _, err := pipe.Exec(octx); err != nil && !errors.Is(err, redis.Nil) { + return st, fmt.Errorf("redis pipeline: %w", err) + } + if v, err := spendCmd.Float64(); err == nil { + st.CurrentSpendUSD = v + } else if !errors.Is(err, redis.Nil) { + return st, fmt.Errorf("cost:agent: %w", err) + } + if v, err := killCmd.Result(); err == nil { + st.Killed = v == 1 + } + return st, nil +} + +// validateKillID rejects ids that would produce a malformed or +// surprising Redis key: empty, absurdly long, or containing +// whitespace / path separators (which would also have been mangled +// by the URL router that delivered them). +func validateKillID(field, v string) error { + switch { + case v == "": + return fmt.Errorf("%s is required", field) + case len(v) > 256: + return fmt.Errorf("%s too long (%d > 256)", field, len(v)) + case strings.ContainsAny(v, " \t\r\n/"): + return fmt.Errorf("%s contains whitespace or '/'", field) + } + return nil +} diff --git a/gateway/internal/auth/kill_test.go b/gateway/internal/auth/kill_test.go new file mode 100644 index 000000000..62617cc13 --- /dev/null +++ b/gateway/internal/auth/kill_test.go @@ -0,0 +1,291 @@ +package auth + +import ( + "context" + "testing" + "time" + + macaroon "github.com/stakwork/stakgraph/gateway/auth/go" + "github.com/stakwork/stakgraph/gateway/internal/pluginctx" + "github.com/stakwork/stakgraph/gateway/internal/redisclient" +) + +func killTestClaims() *macaroon.Claims { + return &macaroon.Claims{ + UserID: testUserID, + AgentName: "coder", + RunID: "r_child", + Nonces: []string{"aaaa000000000000000000000000aaaa", "bbbb000000000000000000000000bbbb"}, + IAT: time.Now().UTC().Format(time.RFC3339), + Chain: []macaroon.ChainLayer{ + {RunID: "r_parent", MaxCostUSD: 5}, + {RunID: "r_child", MaxCostUSD: 2}, + }, + } +} + +func TestCheckRevocations_KillLeafRun(t *testing.T) { + mr := newMiniRedis(t) + mr.Set("bifrost:kill:r_child", "1") + + err := CheckRevocations(context.Background(), killTestClaims()) + if err == nil || err.Code != "run_killed" { + t.Fatalf("want run_killed, got %+v", err) + } + if err.HTTPStatus != 402 { + t.Fatalf("want 402, got %d", err.HTTPStatus) + } +} + +func TestCheckRevocations_KillAncestorRun_KillsDescendant(t *testing.T) { + mr := newMiniRedis(t) + // Only the parent is killed; the leaf's own key is absent. + mr.Set("bifrost:kill:r_parent", "1") + + err := CheckRevocations(context.Background(), killTestClaims()) + if err == nil || err.Code != "run_killed" { + t.Fatalf("want run_killed via ancestor, got %+v", err) + } +} + +func TestCheckRevocations_KillAgent(t *testing.T) { + mr := newMiniRedis(t) + mr.Set("bifrost:kill:agent:coder", "1") + + err := CheckRevocations(context.Background(), killTestClaims()) + if err == nil || err.Code != "agent_killed" { + t.Fatalf("want agent_killed, got %+v", err) + } + if err.HTTPStatus != 402 { + t.Fatalf("want 402, got %d", err.HTTPStatus) + } +} + +func TestCheckRevocations_KillAgent_LeafOnly(t *testing.T) { + mr := newMiniRedis(t) + // A kill on a different agent name must not match. + mr.Set("bifrost:kill:agent:web-search", "1") + + if err := CheckRevocations(context.Background(), killTestClaims()); err != nil { + t.Fatalf("kill on another agent must not match, got %+v", err) + } +} + +func TestCheckRevocations_RevokedBeatsKilled(t *testing.T) { + mr := newMiniRedis(t) + mr.Set("bifrost:kill:r_child", "1") + mr.Set("bifrost:revoke:bbbb000000000000000000000000bbbb", "1") + + err := CheckRevocations(context.Background(), killTestClaims()) + if err == nil || err.Code != "macaroon_revoked" { + t.Fatalf("revocation should win over kill, got %+v", err) + } +} + +func TestCheckRevocations_KillFallsBackToLeafRunID(t *testing.T) { + mr := newMiniRedis(t) + mr.Set("bifrost:kill:r_solo", "1") + + claims := &macaroon.Claims{ + UserID: testUserID, + RunID: "r_solo", // no Chain populated + Nonces: []string{"aaaa000000000000000000000000aaaa"}, + IAT: time.Now().UTC().Format(time.RFC3339), + } + err := CheckRevocations(context.Background(), claims) + if err == nil || err.Code != "run_killed" { + t.Fatalf("want run_killed from Claims.RunID fallback, got %+v", err) + } +} + +func TestKillRun_RoundTrip(t *testing.T) { + mr := newMiniRedis(t) + ctx := context.Background() + + if err := KillRun(ctx, "r_1"); err != nil { + t.Fatal(err) + } + if !mr.Exists("bifrost:kill:r_1") { + t.Fatal("kill key not written") + } + if ttl := mr.TTL("bifrost:kill:r_1"); ttl != killRunTTL { + t.Fatalf("ttl: want %v, got %v", killRunTTL, ttl) + } + if err := UnkillRun(ctx, "r_1"); err != nil { + t.Fatal(err) + } + if mr.Exists("bifrost:kill:r_1") { + t.Fatal("kill key not deleted") + } +} + +func TestKillAgent_RoundTrip(t *testing.T) { + mr := newMiniRedis(t) + ctx := context.Background() + + if err := KillAgent(ctx, "coder"); err != nil { + t.Fatal(err) + } + if ttl := mr.TTL("bifrost:kill:agent:coder"); ttl != killAgentTTL { + t.Fatalf("ttl: want %v, got %v", killAgentTTL, ttl) + } + if err := UnkillAgent(ctx, "coder"); err != nil { + t.Fatal(err) + } + if mr.Exists("bifrost:kill:agent:coder") { + t.Fatal("kill key not deleted") + } +} + +func TestKill_Validation(t *testing.T) { + _ = newMiniRedis(t) + ctx := context.Background() + for _, bad := range []string{"", "has space", "a/b", string(make([]byte, 300))} { + if err := KillRun(ctx, bad); err == nil { + t.Errorf("KillRun(%q): want validation error", bad) + } + if err := KillAgent(ctx, bad); err == nil { + t.Errorf("KillAgent(%q): want validation error", bad) + } + } +} + +func TestKill_RedisUnavailable(t *testing.T) { + redisclient.SetClientForTest(nil) + if err := KillRun(context.Background(), "r_1"); err != ErrRedisUnavailable { + t.Fatalf("want ErrRedisUnavailable, got %v", err) + } + if _, err := GetRunState(context.Background(), "r_1"); err != ErrRedisUnavailable { + t.Fatalf("want ErrRedisUnavailable, got %v", err) + } +} + +func TestGetRunState_Populated(t *testing.T) { + mr := newMiniRedis(t) + mr.HSet("bifrost:cost:run:r_1", "total", "1.25") + mr.HSet("bifrost:steps:run:r_1", "total", "7") + mr.Lpush("bifrost:tools:run:r_1", "read_file") + mr.Lpush("bifrost:tools:run:r_1", "bash") + mr.Set("bifrost:kill:r_1", "1") + mr.SetTTL("bifrost:cost:run:r_1", 90*time.Minute) + + st, err := GetRunState(context.Background(), "r_1") + if err != nil { + t.Fatal(err) + } + if st.CostUSD != 1.25 || st.Steps != 7 || !st.Killed { + t.Fatalf("state: %+v", st) + } + if len(st.Tools) != 2 || st.Tools[0] != "bash" || st.Tools[1] != "read_file" { + t.Fatalf("tools (most recent first): %v", st.Tools) + } + if st.TTLSeconds != int64((90 * time.Minute).Seconds()) { + t.Fatalf("ttl: %d", st.TTLSeconds) + } +} + +func TestGetRunState_Empty(t *testing.T) { + _ = newMiniRedis(t) + st, err := GetRunState(context.Background(), "r_never_called") + if err != nil { + t.Fatal(err) + } + if st.CostUSD != 0 || st.Steps != 0 || st.Killed || len(st.Tools) != 0 { + t.Fatalf("want empty state, got %+v", st) + } + if st.TTLSeconds != -2 { + t.Fatalf("want -2 (no key) ttl sentinel, got %d", st.TTLSeconds) + } +} + +func TestGetAgentState_ConfiguredCapWinsWindow(t *testing.T) { + mr := newMiniRedis(t) + SetConfigForTest(Config{AgentBudgets: map[string]AgentBudget{ + "coder": {CapUSD: 5, Window: "1d"}, + }}) + t.Cleanup(func() { SetConfigForTest(Config{}) }) + + now := time.Now().UTC() + mr.HSet("bifrost:cost:agent:coder:"+now.Format("2006-01-02"), "total", "2.5") + mr.Set("bifrost:kill:agent:coder", "1") + + // ?window=1h must be ignored: the configured 1d bucket is the one + // enforcement reads. + st, err := GetAgentState(context.Background(), "coder", "1h", now) + if err != nil { + t.Fatal(err) + } + if st.Window != "1d" || st.CurrentSpendUSD != 2.5 || !st.Killed { + t.Fatalf("state: %+v", st) + } + if st.ConfiguredCapUSD == nil || *st.ConfiguredCapUSD != 5 { + t.Fatalf("cap: %+v", st.ConfiguredCapUSD) + } +} + +func TestGetAgentState_NoCap_FallbackWindow(t *testing.T) { + _ = newMiniRedis(t) + st, err := GetAgentState(context.Background(), "nobudget", "", time.Now().UTC()) + if err != nil { + t.Fatal(err) + } + if st.Window != "1d" || st.ConfiguredCapUSD != nil || st.Killed || st.CurrentSpendUSD != 0 { + t.Fatalf("state: %+v", st) + } + if _, err := GetAgentState(context.Background(), "nobudget", "1x", time.Now()); err == nil { + t.Fatal("want error for unparseable window") + } +} + +// End-to-end through the hook entry point: a real signed macaroon, +// enforce mode, run killed → 402 short-circuit with the stable code. +func TestApplyToLLMPre_EnforceMode_RunKilled402(t *testing.T) { + reg := newTestRegistry(t) + SetTrustRegistry(reg) + t.Cleanup(func() { SetTrustRegistry(nil) }) + SetConfigForTest(Config{EnforceMacaroons: true}) + t.Cleanup(func() { SetConfigForTest(Config{}) }) + mr := newMiniRedis(t) + + opts := defaultMacaroonOptions(time.Now()) + encoded := buildMacaroon(t, opts) + mr.Set("bifrost:kill:"+opts.runID, "1") + + bctx := newBifrostCtx() + pluginctx.SetRawMacaroon(bctx, encoded) + sc := ApplyToLLMPre(bctx) + if sc == nil || sc.Error == nil { + t.Fatal("expected short-circuit") + } + if sc.Error.StatusCode == nil || *sc.Error.StatusCode != 402 { + t.Fatalf("want 402, got %+v", sc.Error.StatusCode) + } + if sc.Error.Error == nil || sc.Error.Error.Code == nil || *sc.Error.Error.Code != "run_killed" { + t.Fatalf("want run_killed, got %+v", sc.Error.Error) + } + if sc.Error.Type == nil || *sc.Error.Type != "enforcement_rejected" { + t.Fatalf("want enforcement_rejected type, got %+v", sc.Error.Type) + } + if pluginctx.VerifiedClaims(bctx) != nil { + t.Fatal("rejected request must not stamp claims") + } +} + +func TestApplyToLLMPre_ShadowMode_KilledPassesThrough(t *testing.T) { + reg := newTestRegistry(t) + SetTrustRegistry(reg) + t.Cleanup(func() { SetTrustRegistry(nil) }) + SetConfigForTest(Config{EnforceMacaroons: false}) + t.Cleanup(func() { SetConfigForTest(Config{}) }) + mr := newMiniRedis(t) + + opts := defaultMacaroonOptions(time.Now()) + encoded := buildMacaroon(t, opts) + mr.Set("bifrost:kill:agent:coder", "1") + + bctx := newBifrostCtx() + pluginctx.SetRawMacaroon(bctx, encoded) + if sc := ApplyToLLMPre(bctx); sc != nil { + t.Fatalf("shadow mode must not short-circuit, got %+v", sc) + } +} diff --git a/gateway/internal/auth/revocation.go b/gateway/internal/auth/revocation.go index 99c5927d2..0cef667b0 100644 --- a/gateway/internal/auth/revocation.go +++ b/gateway/internal/auth/revocation.go @@ -24,6 +24,15 @@ const ( // timestamp. Any user_authorization.iat strictly before it is // rejected (covers org-level offboarding / user re-issuance). revokeUserBeforePrefix = "revoke_user_before:" + + // killRunPrefix + is the per-run kill switch. Checked + // for the leaf run AND every ancestor in the chain, so killing a + // parent kills every descendant. TTL 1h (hot operational state). + killRunPrefix = "kill:" + + // killAgentPrefix + halts every run whose leaf agent + // (Claims.AgentName, i.e. agents[last]) matches. TTL 24h. + killAgentPrefix = "kill:agent:" ) // pipelineTimeout bounds a single revocation pipeline round-trip. @@ -32,9 +41,12 @@ const ( // so we'd rather fail-closed quickly than block. const pipelineTimeout = 500 * time.Millisecond -// CheckRevocations runs the Redis revocation pipeline against the -// claims produced by Verify. Returns nil on "all clear", or an -// *AdapterError describing which nonce / user-level rule rejected. +// CheckRevocations runs the Redis revocation + kill-switch pipeline +// (phase-6 "Hot path" PIPELINE 1) against the claims produced by +// Verify. Returns nil on "all clear", or an *AdapterError describing +// which nonce / user-level rule / kill switch rejected. Revocations +// are 401s; kills are 402s (the macaroon is still valid — an +// operator chose to stop the run). // // Observability mode: when redisclient.Client() returns nil (no // REDIS_URL configured, or startup ping failed), CheckRevocations @@ -82,6 +94,19 @@ func CheckRevocations(ctx context.Context, claims *macaroon.Claims) *AdapterErro } userBeforeCmd := pipe.Get(pctx, redisclient.Key(revokeUserBeforePrefix+claims.UserID)) + // Kill switches ride the same round-trip. One EXISTS per distinct + // run_id in the chain (outermost first, so a parent kill reports + // the parent's id), plus one for the leaf agent name. + runIDs := chainRunIDs(claims) + killRunCmds := make([]*redis.IntCmd, len(runIDs)) + for i, id := range runIDs { + killRunCmds[i] = pipe.Exists(pctx, redisclient.Key(killRunPrefix+id)) + } + var killAgentCmd *redis.IntCmd + if claims.AgentName != "" { + killAgentCmd = pipe.Exists(pctx, redisclient.Key(killAgentPrefix+claims.AgentName)) + } + if _, err := pipe.Exec(pctx); err != nil && !errors.Is(err, redis.Nil) { // Exec returns the first non-Nil error; redis.Nil is the // expected sentinel when GET misses, so filter it out. @@ -152,9 +177,67 @@ func CheckRevocations(ctx context.Context, claims *macaroon.Claims) *AdapterErro // revocations above still apply). } + // Kill switches, after revocation so a revoked-and-killed + // macaroon reports the stronger (auth) reason. Per-run first, + // then per-agent, matching phase-6 "Hot path" step 2. + for i, cmd := range killRunCmds { + exists, err := cmd.Result() + if err != nil { + return &AdapterError{ + Code: "revocation_check_unavailable", + HTTPStatus: 401, + Message: fmt.Sprintf("redis exists kill:%s: %v", runIDs[i], err), + } + } + if exists == 1 { + return &AdapterError{ + Code: "run_killed", + HTTPStatus: 402, + Message: fmt.Sprintf("run %s was killed by an operator", runIDs[i]), + } + } + } + if killAgentCmd != nil { + exists, err := killAgentCmd.Result() + if err != nil { + return &AdapterError{ + Code: "revocation_check_unavailable", + HTTPStatus: 401, + Message: fmt.Sprintf("redis exists kill:agent:%s: %v", claims.AgentName, err), + } + } + if exists == 1 { + return &AdapterError{ + Code: "agent_killed", + HTTPStatus: 402, + Message: fmt.Sprintf("agent %s was killed by an operator", claims.AgentName), + } + } + } + return nil } +// chainRunIDs returns every distinct run_id in the verified chain, +// outermost (invocation) first, leaf last. Falls back to Claims.RunID +// when the chain is empty so callers that only populate the leaf +// still get their kill checked. +func chainRunIDs(claims *macaroon.Claims) []string { + ids := make([]string, 0, len(claims.Chain)+1) + seen := make(map[string]bool, len(claims.Chain)+1) + for _, layer := range claims.Chain { + if layer.RunID == "" || seen[layer.RunID] { + continue + } + seen[layer.RunID] = true + ids = append(ids, layer.RunID) + } + if claims.RunID != "" && !seen[claims.RunID] { + ids = append(ids, claims.RunID) + } + return ids +} + // revokeCodeFor maps the position of a revoked nonce in the // Claims.Nonces slice to a specific failure code, so an operator // reading a 401 can tell whether it was the user_authorization diff --git a/gateway/plans/phases/phase-6-plugin-enforcement.md b/gateway/plans/phases/phase-6-plugin-enforcement.md index cf7eb94c2..6bd33caf3 100644 --- a/gateway/plans/phases/phase-6-plugin-enforcement.md +++ b/gateway/plans/phases/phase-6-plugin-enforcement.md @@ -28,8 +28,20 @@ > core exposes no pricing manager to plugins — the canonical > logs.db cost is computed by the framework after our hook.) > The verifier now surfaces `Claims.UAIAT` / `UAExp` / `Chain` for -> this. Still open: the PreLLMHook cost/step cap walk and its 402s, -> tool-loop detection, kill switches + admin routes, and the +> this. +> +> **Status (kill switches landed):** PIPELINE 1 is complete — +> `CheckRevocations` now also issues `EXISTS bifrost:kill:` +> for every distinct run_id in `Claims.Chain` and +> `EXISTS bifrost:kill:agent:`, rejecting with +> 402 `run_killed` / `agent_killed` (revocation 401s win when both +> apply). Gated by `enforce_macaroons` like everything else in the +> hook; shadow mode logs. Admin routes (`gateway/internal/adminapi/ +> hotstate.go`, `revoke.go`): `/_plugin/runs/:id/{kill,state}`, +> `/_plugin/agents/:name/{kill,state}` (cookie-or-bearer, CSRF on +> cookie mutations) and `/_plugin/revoke/{nonce,user}/:id` +> (bearer-only). Still open: the PreLLMHook cost/step cap walk +> (PIPELINE 2) and its 402s, tool-loop detection, and the > `/_plugin/config/*` override layer. > > **Status (phase 11 cutover):** Redis bucket keys and hot-path