Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
66 changes: 66 additions & 0 deletions forge-cli/runtime/runner.go
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,7 @@ import (
"os/exec"
"path/filepath"
"strings"
"sync/atomic"
"time"

"github.com/initializ/forge/forge-cli/server"
Expand Down Expand Up @@ -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.
Expand Down Expand Up @@ -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() {

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Finding 1 (should-fix): this gate is applied here and on tasks/sendSubscribe (1688), but the two REST mirrors — POST /tasks/send (runner.go:2251) and POST /tasks/sendSubscribe (2304), both live via registerRESTHandlers — have no r.killed.Load() check. I read both: they decode the body and admit work unconditionally. So a killed agent still accepts new work over REST, silently defeating the accepting-gate contract on half the sync-A2A ingress surface. Add the same guard at the top of each REST handler. (I could not inline-anchor on 2251/2304 themselves — they are unchanged lines, not in this diff — hence the note here.)

return a2a.NewErrorResponse(id, a2a.ErrCodeUnavailable, "agent disabled by kill switch: not accepting new tasks")
}
var params a2a.SendTaskParams
if err := json.Unmarshal(rawParams, &params); err != nil {
return a2a.NewErrorResponse(id, a2a.ErrCodeInvalidParams, "invalid params: "+err.Error())
Expand Down Expand Up @@ -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, &params); err != nil {
server.WriteSSEEvent(w, flusher, "error", a2a.NewErrorResponse(id, a2a.ErrCodeInvalidParams, err.Error())) //nolint:errcheck
Expand Down Expand Up @@ -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, &params) // 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{

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Finding 2 (should-fix): the kill action is recorded only on the ops logger, so it never lands in the tamper-evident audit NDJSON. You already capture caller here — good — but if the agent is idle (cancelled=0) the kill emits NO audit event at all (no per-task invocation_cancelled either), leaving a destructive admin action with no forensic trail and the actor absent from the signed chain. Emit an admin_kill audit event via EmitFromContext (caller / reason / cancelled) unconditionally, independent of in-flight count.

"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
Expand Down Expand Up @@ -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()})
Expand Down Expand Up @@ -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"})
Expand Down
100 changes: 100 additions & 0 deletions forge-cli/runtime/runner_killswitch_test.go
Original file line number Diff line number Diff line change
@@ -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)
}
}
7 changes: 7 additions & 0 deletions forge-core/a2a/jsonrpc.go
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down
10 changes: 10 additions & 0 deletions forge-core/runtime/audit.go
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
34 changes: 33 additions & 1 deletion forge-core/runtime/cancellation.go
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand All @@ -49,7 +57,8 @@ func (r CancellationReason) IsValid() bool {
case CancelReasonWorkflowFailure,
CancelReasonCostLimitExceeded,
CancelReasonTimeout,
CancelReasonExternalSignal:
CancelReasonExternalSignal,
CancelReasonKillSwitch:
return true
}
return false
Expand Down Expand Up @@ -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
Expand Down
51 changes: 51 additions & 0 deletions forge-core/runtime/cancellation_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@ package runtime
import (
"context"
"errors"
"fmt"
"sync"
"testing"
"time"
Expand Down Expand Up @@ -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
Expand Down
Loading