From 286597f53f9bb454803416afac98de9d2ec543d3 Mon Sep 17 00:00:00 2001 From: Adarsh Prashar Date: Sat, 30 May 2026 23:12:25 +0530 Subject: [PATCH] feat: Python SDK (Surface 2) + run-control API + framework adapters Deep control over a governed run from Python. The SDK is a thin, stdlib-only client; the Go core makes every deterministic decision. Go (run-control endpoints the SDK drives): - POST /v1/runs (create with budget), POST /v1/runs/{id}/steps (loop/time enforcement -> 402 on halt), POST /v1/runs/{id}/checkpoints, POST /v1/runs/{id}/cancel, POST /v1/runs/{id}/approvals (request, non-blocking), GET /v1/approvals/{id} (poll). approval.Gate.Create added for the poll model. - api/v1 extended additively; tests for every endpoint incl. loop-budget 402, checkpoint round-trip, approval request->poll->resolve. Python (sdks/python, package `riskkernel`, zero runtime deps): - client.RiskKernel (urllib); 402 -> BudgetExceeded. - Runtime / governed_run / Budget / Run (step, checkpoint, cancel, proxy_config); ApprovalGate + @governed_tool; current_run contextvar; env-configured default. - Adapters (lazy-imported, no hard deps): LangChain CallbackHandler, Claude Agent SDK PreToolUse hook, OpenAI Agents SDK RunHooks. - pyproject (hatchling, Apache-2.0), README, unittest suite against a stdlib stub daemon. Verified end-to-end against the real daemon. - CI: Python SDK workflow on 3.9 / 3.12. Model-call metering stays in the proxy: route the SDK's LLM calls through run.proxy_config(), so the SDK never re-implements governance. --- .github/workflows/python.yml | 32 +++ CHANGELOG.md | 9 + api/v1/openapi.yaml | 140 ++++++++++++ internal/approval/gate.go | 31 +++ internal/httpapi/runs.go | 199 ++++++++++++++++++ internal/httpapi/runs_test.go | 172 +++++++++++++++ internal/httpapi/server.go | 8 +- sdks/python/README.md | 81 +++++++ sdks/python/pyproject.toml | 32 +++ sdks/python/riskkernel/__init__.py | 59 ++++++ sdks/python/riskkernel/adapters/__init__.py | 8 + .../riskkernel/adapters/claude_agent.py | 75 +++++++ sdks/python/riskkernel/adapters/langchain.py | 75 +++++++ .../riskkernel/adapters/openai_agents.py | 49 +++++ sdks/python/riskkernel/approval.py | 86 ++++++++ sdks/python/riskkernel/client.py | 124 +++++++++++ sdks/python/riskkernel/errors.py | 40 ++++ sdks/python/riskkernel/runtime.py | 199 ++++++++++++++++++ sdks/python/tests/__init__.py | 0 sdks/python/tests/test_sdk.py | 152 +++++++++++++ 20 files changed, 1570 insertions(+), 1 deletion(-) create mode 100644 .github/workflows/python.yml create mode 100644 internal/httpapi/runs.go create mode 100644 internal/httpapi/runs_test.go create mode 100644 sdks/python/README.md create mode 100644 sdks/python/pyproject.toml create mode 100644 sdks/python/riskkernel/__init__.py create mode 100644 sdks/python/riskkernel/adapters/__init__.py create mode 100644 sdks/python/riskkernel/adapters/claude_agent.py create mode 100644 sdks/python/riskkernel/adapters/langchain.py create mode 100644 sdks/python/riskkernel/adapters/openai_agents.py create mode 100644 sdks/python/riskkernel/approval.py create mode 100644 sdks/python/riskkernel/client.py create mode 100644 sdks/python/riskkernel/errors.py create mode 100644 sdks/python/riskkernel/runtime.py create mode 100644 sdks/python/tests/__init__.py create mode 100644 sdks/python/tests/test_sdk.py diff --git a/.github/workflows/python.yml b/.github/workflows/python.yml new file mode 100644 index 0000000..a0180e5 --- /dev/null +++ b/.github/workflows/python.yml @@ -0,0 +1,32 @@ +name: Python SDK + +on: + push: + branches: [main] + paths: ["sdks/python/**", ".github/workflows/python.yml"] + pull_request: + branches: [main] + paths: ["sdks/python/**", ".github/workflows/python.yml"] + +permissions: + contents: read + +jobs: + test: + name: test (py${{ matrix.python }}) + runs-on: ubuntu-latest + strategy: + matrix: + python: ["3.9", "3.12"] + defaults: + run: + working-directory: sdks/python + steps: + - uses: actions/checkout@v6 + - uses: actions/setup-python@v5 + with: + python-version: ${{ matrix.python }} + - name: Install (editable, stdlib-only core) + run: pip install -e . + - name: Test + run: python -m unittest discover -s tests -t . -v diff --git a/CHANGELOG.md b/CHANGELOG.md index dfb36b1..43e8fb9 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -62,5 +62,14 @@ surface is governed by [`COMPATIBILITY.md`](COMPATIBILITY.md). (surfaces `pendingApproval` + `waiting_approval` status). Approvals are persisted (migration `00003`) as an audit trail. Webhook is user-configured egress only (see SECURITY.md). +- **Run-control API** — `POST /v1/runs` (create with budget), + `POST /v1/runs/{id}/steps` (loop/time enforcement → 402 on halt), + `POST /v1/runs/{id}/checkpoints`, `POST /v1/runs/{id}/cancel`, + `POST /v1/runs/{id}/approvals` (request → poll), `GET /v1/approvals/{id}`. +- **Python SDK (Surface 2)** — `pip install riskkernel`, a stdlib-only thin client: + `Runtime`, `governed_run`, `Budget`, `Run.step/checkpoint/cancel/proxy_config`, + `ApprovalGate`, `@governed_tool`. Lazy-imported framework adapters for LangChain + (callback handler), the Claude Agent SDK (PreToolUse hook), and the OpenAI Agents + SDK (RunHooks). Verified end-to-end against the daemon; CI on Python 3.9/3.12. [Unreleased]: https://github.com/prashar32/riskkernel/commits/main diff --git a/api/v1/openapi.yaml b/api/v1/openapi.yaml index b7838a5..dc5621b 100644 --- a/api/v1/openapi.yaml +++ b/api/v1/openapi.yaml @@ -123,6 +123,146 @@ paths: schema: $ref: '#/components/schemas/Error' + /v1/runs/{id}/steps: + post: + tags: [runs] + operationId: beginStep + summary: Register a loop iteration (enforce loop + time budgets) + description: | + The Python SDK calls this once per agent loop iteration. It increments the + loop counter and enforces the loop and wall-clock budgets. Returns 402 with + the HaltReason when the run is out of budget. + parameters: + - $ref: '#/components/parameters/RunId' + responses: + '200': + description: Step registered. + content: + application/json: + schema: + type: object + properties: + stepIndex: + type: integer + format: int32 + '402': + description: Run halted — loop or time budget exceeded. + content: + application/json: + schema: + $ref: '#/components/schemas/Error' + '401': + $ref: '#/components/responses/Unauthorized' + '404': + $ref: '#/components/responses/NotFound' + + /v1/runs/{id}/checkpoints: + post: + tags: [checkpoints] + operationId: saveCheckpoint + summary: Save a crash-resumable checkpoint + description: | + Persists an opaque payload (e.g. conversation messages, scratch state) + plus a usage snapshot, so a crashed run can resume from here. + parameters: + - $ref: '#/components/parameters/RunId' + requestBody: + required: true + content: + application/json: + schema: + type: object + properties: + name: + type: string + payload: + type: object + additionalProperties: true + responses: + '201': + description: Checkpoint saved. + '401': + $ref: '#/components/responses/Unauthorized' + '404': + $ref: '#/components/responses/NotFound' + + /v1/runs/{id}/approvals: + post: + tags: [approvals] + operationId: requestApproval + summary: Request approval for a (possibly side-effecting) tool call + description: | + The SDK's @governed_tool path. If policy allows the call, returns status + "approved" immediately. Otherwise creates a pending approval the caller + polls via GET /v1/approvals/{id} (resolved by a human via + POST /v1/runs/{id}/approve). + parameters: + - $ref: '#/components/parameters/RunId' + requestBody: + required: true + content: + application/json: + schema: + type: object + required: [tool] + properties: + tool: + type: string + sideEffect: + type: string + stepIndex: + type: integer + format: int32 + arguments: + type: object + additionalProperties: true + responses: + '200': + description: Allowed by policy — no human approval required. + content: + application/json: + schema: + type: object + properties: + status: + type: string + required: + type: boolean + '201': + description: Approval is pending a human decision. + content: + application/json: + schema: + $ref: '#/components/schemas/ApprovalRequest' + '401': + $ref: '#/components/responses/Unauthorized' + '404': + $ref: '#/components/responses/NotFound' + + /v1/approvals/{id}: + get: + tags: [approvals] + operationId: getApproval + summary: Get a single approval (poll for resolution) + parameters: + - name: id + in: path + required: true + schema: + type: string + format: uuid + responses: + '200': + description: The approval and its current status. + content: + application/json: + schema: + $ref: '#/components/schemas/ApprovalRequest' + '401': + $ref: '#/components/responses/Unauthorized' + '404': + $ref: '#/components/responses/NotFound' + /v1/runs/{id}/cancel: post: tags: [runs] diff --git a/internal/approval/gate.go b/internal/approval/gate.go index fe389bc..d7fa705 100644 --- a/internal/approval/gate.go +++ b/internal/approval/gate.go @@ -113,6 +113,37 @@ func (g *Gate) Request(ctx context.Context, req Request) (Decision, string, erro } } +// Create evaluates policy and, if approval is required, persists a pending +// approval and notifies push channels, returning the record (required=true). +// Unlike Request, it does NOT block — the caller (e.g. the SDK over HTTP) polls +// the approval's status instead. Returns required=false when policy allows the +// call outright. +func (g *Gate) Create(ctx context.Context, req Request) (storage.ApprovalRecord, bool, error) { + if !g.policy.Requires(req.Tool, req.SideEffect) { + return storage.ApprovalRecord{}, false, nil + } + rec := storage.ApprovalRecord{ + ID: g.newID(), + RunID: req.RunID, + StepIndex: req.StepIndex, + Tool: req.Tool, + SideEffect: req.SideEffect, + Arguments: req.Arguments, + Status: storage.ApprovalPending, + CreatedAt: g.now(), + } + if g.store != nil { + if err := g.store.CreateApproval(ctx, rec); err != nil { + return storage.ApprovalRecord{}, false, err + } + } + if g.notifier != nil { + g.notifier.Notify(ctx, rec) + } + g.log.Info("approval required", "id", rec.ID, "run", rec.RunID, "tool", rec.Tool, "side_effect", rec.SideEffect) + return rec, true, nil +} + // Resolve records a human decision and wakes any blocked waiter. Returns // storage.ErrNotFound if the approval is unknown or already resolved. func (g *Gate) Resolve(ctx context.Context, id string, approved bool, reason, by string) error { diff --git a/internal/httpapi/runs.go b/internal/httpapi/runs.go new file mode 100644 index 0000000..2a68627 --- /dev/null +++ b/internal/httpapi/runs.go @@ -0,0 +1,199 @@ +package httpapi + +import ( + "encoding/json" + "errors" + "net/http" + + "github.com/prashar32/riskkernel/internal/approval" + "github.com/prashar32/riskkernel/internal/governor" + "github.com/prashar32/riskkernel/internal/httpx" + "github.com/prashar32/riskkernel/internal/runs" + "github.com/prashar32/riskkernel/internal/storage" +) + +// These endpoints let the Python SDK (Surface 2) drive a governed run directly: +// create it, tick loop iterations, checkpoint state, request tool approval, and +// cancel. Model-call metering/cost/budget is handled by routing the SDK's LLM +// calls through the proxy with the run-id header, so the SDK stays thin. + +type budgetBody struct { + Tokens int64 `json:"tokens"` + Dollars float64 `json:"dollars"` + Loops int32 `json:"loops"` + Seconds int32 `json:"seconds"` +} + +type createRunBody struct { + Name string `json:"name"` + Budget *budgetBody `json:"budget"` + Metadata map[string]string `json:"metadata"` +} + +// handleCreateRun implements POST /v1/runs. +func (s *Server) handleCreateRun(w http.ResponseWriter, r *http.Request) { + var body createRunBody + if err := decodeJSON(w, r, &body); err != nil { + httpx.WriteError(w, http.StatusBadRequest, "bad_request", err.Error()) + return + } + opts := runs.CreateOptions{Name: body.Name, Metadata: body.Metadata} + if body.Budget != nil { + opts.Budget = &governor.Budget{ + Tokens: body.Budget.Tokens, Dollars: body.Budget.Dollars, + Loops: body.Budget.Loops, Seconds: body.Budget.Seconds, + } + } + run := s.runs.Create(opts) + httpx.WriteJSON(w, http.StatusCreated, runViewFromManager(run)) +} + +// handleBeginStep implements POST /v1/runs/{id}/steps — registers a loop +// iteration and enforces the loop + time budgets. 402 when the budget is spent. +func (s *Server) handleBeginStep(w http.ResponseWriter, r *http.Request) { + run, ok := s.runs.Get(r.PathValue("id")) + if !ok { + httpx.WriteError(w, http.StatusNotFound, "not_found", "run not found") + return + } + step, err := run.BeginStep() + if err != nil { + writeHalt(w, err) + return + } + httpx.WriteJSON(w, http.StatusOK, map[string]any{"stepIndex": step}) +} + +type checkpointBody struct { + Name string `json:"name"` + Payload map[string]any `json:"payload"` +} + +// handleSaveCheckpoint implements POST /v1/runs/{id}/checkpoints. +func (s *Server) handleSaveCheckpoint(w http.ResponseWriter, r *http.Request) { + runID := r.PathValue("id") + if _, ok := s.runs.Get(runID); !ok { + httpx.WriteError(w, http.StatusNotFound, "not_found", "run not found") + return + } + var body checkpointBody + if err := decodeJSON(w, r, &body); err != nil { + httpx.WriteError(w, http.StatusBadRequest, "bad_request", err.Error()) + return + } + if err := s.runs.Checkpoint(runID, body.Name, body.Payload); err != nil { + httpx.WriteError(w, http.StatusInternalServerError, "internal_error", err.Error()) + return + } + httpx.WriteJSON(w, http.StatusCreated, map[string]any{"ok": true}) +} + +// handleCancelRun implements POST /v1/runs/{id}/cancel (kill switch). +func (s *Server) handleCancelRun(w http.ResponseWriter, r *http.Request) { + run, ok := s.runs.Get(r.PathValue("id")) + if !ok { + httpx.WriteError(w, http.StatusNotFound, "not_found", "run not found") + return + } + run.Cancel() + httpx.WriteJSON(w, http.StatusOK, runViewFromManager(run)) +} + +type toolApprovalBody struct { + Tool string `json:"tool"` + SideEffect string `json:"sideEffect"` + StepIndex int32 `json:"stepIndex"` + Arguments map[string]any `json:"arguments"` +} + +// handleRequestApproval implements POST /v1/runs/{id}/approvals — the SDK's +// @governed_tool path. If policy allows the call, returns status "approved" +// immediately; otherwise creates a pending approval the SDK polls via +// GET /v1/approvals/{id}. +func (s *Server) handleRequestApproval(w http.ResponseWriter, r *http.Request) { + runID := r.PathValue("id") + if _, ok := s.runs.Get(runID); !ok { + httpx.WriteError(w, http.StatusNotFound, "not_found", "run not found") + return + } + var body toolApprovalBody + if err := decodeJSON(w, r, &body); err != nil { + httpx.WriteError(w, http.StatusBadRequest, "bad_request", err.Error()) + return + } + if body.Tool == "" { + httpx.WriteError(w, http.StatusBadRequest, "bad_request", "tool is required") + return + } + rec, required, err := s.approvals.Create(r.Context(), approval.Request{ + RunID: runID, StepIndex: body.StepIndex, Tool: body.Tool, + SideEffect: body.SideEffect, Arguments: body.Arguments, + }) + if err != nil { + httpx.WriteError(w, http.StatusInternalServerError, "internal_error", err.Error()) + return + } + if !required { + httpx.WriteJSON(w, http.StatusOK, map[string]any{"status": "approved", "required": false}) + return + } + httpx.WriteJSON(w, http.StatusCreated, approvalView(rec)) +} + +// handleGetApproval implements GET /v1/approvals/{id} — the SDK polls this for a +// pending approval's resolution. +func (s *Server) handleGetApproval(w http.ResponseWriter, r *http.Request) { + a, err := s.approvals.Get(r.Context(), r.PathValue("id")) + if err != nil { + if errors.Is(err, storage.ErrNotFound) { + httpx.WriteError(w, http.StatusNotFound, "not_found", "approval not found") + return + } + httpx.WriteError(w, http.StatusInternalServerError, "internal_error", err.Error()) + return + } + httpx.WriteJSON(w, http.StatusOK, approvalView(a)) +} + +// --- helpers --- + +func decodeJSON(w http.ResponseWriter, r *http.Request, v any) error { + dec := json.NewDecoder(http.MaxBytesReader(w, r.Body, 1<<20)) + if err := dec.Decode(v); err != nil { + return errors.New("invalid JSON: " + err.Error()) + } + return nil +} + +// writeHalt translates a governor.HaltError into a 402 response. +func writeHalt(w http.ResponseWriter, err error) { + var he *governor.HaltError + if errors.As(err, &he) { + httpx.WriteError(w, http.StatusPaymentRequired, string(he.Reason), "run halted: "+string(he.Reason)) + return + } + httpx.WriteError(w, http.StatusInternalServerError, "internal_error", err.Error()) +} + +func runViewFromManager(run *runs.Run) map[string]any { + v := run.View() + return map[string]any{ + "id": v.ID, + "name": v.Name, + "status": v.Status, + "haltReason": string(v.HaltReason), + "budget": map[string]any{ + "tokens": v.Budget.Tokens, "dollars": v.Budget.Dollars, + "loops": v.Budget.Loops, "seconds": v.Budget.Seconds, + }, + "usage": map[string]any{ + "tokens": v.Usage.Tokens(), + "promptTokens": v.Usage.PromptTokens, + "completionTokens": v.Usage.CompletionTokens, + "dollars": v.Usage.Dollars, + "loops": v.Usage.Loops, + }, + "createdAt": v.CreatedAt, + "updatedAt": v.UpdatedAt, + } +} diff --git a/internal/httpapi/runs_test.go b/internal/httpapi/runs_test.go new file mode 100644 index 0000000..646c743 --- /dev/null +++ b/internal/httpapi/runs_test.go @@ -0,0 +1,172 @@ +package httpapi + +import ( + "encoding/json" + "net/http" + "net/http/httptest" + "strings" + "testing" +) + +func do(t *testing.T, h http.Handler, method, path, body string) *httptest.ResponseRecorder { + t.Helper() + var r *http.Request + if body == "" { + r = httptest.NewRequest(method, path, nil) + } else { + r = httptest.NewRequest(method, path, strings.NewReader(body)) + } + w := httptest.NewRecorder() + h.ServeHTTP(w, r) + return w +} + +func TestCreateRun(t *testing.T) { + srv, _, _ := newTestServer(t, "") + h := srv.Handler() + w := do(t, h, http.MethodPost, "/v1/runs", `{"name":"sdk-run","budget":{"tokens":1000,"loops":3}}`) + if w.Code != http.StatusCreated { + t.Fatalf("status = %d, body=%s", w.Code, w.Body.String()) + } + var run map[string]any + _ = json.Unmarshal(w.Body.Bytes(), &run) + if run["id"] == "" || run["status"] != "running" { + t.Fatalf("run = %v", run) + } + budget := run["budget"].(map[string]any) + if budget["tokens"].(float64) != 1000 { + t.Errorf("budget = %v", budget) + } +} + +func TestBeginStep_LoopBudget(t *testing.T) { + srv, _, _ := newTestServer(t, "") + h := srv.Handler() + w := do(t, h, http.MethodPost, "/v1/runs", `{"budget":{"loops":2}}`) + var run map[string]any + _ = json.Unmarshal(w.Body.Bytes(), &run) + id := run["id"].(string) + + for i := 1; i <= 2; i++ { + w := do(t, h, http.MethodPost, "/v1/runs/"+id+"/steps", "") + if w.Code != http.StatusOK { + t.Fatalf("step %d status = %d", i, w.Code) + } + } + // 3rd step exceeds the loop budget. + w = do(t, h, http.MethodPost, "/v1/runs/"+id+"/steps", "") + if w.Code != http.StatusPaymentRequired { + t.Fatalf("3rd step status = %d, want 402; body=%s", w.Code, w.Body.String()) + } + var errBody struct{ Code string } + _ = json.Unmarshal(w.Body.Bytes(), &errBody) + if errBody.Code != "loop_budget_exceeded" { + t.Errorf("error code = %q", errBody.Code) + } +} + +func TestCheckpointRoundTrip(t *testing.T) { + srv, _, _ := newTestServer(t, "") + h := srv.Handler() + w := do(t, h, http.MethodPost, "/v1/runs", `{"name":"cp"}`) + var run map[string]any + _ = json.Unmarshal(w.Body.Bytes(), &run) + id := run["id"].(string) + + w = do(t, h, http.MethodPost, "/v1/runs/"+id+"/checkpoints", + `{"name":"after-plan","payload":{"messages":["hi"],"cursor":3}}`) + if w.Code != http.StatusCreated { + t.Fatalf("checkpoint status = %d, body=%s", w.Code, w.Body.String()) + } + + w = do(t, h, http.MethodGet, "/v1/checkpoints/"+id, "") + if w.Code != http.StatusOK { + t.Fatalf("get checkpoint status = %d", w.Code) + } + var cp map[string]any + _ = json.Unmarshal(w.Body.Bytes(), &cp) + payload, _ := cp["payload"].(map[string]any) + if payload["cursor"].(float64) != 3 { + t.Fatalf("checkpoint payload not persisted: %v", cp) + } +} + +func TestCancelRun(t *testing.T) { + srv, _, _ := newTestServer(t, "") + h := srv.Handler() + w := do(t, h, http.MethodPost, "/v1/runs", `{}`) + var run map[string]any + _ = json.Unmarshal(w.Body.Bytes(), &run) + id := run["id"].(string) + + w = do(t, h, http.MethodPost, "/v1/runs/"+id+"/cancel", `{"reason":"manual"}`) + if w.Code != http.StatusOK { + t.Fatalf("cancel status = %d", w.Code) + } + var got map[string]any + _ = json.Unmarshal(w.Body.Bytes(), &got) + if got["status"] != "cancelled" { + t.Fatalf("status after cancel = %v", got["status"]) + } +} + +func TestRequestApproval_PollResolve(t *testing.T) { + srv, _, _ := newTestServer(t, "") // gate is DefaultSafe:true + h := srv.Handler() + w := do(t, h, http.MethodPost, "/v1/runs", `{}`) + var run map[string]any + _ = json.Unmarshal(w.Body.Bytes(), &run) + id := run["id"].(string) + + // A side-effecting tool needs approval → 201 pending. + w = do(t, h, http.MethodPost, "/v1/runs/"+id+"/approvals", + `{"tool":"mcp://shell","sideEffect":"exec","arguments":{"cmd":"ls"}}`) + if w.Code != http.StatusCreated { + t.Fatalf("request approval status = %d, body=%s", w.Code, w.Body.String()) + } + var ap map[string]any + _ = json.Unmarshal(w.Body.Bytes(), &ap) + apID := ap["id"].(string) + if ap["status"] != "pending" { + t.Fatalf("approval status = %v", ap["status"]) + } + + // Poll → still pending. + w = do(t, h, http.MethodGet, "/v1/approvals/"+apID, "") + _ = json.Unmarshal(w.Body.Bytes(), &ap) + if ap["status"] != "pending" { + t.Fatalf("poll status = %v", ap["status"]) + } + + // Resolve via the HITL endpoint, then poll → approved. + w = do(t, h, http.MethodPost, "/v1/runs/"+id+"/approve", + `{"approvalId":"`+apID+`","decision":"approve","decidedBy":"tester"}`) + if w.Code != http.StatusOK { + t.Fatalf("approve status = %d", w.Code) + } + w = do(t, h, http.MethodGet, "/v1/approvals/"+apID, "") + _ = json.Unmarshal(w.Body.Bytes(), &ap) + if ap["status"] != "approved" { + t.Fatalf("final status = %v", ap["status"]) + } +} + +func TestRequestApproval_NotRequired(t *testing.T) { + srv, _, _ := newTestServer(t, "") + h := srv.Handler() + w := do(t, h, http.MethodPost, "/v1/runs", `{}`) + var run map[string]any + _ = json.Unmarshal(w.Body.Bytes(), &run) + id := run["id"].(string) + + // Read-only tool (no side effect) → auto-approved, no pending row. + w = do(t, h, http.MethodPost, "/v1/runs/"+id+"/approvals", `{"tool":"mcp://fs","sideEffect":""}`) + if w.Code != http.StatusOK { + t.Fatalf("status = %d, want 200 (auto-approve)", w.Code) + } + var ap map[string]any + _ = json.Unmarshal(w.Body.Bytes(), &ap) + if ap["status"] != "approved" || ap["required"] != false { + t.Fatalf("auto-approve response = %v", ap) + } +} diff --git a/internal/httpapi/server.go b/internal/httpapi/server.go index 19645d7..076953c 100644 --- a/internal/httpapi/server.go +++ b/internal/httpapi/server.go @@ -56,14 +56,20 @@ func (s *Server) Handler() http.Handler { if s.runs != nil { mux.HandleFunc("GET /v1/checkpoints/{run_id}", s.requireAuth(s.handleGetCheckpoint)) mux.HandleFunc("GET /v1/runs/{id}", s.requireAuth(s.handleGetRun)) + // Run lifecycle control (Surface 2 — the Python SDK drives these). + mux.HandleFunc("POST /v1/runs", s.requireAuth(s.handleCreateRun)) + mux.HandleFunc("POST /v1/runs/{id}/steps", s.requireAuth(s.handleBeginStep)) + mux.HandleFunc("POST /v1/runs/{id}/checkpoints", s.requireAuth(s.handleSaveCheckpoint)) + mux.HandleFunc("POST /v1/runs/{id}/cancel", s.requireAuth(s.handleCancelRun)) } if s.approvals != nil { mux.HandleFunc("POST /v1/runs/{id}/approve", s.requireAuth(s.handleApprove)) + mux.HandleFunc("POST /v1/runs/{id}/approvals", s.requireAuth(s.handleRequestApproval)) mux.HandleFunc("GET /v1/approvals", s.requireAuth(s.handleListApprovals)) + mux.HandleFunc("GET /v1/approvals/{id}", s.requireAuth(s.handleGetApproval)) // Local admin web page (Surface: human-in-the-loop, pull channel). mux.HandleFunc("GET /admin/approvals", s.requireAuth(s.handleAdminApprovalsPage)) } - // Further /v1 routes (full runs API) land in later build steps. return s.recoverer(mux) } diff --git a/sdks/python/README.md b/sdks/python/README.md new file mode 100644 index 0000000..59788bb --- /dev/null +++ b/sdks/python/README.md @@ -0,0 +1,81 @@ +# riskkernel (Python SDK) + +The Python SDK for [RiskKernel](https://github.com/prashar32/riskkernel) — **Surface 2**, deep control over a governed agent run. + +It is a **thin client** over the self-hosted RiskKernel daemon. Every deterministic +decision — budgets, loop/time halts, approval policy — happens in the Go core. The +SDK just makes governed runs ergonomic from Python. **Core install is stdlib-only** +(no third-party dependencies). + +```bash +pip install riskkernel +``` + +## Quickstart + +```python +import riskkernel as rk + +rt = rk.Runtime(base_url="http://localhost:7070") # your daemon + +with rt.governed_run(name="research", + budget=rt.budget(dollars=1.00, loops=20, seconds=300)) as run: + # Route your LLM client through the governing proxy so every model call is + # metered, priced, and budget-enforced under this run: + cfg = run.proxy_config() + # cfg["base_url"] -> http://localhost:7070/v1 + # cfg["headers"] -> {"X-RiskKernel-Run-Id": ""} + + for _ in range(100): + run.step() # raises rk.BudgetExceeded when loops/time run out + # ... your agent reasoning + tool calls ... + run.checkpoint("after-step", {"messages": messages}) +``` + +When the governor halts the run (token / dollar / loop / time budget), the next +`run.step()` — or a proxied model call — raises `rk.BudgetExceeded`. + +## Human-in-the-loop tools + +Gate side-effecting tools on human approval (the daemon's policy decides what needs +it; the call blocks until a human resolves it via CLI / web / webhook): + +```python +from riskkernel import governed_tool, ApprovalGate + +@governed_tool(side_effect="write") +def write_file(path, content): + ... # only runs if approved; else rk.ApprovalDenied + +# or explicitly: +gate = ApprovalGate(run) +if gate.allow("mcp://shell", side_effect="exec", arguments={"cmd": cmd}): + run_shell(cmd) +``` + +## Framework adapters + +Lazy-imported, so you only pay for what you use: + +```python +# LangChain / LangGraph — enforces loop/time budgets per LLM call +from riskkernel.adapters.langchain import RiskKernelCallbackHandler +llm.invoke(prompt, config={"callbacks": [RiskKernelCallbackHandler(run)]}) + +# Claude Agent SDK — PreToolUse approval hook +from riskkernel.adapters.claude_agent import make_pre_tool_use_hook +hook = make_pre_tool_use_hook(run, side_effect_for={"Bash": "exec", "Write": "write"}) + +# OpenAI Agents SDK — RunHooks (steps + tool approval) +from riskkernel.adapters.openai_agents import RiskKernelRunHooks +hooks = RiskKernelRunHooks(run, gate_tools=True) +``` + +## Configuration + +`Runtime(base_url=..., token=...)`, or the env vars `RISKKERNEL_BASE_URL` and +`RISKKERNEL_API_TOKEN` (used by the decorator/convenience API and `default_runtime()`). + +## License + +Apache-2.0. diff --git a/sdks/python/pyproject.toml b/sdks/python/pyproject.toml new file mode 100644 index 0000000..45bf4a8 --- /dev/null +++ b/sdks/python/pyproject.toml @@ -0,0 +1,32 @@ +[build-system] +requires = ["hatchling"] +build-backend = "hatchling.build" + +[project] +name = "riskkernel" +version = "0.1.0.dev0" +description = "Thin Python client for the RiskKernel reliability runtime (Surface 2)." +readme = "README.md" +requires-python = ">=3.9" +license = "Apache-2.0" +authors = [{ name = "Adarsh Prashar" }] +keywords = ["llm", "agents", "governance", "reliability", "budget", "guardrails"] +classifiers = [ + "Development Status :: 3 - Alpha", + "Intended Audience :: Developers", + "License :: OSI Approved :: Apache Software License", + "Programming Language :: Python :: 3", + "Topic :: Software Development :: Libraries", +] +dependencies = [] # core is stdlib-only + +[project.urls] +Homepage = "https://github.com/prashar32/riskkernel" +Source = "https://github.com/prashar32/riskkernel" + +[project.optional-dependencies] +langchain = ["langchain-core"] +dev = ["pytest"] + +[tool.hatch.build.targets.wheel] +packages = ["riskkernel"] diff --git a/sdks/python/riskkernel/__init__.py b/sdks/python/riskkernel/__init__.py new file mode 100644 index 0000000..6fad48f --- /dev/null +++ b/sdks/python/riskkernel/__init__.py @@ -0,0 +1,59 @@ +"""RiskKernel Python SDK — Surface 2 (deep control). + +A thin client over the self-hosted RiskKernel daemon. The Go core makes every +deterministic decision (budgets, halts, approval policy); this package just makes +governed runs ergonomic from Python. + +Quickstart:: + + import riskkernel as rk + + rt = rk.Runtime(base_url="http://localhost:7070") + with rt.governed_run(name="research", budget=rt.budget(dollars=1.00, loops=20)) as run: + cfg = run.proxy_config() # route your LLM client through the governing proxy + for _ in range(100): + run.step() # raises BudgetExceeded when loops/time run out + ... # your agent reasoning + tool calls + run.checkpoint("after-step", {"messages": messages}) +""" + +from .client import RiskKernel +from .errors import ( + APIError, + ApprovalDenied, + ApprovalTimeout, + BudgetExceeded, + RiskKernelError, +) +from .approval import ApprovalGate, governed_tool +from .runtime import ( + Budget, + Decision, + Run, + Runtime, + configure, + current_run, + default_runtime, + governed_run, +) + +__version__ = "0.1.0.dev0" + +__all__ = [ + "RiskKernel", + "Runtime", + "Run", + "Budget", + "Decision", + "ApprovalGate", + "governed_run", + "governed_tool", + "current_run", + "configure", + "default_runtime", + "RiskKernelError", + "APIError", + "BudgetExceeded", + "ApprovalDenied", + "ApprovalTimeout", +] diff --git a/sdks/python/riskkernel/adapters/__init__.py b/sdks/python/riskkernel/adapters/__init__.py new file mode 100644 index 0000000..39fc6e4 --- /dev/null +++ b/sdks/python/riskkernel/adapters/__init__.py @@ -0,0 +1,8 @@ +"""Framework adapters that bind a RiskKernel governed run to popular agent +frameworks. Each adapter lazily imports its framework, so the core SDK has no +third-party dependencies and you only pay for what you use. + +- ``langchain`` — a CallbackHandler (loop/time enforcement per LLM call). +- ``claude_agent`` — a PreToolUse hook for the Claude Agent SDK (approval gate). +- ``openai_agents`` — RunHooks for the OpenAI Agents SDK (steps + approval gate). +""" diff --git a/sdks/python/riskkernel/adapters/claude_agent.py b/sdks/python/riskkernel/adapters/claude_agent.py new file mode 100644 index 0000000..c115727 --- /dev/null +++ b/sdks/python/riskkernel/adapters/claude_agent.py @@ -0,0 +1,75 @@ +"""Claude Agent SDK adapter: a PreToolUse hook that routes side-effecting tool +calls through the RiskKernel approval gate. Maps cleanly to the Claude Agent SDK's +permission model (``permissionDecision: "deny"`` blocks a tool). + + from riskkernel.adapters.claude_agent import make_pre_tool_use_hook + hook = make_pre_tool_use_hook(run, side_effect_for={"Bash": "exec", "Write": "write"}) + # register `hook` as your PreToolUse hook in the Claude Agent SDK options. + +The hook signature follows the Claude Agent SDK: it receives the hook input +(containing the tool name and input) and returns a decision dict. Because SDK +versions differ slightly, the returned shape is the documented +``hookSpecificOutput`` / ``permissionDecision`` form; adjust if your version +differs. +""" + +from __future__ import annotations + +from typing import Any, Callable, Dict, Optional + +from ..runtime import Run + + +def make_pre_tool_use_hook( + run: Run, + side_effect_for: Optional[Dict[str, str]] = None, + default_side_effect: str = "write", + timeout: Optional[float] = None, +) -> Callable[..., dict]: + """Build a PreToolUse hook bound to a governed run. + + Args: + run: the governed Run. + side_effect_for: map of tool name -> side-effect label. Tools not listed + use ``default_side_effect``. A tool mapped to "" (empty) is treated as + read-only and never gated. + default_side_effect: side effect for unlisted tools. + timeout: max seconds to await a human decision. + """ + side_effect_for = side_effect_for or {} + + def hook(input_data: Any = None, *args: Any, **kwargs: Any) -> dict: + tool_name, tool_input = _extract(input_data, kwargs) + side_effect = side_effect_for.get(tool_name, default_side_effect) + decision = run.approve( + tool_name or "tool", side_effect=side_effect, + arguments={"input": _stringify(tool_input)}, timeout=timeout, + ) + if decision.approved: + return {} # allow (no decision == proceed) + return { + "hookSpecificOutput": { + "hookEventName": "PreToolUse", + "permissionDecision": "deny", + "permissionDecisionReason": decision.reason or "denied via RiskKernel approval gate", + } + } + + return hook + + +def _extract(input_data: Any, kwargs: dict): + """Pull tool name + input out of the hook payload across SDK shapes.""" + data = input_data if isinstance(input_data, dict) else kwargs + name = data.get("tool_name") or data.get("toolName") or data.get("name") or "" + tinput = data.get("tool_input") or data.get("toolInput") or data.get("input") + return name, tinput + + +def _stringify(v: Any) -> Any: + try: + import json + json.dumps(v) + return v + except Exception: + return repr(v) diff --git a/sdks/python/riskkernel/adapters/langchain.py b/sdks/python/riskkernel/adapters/langchain.py new file mode 100644 index 0000000..462c4be --- /dev/null +++ b/sdks/python/riskkernel/adapters/langchain.py @@ -0,0 +1,75 @@ +"""LangChain / LangGraph adapter: a callback handler that enforces a governed +run's loop and time budgets, ticking one step per LLM call. Point your LangChain +LLM at the governing proxy (``run.proxy_config()``) for token/cost/budget metering; +this handler adds the outer-loop enforcement the proxy can't see. + + from riskkernel.adapters.langchain import RiskKernelCallbackHandler + handler = RiskKernelCallbackHandler(run) + llm.invoke(prompt, config={"callbacks": [handler]}) + +A BudgetExceeded raised here propagates out of the LangChain call, halting the +chain — exactly when the run is out of budget. +""" + +from __future__ import annotations + +from typing import Any, Optional + +from ..approval import ApprovalGate +from ..runtime import Run + + +def _base_handler(): + # Inherit the real base class when LangChain is installed (best integration); + # otherwise fall back to object so the module still imports. + try: + from langchain_core.callbacks import BaseCallbackHandler # type: ignore + return BaseCallbackHandler + except Exception: + try: + from langchain.callbacks.base import BaseCallbackHandler # type: ignore + return BaseCallbackHandler + except Exception: + return object + + +class RiskKernelCallbackHandler(_base_handler()): # type: ignore[misc] + """Enforces loop/time budgets and (optionally) gates tools on approval. + + Args: + run: the governed Run. + gate_tools: if True, every tool call must pass the approval gate. + tool_side_effect: side-effect label reported for gated tools. + """ + + def __init__(self, run: Run, gate_tools: bool = False, + tool_side_effect: str = "tool"): + self.run = run + self.gate_tools = gate_tools + self.tool_side_effect = tool_side_effect + self._gate = ApprovalGate(run) + + # One LLM call == one governed step. Raises BudgetExceeded when spent. + def on_llm_start(self, serialized: Any, prompts: Any, **kwargs: Any) -> None: + self.run.step() + + def on_chat_model_start(self, serialized: Any, messages: Any, **kwargs: Any) -> None: + self.run.step() + + def on_tool_start(self, serialized: Any, input_str: Any, **kwargs: Any) -> None: + if not self.gate_tools: + return + name = "" + if isinstance(serialized, dict): + name = serialized.get("name", "") + self._gate.require(name or "tool", side_effect=self.tool_side_effect, + arguments={"input": _stringify(input_str)}) + + +def _stringify(v: Any) -> Any: + try: + import json + json.dumps(v) + return v + except Exception: + return repr(v) diff --git a/sdks/python/riskkernel/adapters/openai_agents.py b/sdks/python/riskkernel/adapters/openai_agents.py new file mode 100644 index 0000000..c93f1b8 --- /dev/null +++ b/sdks/python/riskkernel/adapters/openai_agents.py @@ -0,0 +1,49 @@ +"""OpenAI Agents SDK adapter: lifecycle hooks that tick a governed step per agent +turn and gate tools through the approval gate. + + from riskkernel.adapters.openai_agents import RiskKernelRunHooks + hooks = RiskKernelRunHooks(run, gate_tools=True) + await Runner.run(agent, input, hooks=hooks) + +The OpenAI Agents SDK calls ``on_agent_start``/``on_tool_start`` (async). We tick a +step on each agent start (loop/time enforcement) and, when ``gate_tools`` is set, +await approval before a tool runs — raising ApprovalDenied to block it. +""" + +from __future__ import annotations + +from typing import Any, Optional + +from ..approval import ApprovalGate +from ..runtime import Run + + +def _base_hooks(): + try: + from agents import RunHooks # type: ignore (openai-agents) + return RunHooks + except Exception: + return object + + +class RiskKernelRunHooks(_base_hooks()): # type: ignore[misc] + """RunHooks that bind a governed run to an OpenAI Agents run.""" + + def __init__(self, run: Run, gate_tools: bool = False, + tool_side_effect: str = "tool", timeout: Optional[float] = None): + self.run = run + self.gate_tools = gate_tools + self.tool_side_effect = tool_side_effect + self.timeout = timeout + self._gate = ApprovalGate(run) + + async def on_agent_start(self, context: Any = None, agent: Any = None, **kwargs: Any) -> None: + # One agent turn == one governed step (enforces loop/time budgets). + self.run.step() + + async def on_tool_start(self, context: Any = None, agent: Any = None, + tool: Any = None, **kwargs: Any) -> None: + if not self.gate_tools: + return + name = getattr(tool, "name", None) or str(tool) + self._gate.require(name, side_effect=self.tool_side_effect, timeout=self.timeout) diff --git a/sdks/python/riskkernel/approval.py b/sdks/python/riskkernel/approval.py new file mode 100644 index 0000000..39e9bef --- /dev/null +++ b/sdks/python/riskkernel/approval.py @@ -0,0 +1,86 @@ +"""Human-in-the-loop approval helpers for the SDK.""" + +from __future__ import annotations + +import functools +from typing import Any, Callable, Optional + +from .errors import ApprovalDenied, RiskKernelError +from .runtime import Decision, Run, current_run + + +class ApprovalGate: + """Gates side-effecting actions on human approval. Wraps a Run and asks the + daemon (deterministic policy) whether a call needs approval, then blocks until + a human resolves it. + + gate = ApprovalGate(run) + if gate.allow("mcp://shell", side_effect="exec", arguments={"cmd": cmd}): + run_shell(cmd) + """ + + def __init__(self, run: Optional[Run] = None): + self._run = run + + def _resolve_run(self) -> Run: + run = self._run or current_run() + if run is None: + raise RiskKernelError("ApprovalGate used outside a governed run") + return run + + def decide(self, tool: str, side_effect: str = "", arguments: Optional[dict] = None, + step_index: int = 0, timeout: Optional[float] = None) -> Decision: + """Return the Decision (blocking until resolved). Does not raise on denial.""" + return self._resolve_run().approve( + tool, side_effect=side_effect, arguments=arguments, + step_index=step_index, timeout=timeout, + ) + + def allow(self, tool: str, side_effect: str = "", arguments: Optional[dict] = None, + step_index: int = 0, timeout: Optional[float] = None) -> bool: + """Convenience boolean: True if approved.""" + return self.decide(tool, side_effect, arguments, step_index, timeout).approved + + def require(self, tool: str, side_effect: str = "", arguments: Optional[dict] = None, + step_index: int = 0, timeout: Optional[float] = None) -> None: + """Raise ApprovalDenied if not approved (use to guard before a side effect).""" + d = self.decide(tool, side_effect, arguments, step_index, timeout) + if not d.approved: + raise ApprovalDenied(tool, d.reason) + + +def governed_tool(_fn: Optional[Callable] = None, *, tool: Optional[str] = None, + side_effect: str = "write", timeout: Optional[float] = None): + """Decorator for a side-effecting tool function: before it runs, ask the + approval gate (under the current governed run). Raises ApprovalDenied if the + human says no. + + @governed_tool(side_effect="write") + def write_file(path, content): ... + """ + + def decorate(fn: Callable) -> Callable: + tool_name = tool or fn.__name__ + + @functools.wraps(fn) + def wrapper(*args: Any, **kwargs: Any) -> Any: + ApprovalGate().require( + tool_name, side_effect=side_effect, + arguments={"args": _safe(args), "kwargs": _safe(kwargs)}, + timeout=timeout, + ) + return fn(*args, **kwargs) + + return wrapper + + return decorate(_fn) if _fn is not None else decorate + + +def _safe(obj: Any) -> Any: + """Best-effort JSON-able rendering of call arguments for the approver to read.""" + try: + import json + json.dumps(obj) + return obj + except Exception: + return repr(obj) diff --git a/sdks/python/riskkernel/client.py b/sdks/python/riskkernel/client.py new file mode 100644 index 0000000..ed96cd0 --- /dev/null +++ b/sdks/python/riskkernel/client.py @@ -0,0 +1,124 @@ +"""Thin HTTP client for the RiskKernel daemon's /v1 API. + +Stdlib only (urllib) — no third-party dependencies, so ``pip install riskkernel`` +stays light and auditable. This client carries NO governance logic: the Go daemon +makes every deterministic decision. The client just relays calls and surfaces the +daemon's verdicts (e.g. a 402 becomes ``BudgetExceeded``). +""" + +from __future__ import annotations + +import json +import urllib.error +import urllib.request +from typing import Any, Optional + +from .errors import APIError, BudgetExceeded + + +class RiskKernel: + """Client for a running RiskKernel daemon.""" + + def __init__( + self, + base_url: str = "http://localhost:7070", + token: Optional[str] = None, + timeout: float = 30.0, + ): + self.base_url = base_url.rstrip("/") + self.token = token + self.timeout = timeout + + # --- low-level --- + + def _request(self, method: str, path: str, body: Optional[dict] = None) -> Any: + url = self.base_url + path + data = None + headers = {"Accept": "application/json"} + if body is not None: + data = json.dumps(body).encode("utf-8") + headers["Content-Type"] = "application/json" + if self.token: + headers["Authorization"] = "Bearer " + self.token + + req = urllib.request.Request(url, data=data, headers=headers, method=method) + try: + with urllib.request.urlopen(req, timeout=self.timeout) as resp: + raw = resp.read() + return json.loads(raw) if raw else {} + except urllib.error.HTTPError as e: + raw = e.read() + payload = {} + try: + payload = json.loads(raw) + except Exception: + pass + code = payload.get("code", "") + message = payload.get("message", e.reason or "") + # 402 == the governor halted the run on a budget. + if e.code == 402: + raise BudgetExceeded(code, message) from None + raise APIError(e.code, code, message) from None + except urllib.error.URLError as e: + raise APIError(0, "connection_error", + f"cannot reach daemon at {self.base_url}: {e.reason}") from None + + # --- runs --- + + def create_run(self, name: Optional[str] = None, budget: Optional[dict] = None, + metadata: Optional[dict] = None) -> dict: + body: dict = {} + if name: + body["name"] = name + if budget: + body["budget"] = budget + if metadata: + body["metadata"] = metadata + return self._request("POST", "/v1/runs", body) + + def get_run(self, run_id: str) -> dict: + return self._request("GET", f"/v1/runs/{run_id}") + + def begin_step(self, run_id: str) -> int: + """Register a loop iteration. Raises BudgetExceeded (402) if the loop or + time budget is spent.""" + out = self._request("POST", f"/v1/runs/{run_id}/steps", {}) + return int(out.get("stepIndex", 0)) + + def checkpoint(self, run_id: str, name: str = "", payload: Optional[dict] = None) -> None: + self._request("POST", f"/v1/runs/{run_id}/checkpoints", + {"name": name, "payload": payload or {}}) + + def latest_checkpoint(self, run_id: str) -> Optional[dict]: + try: + return self._request("GET", f"/v1/checkpoints/{run_id}") + except APIError as e: + if e.status == 404: + return None + raise + + def cancel(self, run_id: str, reason: str = "") -> dict: + return self._request("POST", f"/v1/runs/{run_id}/cancel", {"reason": reason}) + + # --- approvals --- + + def request_approval(self, run_id: str, tool: str, side_effect: str = "", + arguments: Optional[dict] = None, step_index: int = 0) -> dict: + """Request approval for a (possibly side-effecting) tool call. Returns a + dict with ``status``: ``approved`` (allowed by policy) or ``pending`` + (a human must decide; poll ``get_approval``).""" + return self._request("POST", f"/v1/runs/{run_id}/approvals", { + "tool": tool, "sideEffect": side_effect, + "arguments": arguments or {}, "stepIndex": step_index, + }) + + def get_approval(self, approval_id: str) -> dict: + return self._request("GET", f"/v1/approvals/{approval_id}") + + def resolve_approval(self, run_id: str, approval_id: str, approve: bool, + reason: str = "", decided_by: str = "sdk") -> dict: + return self._request("POST", f"/v1/runs/{run_id}/approve", { + "approvalId": approval_id, + "decision": "approve" if approve else "deny", + "reason": reason, "decidedBy": decided_by, + }) diff --git a/sdks/python/riskkernel/errors.py b/sdks/python/riskkernel/errors.py new file mode 100644 index 0000000..bdd8727 --- /dev/null +++ b/sdks/python/riskkernel/errors.py @@ -0,0 +1,40 @@ +"""Exceptions raised by the RiskKernel SDK.""" + +from __future__ import annotations + + +class RiskKernelError(Exception): + """Base class for all SDK errors.""" + + +class APIError(RiskKernelError): + """The daemon returned an unexpected (non-2xx) response.""" + + def __init__(self, status: int, code: str = "", message: str = ""): + self.status = status + self.code = code + self.message = message + super().__init__(f"riskkernel API error {status} {code}: {message}") + + +class BudgetExceeded(RiskKernelError): + """A governed run hit one of its hard budgets (the deterministic governor + halted it). ``reason`` is the machine-readable HaltReason, e.g. + ``token_budget_exceeded`` or ``loop_budget_exceeded``.""" + + def __init__(self, reason: str, message: str = ""): + self.reason = reason + super().__init__(message or f"run halted: {reason}") + + +class ApprovalDenied(RiskKernelError): + """A human denied a side-effecting tool call gated by the approval gate.""" + + def __init__(self, tool: str, reason: str = ""): + self.tool = tool + self.reason = reason + super().__init__(f"approval denied for {tool}" + (f": {reason}" if reason else "")) + + +class ApprovalTimeout(RiskKernelError): + """No human resolved a pending approval within the configured timeout.""" diff --git a/sdks/python/riskkernel/runtime.py b/sdks/python/riskkernel/runtime.py new file mode 100644 index 0000000..d8dd856 --- /dev/null +++ b/sdks/python/riskkernel/runtime.py @@ -0,0 +1,199 @@ +"""High-level governed-run ergonomics over the thin client. + +The runtime is still thin: budgets, loop/time enforcement, checkpoints, and +approval decisions all happen in the Go daemon. This module just makes them +pleasant to use from Python (context managers, decorators, a current-run var). +""" + +from __future__ import annotations + +import contextvars +import functools +import os +import time +from contextlib import contextmanager +from dataclasses import dataclass +from typing import Any, Callable, Optional + +from .client import RiskKernel +from .errors import ApprovalTimeout + +# The run currently in scope (set by governed_run), so @governed_tool and +# checkpoint() can find it without threading it through every call. +_current_run: contextvars.ContextVar[Optional["Run"]] = contextvars.ContextVar( + "riskkernel_current_run", default=None +) + + +def current_run() -> Optional["Run"]: + """Return the governed run currently in scope, or None.""" + return _current_run.get() + + +@dataclass +class Budget: + """Hard per-run limits. Any field left None is unlimited for that dimension.""" + + tokens: Optional[int] = None + dollars: Optional[float] = None + loops: Optional[int] = None + seconds: Optional[int] = None + + def to_dict(self) -> dict: + out: dict = {} + if self.tokens is not None: + out["tokens"] = self.tokens + if self.dollars is not None: + out["dollars"] = self.dollars + if self.loops is not None: + out["loops"] = self.loops + if self.seconds is not None: + out["seconds"] = self.seconds + return out + + +@dataclass +class Decision: + approved: bool + required: bool = True + reason: str = "" + by: str = "" + + +class Run: + """A governed run bound to a daemon run id.""" + + def __init__(self, client: RiskKernel, data: dict, + poll_interval: float = 2.0, timeout: Optional[float] = None): + self._client = client + self._poll = poll_interval + self._timeout = timeout + self.id: str = data["id"] + self.data = data + + def step(self) -> int: + """Register a loop iteration; raises BudgetExceeded if loop/time budget is spent.""" + return self._client.begin_step(self.id) + + def checkpoint(self, name: str = "", payload: Optional[dict] = None) -> None: + self._client.checkpoint(self.id, name, payload) + + def latest_checkpoint(self) -> Optional[dict]: + return self._client.latest_checkpoint(self.id) + + def cancel(self, reason: str = "") -> dict: + return self._client.cancel(self.id, reason) + + def status(self) -> dict: + return self._client.get_run(self.id) + + def proxy_config(self) -> dict: + """Config for routing this run's model calls through the governing proxy. + Point your LLM client's base URL here and send the header so every call is + metered, priced, and budget-enforced under this run.""" + return { + "base_url": self._client.base_url + "/v1", + "headers": {"X-RiskKernel-Run-Id": self.id}, + } + + def approve(self, tool: str, side_effect: str = "", arguments: Optional[dict] = None, + step_index: int = 0, poll_interval: Optional[float] = None, + timeout: Optional[float] = None) -> Decision: + """Request approval for a tool call, blocking (polling) until a human + resolves it. Returns a Decision; raises ApprovalTimeout if none arrives.""" + res = self._client.request_approval(self.id, tool, side_effect, arguments, step_index) + if res.get("status") == "approved": + return Decision(True, bool(res.get("required", True))) + approval_id = res["id"] + interval = poll_interval if poll_interval is not None else self._poll + limit = timeout if timeout is not None else self._timeout + deadline = (time.monotonic() + limit) if limit is not None else None + while True: + a = self._client.get_approval(approval_id) + st = a.get("status") + if st == "approved": + return Decision(True, True, a.get("reason", ""), a.get("decidedBy", "")) + if st == "denied": + return Decision(False, True, a.get("reason", ""), a.get("decidedBy", "")) + if deadline is not None and time.monotonic() > deadline: + raise ApprovalTimeout(f"no decision for approval {approval_id} within {limit}s") + time.sleep(interval) + + +class Runtime: + """Entry point: holds a client and default approval-polling settings.""" + + def __init__(self, client: Optional[RiskKernel] = None, + base_url: str = "http://localhost:7070", token: Optional[str] = None, + approval_poll_interval: float = 2.0, + approval_timeout: Optional[float] = None): + self.client = client or RiskKernel(base_url, token) + self._poll = approval_poll_interval + self._timeout = approval_timeout + + def budget(self, tokens: Optional[int] = None, dollars: Optional[float] = None, + loops: Optional[int] = None, seconds: Optional[int] = None) -> Budget: + return Budget(tokens, dollars, loops, seconds) + + @contextmanager + def governed_run(self, name: Optional[str] = None, + budget: Optional[Budget | dict] = None, + metadata: Optional[dict] = None, cancel_on_error: bool = True): + """Context manager that opens a governed run, sets it as the current run, + and cancels it if the body raises (unless cancel_on_error=False).""" + b = budget.to_dict() if isinstance(budget, Budget) else budget + data = self.client.create_run(name=name, budget=b, metadata=metadata) + run = Run(self.client, data, self._poll, self._timeout) + token = _current_run.set(run) + try: + yield run + except Exception: + if cancel_on_error: + try: + run.cancel("error") + except Exception: + pass + raise + finally: + _current_run.reset(token) + + +# Module-level default runtime, configured from the environment, for the +# decorator/convenience API. +_default_runtime: Optional[Runtime] = None + + +def default_runtime() -> Runtime: + global _default_runtime + if _default_runtime is None: + _default_runtime = Runtime( + base_url=os.environ.get("RISKKERNEL_BASE_URL", "http://localhost:7070"), + token=os.environ.get("RISKKERNEL_API_TOKEN"), + ) + return _default_runtime + + +def configure(runtime: Runtime) -> None: + """Override the module-level default runtime (used by the decorators).""" + global _default_runtime + _default_runtime = runtime + + +def governed_run(_fn: Optional[Callable] = None, *, name: Optional[str] = None, + budget: Optional[Budget | dict] = None, runtime: Optional[Runtime] = None): + """Decorator: run the wrapped function inside a governed run. The run is + available via current_run() (and passed as the ``run`` kwarg if the function + declares one).""" + + def decorate(fn: Callable) -> Callable: + @functools.wraps(fn) + def wrapper(*args: Any, **kwargs: Any) -> Any: + rt = runtime or default_runtime() + run_name = name or fn.__name__ + with rt.governed_run(name=run_name, budget=budget) as run: + if "run" in fn.__code__.co_varnames and "run" not in kwargs: + kwargs["run"] = run + return fn(*args, **kwargs) + return wrapper + + return decorate(_fn) if _fn is not None else decorate diff --git a/sdks/python/tests/__init__.py b/sdks/python/tests/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/sdks/python/tests/test_sdk.py b/sdks/python/tests/test_sdk.py new file mode 100644 index 0000000..6d8c843 --- /dev/null +++ b/sdks/python/tests/test_sdk.py @@ -0,0 +1,152 @@ +"""SDK tests against a stdlib stub daemon — no Go binary, no third-party deps.""" + +import json +import threading +import unittest +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer + +import riskkernel as rk +from riskkernel.errors import ApprovalDenied, BudgetExceeded + + +class _State: + def __init__(self): + self.steps = 0 + self.loop_budget = 2 + self.approval_polls = 0 + self.last_checkpoint = None + + +STATE = _State() + + +class StubHandler(BaseHTTPRequestHandler): + def log_message(self, *a): # silence + pass + + def _send(self, code, obj): + body = json.dumps(obj).encode() + self.send_response(code) + self.send_header("Content-Type", "application/json") + self.send_header("Content-Length", str(len(body))) + self.end_headers() + self.wfile.write(body) + + def _read(self): + n = int(self.headers.get("Content-Length", 0) or 0) + raw = self.rfile.read(n) if n else b"" + return json.loads(raw) if raw else {} + + def do_GET(self): + p = self.path + if p.startswith("/v1/approvals/"): + STATE.approval_polls += 1 + status = "approved" if STATE.approval_polls >= 2 else "pending" + return self._send(200, {"id": "ap-1", "status": status, "decidedBy": "tester"}) + if p.startswith("/v1/checkpoints/"): + return self._send(200, {"runId": "run-1", "stepIndex": 1, + "payload": STATE.last_checkpoint or {}}) + return self._send(404, {"code": "not_found", "message": "no"}) + + def do_POST(self): + p = self.path + body = self._read() + if p == "/v1/runs": + return self._send(201, {"id": "run-1", "status": "running", + "budget": body.get("budget", {}), + "usage": {"tokens": 0, "loops": 0}}) + if p == "/v1/runs/run-1/steps": + STATE.steps += 1 + if STATE.steps > STATE.loop_budget: + return self._send(402, {"code": "loop_budget_exceeded", + "message": "run halted: loop_budget_exceeded"}) + return self._send(200, {"stepIndex": STATE.steps}) + if p == "/v1/runs/run-1/checkpoints": + STATE.last_checkpoint = body.get("payload") + return self._send(201, {"ok": True}) + if p == "/v1/runs/run-1/cancel": + return self._send(200, {"id": "run-1", "status": "cancelled"}) + if p == "/v1/runs/run-1/approvals": + if not body.get("sideEffect"): + return self._send(200, {"status": "approved", "required": False}) + return self._send(201, {"id": "ap-1", "status": "pending"}) + if p == "/v1/runs/run-1/approve": + return self._send(200, {"id": "run-1", "status": "running"}) + return self._send(404, {"code": "not_found", "message": "no"}) + + +class SDKTest(unittest.TestCase): + @classmethod + def setUpClass(cls): + cls.server = ThreadingHTTPServer(("127.0.0.1", 0), StubHandler) + cls.port = cls.server.server_address[1] + cls.thread = threading.Thread(target=cls.server.serve_forever, daemon=True) + cls.thread.start() + cls.base = f"http://127.0.0.1:{cls.port}" + + @classmethod + def tearDownClass(cls): + cls.server.shutdown() + + def setUp(self): + global STATE + STATE = _State() # the handler reads this module global + self.rt = rk.Runtime(base_url=self.base, approval_poll_interval=0.01) + + def test_governed_run_and_step_budget(self): + with self.rt.governed_run(name="t", budget=self.rt.budget(loops=2)) as run: + self.assertEqual(run.id, "run-1") + self.assertEqual(run.step(), 1) + self.assertEqual(run.step(), 2) + with self.assertRaises(BudgetExceeded) as cm: + run.step() + self.assertEqual(cm.exception.reason, "loop_budget_exceeded") + + def test_checkpoint_roundtrip(self): + with self.rt.governed_run(name="t") as run: + run.checkpoint("after", {"cursor": 7}) + cp = run.latest_checkpoint() + self.assertEqual(cp["payload"]["cursor"], 7) + + def test_budget_to_dict(self): + b = self.rt.budget(tokens=100, dollars=1.5) + self.assertEqual(b.to_dict(), {"tokens": 100, "dollars": 1.5}) + + def test_approval_not_required(self): + with self.rt.governed_run(name="t") as run: + d = run.approve("mcp://fs", side_effect="") # read-only + self.assertTrue(d.approved) + self.assertFalse(d.required) + + def test_approval_pending_then_approved(self): + with self.rt.governed_run(name="t") as run: + d = run.approve("mcp://shell", side_effect="exec", arguments={"cmd": "ls"}) + self.assertTrue(d.approved) # stub flips to approved on 2nd poll + + def test_governed_tool_denied(self): + # Make the stub deny: flip approval to "denied" by overriding poll result. + with self.rt.governed_run(name="t") as run: + # Monkeypatch the client to return a denied approval. + run._client.get_approval = lambda _id: {"id": "ap-1", "status": "denied", "reason": "no"} + + @rk.governed_tool(side_effect="write", tool="danger") + def danger(): + return "ran" + + with self.assertRaises(ApprovalDenied): + danger() + + def test_cancel(self): + with self.rt.governed_run(name="t") as run: + out = run.cancel("done") + self.assertEqual(out["status"], "cancelled") + + def test_proxy_config(self): + with self.rt.governed_run(name="t") as run: + cfg = run.proxy_config() + self.assertTrue(cfg["base_url"].endswith("/v1")) + self.assertEqual(cfg["headers"]["X-RiskKernel-Run-Id"], "run-1") + + +if __name__ == "__main__": + unittest.main()