diff --git a/forge-cli/runtime/runner.go b/forge-cli/runtime/runner.go index 9dce2b68..d1c60284 100644 --- a/forge-cli/runtime/runner.go +++ b/forge-cli/runtime/runner.go @@ -10,6 +10,7 @@ import ( "os/exec" "path/filepath" "strings" + "sync/atomic" "time" "github.com/initializ/forge/forge-cli/server" @@ -185,6 +186,7 @@ type Runner struct { standaloneSubjectStore mcp.SubjectTokenStore // #332 shared per-subject token cache: standalone resolver reads, callback writes; nil unless a standalone type:user server exists taskStore *a2a.TaskStore // shared task store, populated once srv is built; read by defer hook when it fires platformCommandGuard *coreruntime.PlatformCommandGuard // #238 (ASI02) operator-authored command deny, applied to every tool call; empty when no layer declares denied_command_patterns + killed atomic.Bool // kill switch: set by the admin/kill handler; when true, tasks/send + tasks/sendSubscribe refuse new work (in-flight work is cancelled via cancelRegistry.CancelAll, then the platform scales the workload to zero) } // NewRunner creates a Runner from the given config. @@ -1635,6 +1637,9 @@ func (r *Runner) registerHandlers(srv *server.Server, executor coreruntime.Agent // JSON-RPC path goes through the same audit + accumulator wiring as // REST POST /tasks/send. See issue #87 / FWS-3. srv.RegisterHandler("tasks/send", func(ctx context.Context, id any, rawParams json.RawMessage) *a2a.JSONRPCResponse { + if r.killed.Load() { + return a2a.NewErrorResponse(id, a2a.ErrCodeUnavailable, "agent disabled by kill switch: not accepting new tasks") + } var params a2a.SendTaskParams if err := json.Unmarshal(rawParams, ¶ms); err != nil { return a2a.NewErrorResponse(id, a2a.ErrCodeInvalidParams, "invalid params: "+err.Error()) @@ -1680,6 +1685,10 @@ func (r *Runner) registerHandlers(srv *server.Server, executor coreruntime.Agent // tasks/sendSubscribe — SSE streaming srv.RegisterSSEHandler("tasks/sendSubscribe", func(ctx context.Context, id any, rawParams json.RawMessage, w http.ResponseWriter, flusher http.Flusher) { + if r.killed.Load() { + server.WriteSSEEvent(w, flusher, "error", a2a.NewErrorResponse(id, a2a.ErrCodeUnavailable, "agent disabled by kill switch: not accepting new tasks")) //nolint:errcheck + return + } var params a2a.SendTaskParams if err := json.Unmarshal(rawParams, ¶ms); err != nil { server.WriteSSEEvent(w, flusher, "error", a2a.NewErrorResponse(id, a2a.ErrCodeInvalidParams, err.Error())) //nolint:errcheck @@ -1943,6 +1952,55 @@ func (r *Runner) registerHandlers(srv *server.Server, executor coreruntime.Agent // store has so the orchestrator reads the actual outcome. return a2a.NewResponse(id, task) }) + + // admin/kill — the agent kill switch. Flips the accepting gate so + // tasks/send + tasks/sendSubscribe refuse new work, then cancels + // EVERY in-flight invocation via cancelRegistry.CancelAll. Each + // cancelled invocation emits its own invocation_cancelled audit + // event with reason=kill_switch. The platform (agent-builder) calls + // this over the in-cluster A2A channel, then scales the workload to + // zero regardless of the outcome here. + // + // Auth: the server-wide AuthMiddleware already gates every JSON-RPC + // method, so only an authenticated caller reaches this handler; the + // primary access control is agent-builder's admin-RBAC on the + // /kill endpoint. TODO(kill-switch hardening): additionally restrict + // to the platform/agent-runtime identity via the verified role claim. + // Idempotent: a second kill just re-signals an empty registry (0). + srv.RegisterHandler("admin/kill", func(ctx context.Context, id any, rawParams json.RawMessage) *a2a.JSONRPCResponse { + var params struct { + Reason string `json:"reason"` + } + _ = json.Unmarshal(rawParams, ¶ms) // reason optional; body may be empty + reason := coreruntime.CancellationReason(params.Reason) + if reason == "" { + reason = coreruntime.CancelReasonKillSwitch + } + r.killed.Store(true) + cancelled := r.cancelRegistry.CancelAll(reason) + caller := "" + if idn := auth.IdentityFromContext(ctx); idn != nil { + caller = idn.Email + } + // Record the kill in the tamper-evident audit chain UNCONDITIONALLY — + // even when nothing was in flight (cancelled==0), so a destructive + // admin action never lacks a forensic record + actor. The cancelled + // invocations additionally each emit invocation_cancelled(kill_switch). + auditLogger.EmitFromContext(ctx, coreruntime.AuditEvent{ + Event: coreruntime.AuditAdminKilled, + Fields: map[string]any{ + "caller": caller, + "reason": string(reason), + "cancelled": cancelled, + }, + }) + r.logger.Info("admin/kill", map[string]any{ + "cancelled": cancelled, + "reason": string(reason), + "caller": caller, + }) + return a2a.NewResponse(id, map[string]any{"killed": true, "cancelled": cancelled}) + }) } // registerInvocationSeq exposes this invocation's sequence counter by @@ -2203,6 +2261,10 @@ func (r *Runner) registerRESTHandlers(srv *server.Server, executor coreruntime.A // POST /tasks/send — synchronous REST endpoint srv.RegisterHTTPHandler("POST /tasks/send", func(w http.ResponseWriter, req *http.Request) { + if r.killed.Load() { + writeJSON(w, http.StatusServiceUnavailable, map[string]string{"error": "agent disabled by kill switch: not accepting new tasks"}) + return + } var body restTaskRequest if err := json.NewDecoder(req.Body).Decode(&body); err != nil { writeJSON(w, http.StatusBadRequest, map[string]string{"error": "invalid request body: " + err.Error()}) @@ -2256,6 +2318,10 @@ func (r *Runner) registerRESTHandlers(srv *server.Server, executor coreruntime.A // POST /tasks/sendSubscribe — SSE streaming REST endpoint srv.RegisterHTTPHandler("POST /tasks/sendSubscribe", func(w http.ResponseWriter, req *http.Request) { + if r.killed.Load() { + writeJSON(w, http.StatusServiceUnavailable, map[string]string{"error": "agent disabled by kill switch: not accepting new tasks"}) + return + } flusher, ok := w.(http.Flusher) if !ok { writeJSON(w, http.StatusInternalServerError, map[string]string{"error": "streaming not supported"}) diff --git a/forge-cli/runtime/runner_killswitch_test.go b/forge-cli/runtime/runner_killswitch_test.go new file mode 100644 index 00000000..56af6d6e --- /dev/null +++ b/forge-cli/runtime/runner_killswitch_test.go @@ -0,0 +1,100 @@ +package runtime + +import ( + "context" + "encoding/json" + "fmt" + "net/http" + "testing" + "time" + + "github.com/initializ/forge/forge-core/a2a" + "github.com/initializ/forge/forge-core/auth" + "github.com/initializ/forge/forge-core/types" +) + +// TestRunner_KillSwitch_RefusesNewWorkOnEveryIngress is the kill-switch +// contract test: once admin/kill trips the gate, NONE of the four new-work +// ingress paths may admit a task. It guards specifically against Finding 1 of +// the #439 review — the JSON-RPC gate landing but the two REST mirrors +// (POST /tasks/send, POST /tasks/sendSubscribe) staying open. +func TestRunner_KillSwitch_RefusesNewWorkOnEveryIngress(t *testing.T) { + dir := t.TempDir() + cfg := &types.ForgeConfig{ + AgentID: "kill-switch-gate", + Version: "0.1.0", + Framework: "forge", + Entrypoint: "python main.py", + Tools: []types.ToolRef{{Name: "search"}}, + } + port, err := findFreePort() + if err != nil { + t.Fatal(err) + } + runner, err := NewRunner(RunnerConfig{Config: cfg, WorkDir: dir, Port: port, MockTools: true}) + if err != nil { + t.Fatalf("NewRunner: %v", err) + } + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + go func() { _ = runner.Run(ctx) }() + baseURL := fmt.Sprintf("http://localhost:%d", port) + waitForServer(t, baseURL, 5*time.Second) + token, _ := auth.LoadToken(dir) + + // Trip the kill switch. Idle agent → cancelled=0, but the call must still + // succeed and flip the gate (and, per Finding 2, emit admin_killed regardless). + killBody := []byte(`{"jsonrpc":"2.0","id":1,"method":"admin/kill","params":{"reason":"test"}}`) + resp, err := authPost(baseURL+"/", token, killBody) + if err != nil { + t.Fatalf("admin/kill: %v", err) + } + var killResp a2a.JSONRPCResponse + _ = json.NewDecoder(resp.Body).Decode(&killResp) + _ = resp.Body.Close() + if killResp.Error != nil { + t.Fatalf("admin/kill returned error: %+v", killResp.Error) + } + result, _ := killResp.Result.(map[string]any) + if killed, _ := result["killed"].(bool); !killed { + t.Fatalf("admin/kill result should report killed=true, got %v", killResp.Result) + } + + send := []byte(`{"jsonrpc":"2.0","id":2,"method":"tasks/send","params":{"id":"t-after-kill","message":{"role":"user","parts":[{"kind":"text","text":"hi"}]}}}`) + + // 1. JSON-RPC tasks/send → Unavailable error, work NOT admitted. + r1, err := authPost(baseURL+"/", token, send) + if err != nil { + t.Fatalf("tasks/send: %v", err) + } + var rpc a2a.JSONRPCResponse + _ = json.NewDecoder(r1.Body).Decode(&rpc) + _ = r1.Body.Close() + if rpc.Error == nil { + t.Fatalf("JSON-RPC tasks/send after kill must error; got result=%v", rpc.Result) + } + if rpc.Error.Code != a2a.ErrCodeUnavailable { + t.Errorf("JSON-RPC error.code = %d, want %d (Unavailable)", rpc.Error.Code, a2a.ErrCodeUnavailable) + } + + // 2. REST POST /tasks/send → 503, work NOT admitted (the Finding-1 gap). + rest := []byte(`{"task":{"id":"t-rest-after-kill","message":{"role":"user","parts":[{"kind":"text","text":"hi"}]}}}`) + r2, err := authPost(baseURL+"/tasks/send", token, rest) + if err != nil { + t.Fatalf("REST tasks/send: %v", err) + } + _ = r2.Body.Close() + if r2.StatusCode != http.StatusServiceUnavailable { + t.Errorf("REST POST /tasks/send after kill: status = %d, want 503", r2.StatusCode) + } + + // 3. REST POST /tasks/sendSubscribe → 503. + r3, err := authPost(baseURL+"/tasks/sendSubscribe", token, rest) + if err != nil { + t.Fatalf("REST tasks/sendSubscribe: %v", err) + } + _ = r3.Body.Close() + if r3.StatusCode != http.StatusServiceUnavailable { + t.Errorf("REST POST /tasks/sendSubscribe after kill: status = %d, want 503", r3.StatusCode) + } +} diff --git a/forge-core/a2a/jsonrpc.go b/forge-core/a2a/jsonrpc.go index 5a0fadd6..1bcbec52 100644 --- a/forge-core/a2a/jsonrpc.go +++ b/forge-core/a2a/jsonrpc.go @@ -9,6 +9,13 @@ const ( ErrCodeMethodNotFound = -32601 ErrCodeInvalidParams = -32602 ErrCodeInternal = -32603 + + // ErrCodeUnavailable is a server-defined code (JSON-RPC reserves + // -32000..-32099 for implementation-defined server errors). Signals a + // deliberate refusal to serve — e.g. the agent kill switch is tripped + // and the runtime is not accepting new tasks — as distinct from an + // unexpected internal fault (ErrCodeInternal). + ErrCodeUnavailable = -32000 ) // JSONRPCRequest is an incoming JSON-RPC 2.0 request. diff --git a/forge-core/runtime/audit.go b/forge-core/runtime/audit.go index 29ff65c3..2ea5ce63 100644 --- a/forge-core/runtime/audit.go +++ b/forge-core/runtime/audit.go @@ -205,6 +205,16 @@ const ( // calls completed before the cancel signal. See issue #88 / FWS-4. AuditInvocationCancelled = "invocation_cancelled" + // AuditAdminKilled is emitted when the agent kill switch is tripped via + // the admin/kill A2A verb. Recorded UNCONDITIONALLY — even when no + // invocation was in flight (Fields["cancelled"] == 0) — so a + // destructive admin action always has a forensic record in the + // tamper-evident chain, with the actor. Fields: caller (verified + // email, "" if unauthenticated context), reason, cancelled (count of + // in-flight invocations signalled). Each of those invocations also + // emits its own invocation_cancelled with reason=kill_switch. + AuditAdminKilled = "admin_killed" + // AuditTaskAdmissionDenied is emitted when the admission middleware // rejects an inbound A2A invocation based on a platform-side quota // / cost-limit decision (issue #201). Carries the platform's diff --git a/forge-core/runtime/cancellation.go b/forge-core/runtime/cancellation.go index 9dd2fda2..73696ab1 100644 --- a/forge-core/runtime/cancellation.go +++ b/forge-core/runtime/cancellation.go @@ -38,6 +38,14 @@ const ( // cancel, debugging stop, anything else not covered by the more // specific reasons. CancelReasonExternalSignal CancellationReason = "external_signal" + + // CancelReasonKillSwitch is set when an operator trips the agent + // kill switch: every in-flight invocation is cancelled via + // CancellationRegistry.CancelAll and the workload is then scaled to + // zero by the platform. Distinct from external_signal so the + // invocation_cancelled audit event attributes the stop to the kill + // switch specifically (not a per-task operator cancel). + CancelReasonKillSwitch CancellationReason = "kill_switch" ) // IsValid reports whether r is one of the documented reason values. @@ -49,7 +57,8 @@ func (r CancellationReason) IsValid() bool { case CancelReasonWorkflowFailure, CancelReasonCostLimitExceeded, CancelReasonTimeout, - CancelReasonExternalSignal: + CancelReasonExternalSignal, + CancelReasonKillSwitch: return true } return false @@ -167,6 +176,29 @@ func (r *CancellationRegistry) Cancel(taskID string, reason CancellationReason) return true } +// CancelAll signals every in-flight invocation with reason and returns the +// number signalled. This is the kill-switch primitive: the admin/kill handler +// calls it to abort all active work on the agent at once, before the platform +// scales the workload to zero. +// +// The cancel funcs are snapshotted under the lock and invoked outside it — +// matching Cancel's contention profile and avoiding any chance of a cancel +// callback re-entering the registry under the held lock. Each cancelled +// invocation's own deferred release() then pops its entry as executeTask +// unwinds; CancelAll does not delete entries itself. +func (r *CancellationRegistry) CancelAll(reason CancellationReason) int { + r.mu.Lock() + cancels := make([]context.CancelCauseFunc, 0, len(r.entries)) + for _, e := range r.entries { + cancels = append(cancels, e.cancel) + } + r.mu.Unlock() + for _, cancel := range cancels { + cancel(&cancelledByOrchestrator{Reason: reason}) + } + return len(cancels) +} + // Len returns the number of in-flight registrations. Exposed for // tests and operational observability — there is no per-task lookup // API by design (the handler only needs Cancel; the executeTask diff --git a/forge-core/runtime/cancellation_test.go b/forge-core/runtime/cancellation_test.go index 8d0964c9..e1fc5312 100644 --- a/forge-core/runtime/cancellation_test.go +++ b/forge-core/runtime/cancellation_test.go @@ -3,6 +3,7 @@ package runtime import ( "context" "errors" + "fmt" "sync" "testing" "time" @@ -39,6 +40,56 @@ func TestCancellationRegistry_RegisterCancelReleaseLifecycle(t *testing.T) { } } +func TestCancellationRegistry_CancelAll_SignalsEveryInflightWithReason(t *testing.T) { + // The kill-switch primitive: CancelAll must cancel every registered + // invocation, stamp the given reason on each cause, and report the count. + reg := NewCancellationRegistry() + const n = 5 + ctxs := make([]context.Context, n) + releases := make([]func(), n) + for i := 0; i < n; i++ { + ctx, cancel := context.WithCancelCause(context.Background()) + ctxs[i] = ctx + releases[i] = reg.Register(fmt.Sprintf("task-%d", i), cancel) + } + if reg.Len() != n { + t.Fatalf("Len()=%d, want %d", reg.Len(), n) + } + + if got := reg.CancelAll(CancelReasonKillSwitch); got != n { + t.Errorf("CancelAll returned %d, want %d", got, n) + } + for i, ctx := range ctxs { + if ctx.Err() == nil { + t.Errorf("ctx[%d] should be cancelled after CancelAll", i) + } + if r := CancellationReasonFromCause(ctx); r != CancelReasonKillSwitch { + t.Errorf("ctx[%d] reason=%q, want kill_switch", i, r) + } + } + // CancelAll does not delete entries; each invocation's own release() does. + for _, release := range releases { + release() + } + if reg.Len() != 0 { + t.Errorf("after releases: Len()=%d, want 0", reg.Len()) + } +} + +func TestCancellationRegistry_CancelAll_EmptyIsZero(t *testing.T) { + // A second kill (or a kill with no in-flight work) is a no-op returning 0. + reg := NewCancellationRegistry() + if got := reg.CancelAll(CancelReasonKillSwitch); got != 0 { + t.Errorf("CancelAll on empty registry returned %d, want 0", got) + } +} + +func TestCancellationReason_KillSwitchIsValid(t *testing.T) { + if !CancelReasonKillSwitch.IsValid() { + t.Errorf("CancelReasonKillSwitch should be a documented (valid) reason") + } +} + func TestCancellationRegistry_CancelUnknownTaskIsIdempotent(t *testing.T) { // Cancel-after-complete must be a no-op (returns false) so the // orchestrator can issue cancels optimistically without races