From d4d5997654a06f505d2cbbf5bf798332fb4433d9 Mon Sep 17 00:00:00 2001 From: matt Date: Tue, 23 Jun 2026 15:58:49 -0700 Subject: [PATCH 1/7] Tell agents to read in their rules to reduce the number of denied calls --- internal/api/auth_test.go | 11 ++- internal/api/handlers.go | 96 ++++++++++++++++----- internal/api/handlers_test.go | 157 ++++++++++++++++++++++++++++++++++ 3 files changed, 240 insertions(+), 24 deletions(-) diff --git a/internal/api/auth_test.go b/internal/api/auth_test.go index fd7f973..1012dcb 100644 --- a/internal/api/auth_test.go +++ b/internal/api/auth_test.go @@ -185,10 +185,14 @@ func TestMCPAcceptsValidTokenAndPlumbsAgentID(t *testing.T) { func TestAgentRulesRequiresAuthAndUsesTokenAgentID(t *testing.T) { rig := newAuthTestRig(t) + tokenAgent := store.AgentRecord{ID: "agent-cuid-007", AgentIDs: `["agent-007"]`} + otherAgent := store.AgentRecord{ID: "agent-cuid-other", AgentIDs: `["other"]`} rules := &stubRulesRepo{rules: []store.Rule{ - {ID: "auto-rule", Action: invocation.RuleActionAutoApprove, ServerPatterns: []string{"amp"}, ToolPatterns: []string{"Read"}, Enabled: true, Order: 0}, + {ID: "other-deny", Action: invocation.RuleActionAutoDeny, ServerPatterns: []string{"amp"}, ToolPatterns: []string{"Read"}, AgentCUIDs: []string{otherAgent.ID}, Enabled: true, Order: 0}, + {ID: "auto-rule", Action: invocation.RuleActionAutoApprove, ServerPatterns: []string{"amp"}, ToolPatterns: []string{"Read"}, AgentCUIDs: []string{tokenAgent.ID}, Enabled: true, Order: 1}, }} - h := NewHandler(&stubService{}, stubServerService{}, nil, rules, nil, nil, nil, nil, nil, nil) + agents := &stubAgentsRepo{records: []store.AgentRecord{tokenAgent, otherAgent}} + h := NewHandler(&stubService{}, stubServerService{}, nil, rules, agents, nil, nil, nil, nil, nil) h.SetAuthValidator(rig.v) handler := h.Routes() @@ -217,6 +221,9 @@ func TestAgentRulesRequiresAuthAndUsesTokenAgentID(t *testing.T) { if resp.Action != invocation.RuleActionAutoApprove { t.Fatalf("expected auto_approve action, got %q", resp.Action) } + if len(resp.Items) != 1 || resp.Items[0].ID != "auto-rule" { + t.Fatalf("expected only token agent rule to be visible, got %#v", resp.Items) + } } func TestAgentRuntimeEndpointsRequireAuthWhenValidatorConfigured(t *testing.T) { diff --git a/internal/api/handlers.go b/internal/api/handlers.go index 6835267..d01509f 100644 --- a/internal/api/handlers.go +++ b/internal/api/handlers.go @@ -616,7 +616,10 @@ func parseAgentIDs(raw string) []string { // MCP clients can surface it to the model as system-level guidance. const atryumInitializeInstructions = "This MCP server is gated by the Atryum harness. " + "Atryum may approve, deny, or request human approval for tool calls according to configured rules. " + - "Those rules may change between conversation turns, so treat each tool call as subject to the current Atryum policy." + "Those rules may change between conversation turns, so treat each tool call as subject to the current Atryum policy. " + + "To see the static approval rules that currently apply to you, issue an HTTP GET to /api/v1/agent/rules " + + "(optionally with ?server={server}&tool={tool} to preview the disposition for a specific tool); " + + "the response is advisory only, as AI-evaluation and human-approval outcomes are decided during the actual gated call." type AgentRule struct { ID string `json:"id"` @@ -624,17 +627,21 @@ type AgentRule struct { ServerPatterns []string `json:"server_patterns"` ToolPatterns []string `json:"tool_patterns"` Description string `json:"description,omitempty"` + Guidance string `json:"guidance,omitempty"` Order int `json:"order"` } type AgentRulesResponse struct { - AgentID string `json:"agent_id,omitempty"` - Server string `json:"server,omitempty"` - Tool string `json:"tool,omitempty"` - DefaultAction string `json:"default_action"` - Action string `json:"action,omitempty"` - MatchedRuleID *string `json:"matched_rule_id,omitempty"` - Items []AgentRule `json:"items"` + AgentID string `json:"agent_id,omitempty"` + Server string `json:"server,omitempty"` + Tool string `json:"tool,omitempty"` + DefaultAction string `json:"default_action"` + Action string `json:"action,omitempty"` + MatchedRuleID *string `json:"matched_rule_id,omitempty"` + GeneratedAt time.Time `json:"generated_at"` + EvaluationMode string `json:"evaluation_mode"` + Explanation string `json:"explanation"` + Items []AgentRule `json:"items"` } // ───────────────────────────────────────────────────────────────────────────── @@ -999,7 +1006,7 @@ func (h *Handler) agentRules(w http.ResponseWriter, r *http.Request) { return } if h.rulesRepo == nil { - writeJSON(w, http.StatusOK, AgentRulesResponse{DefaultAction: invocation.RuleActionHumanApproval, Items: []AgentRule{}}) + writeJSON(w, http.StatusOK, newAgentRulesResponse("", "", "")) return } @@ -1025,13 +1032,7 @@ func (h *Handler) agentRules(w http.ResponseWriter, r *http.Request) { } func (h *Handler) buildAgentRulesResponse(ctx context.Context, agentID, server, tool string) (AgentRulesResponse, error) { - resp := AgentRulesResponse{ - AgentID: agentID, - Server: server, - Tool: tool, - DefaultAction: invocation.RuleActionHumanApproval, - Items: []AgentRule{}, - } + resp := newAgentRulesResponse(agentID, server, tool) if h.rulesRepo == nil { return resp, nil } @@ -1043,7 +1044,7 @@ func (h *Handler) buildAgentRulesResponse(ctx context.Context, agentID, server, agentCUID := h.resolveAgentRecordForRules(ctx, agentID) for _, rule := range rules { - if !apiRuleMatches(rule, server, tool, agentCUID, false) { + if !apiRuleVisibleToAgent(rule, agentCUID) { continue } resp.Items = append(resp.Items, AgentRule{ @@ -1052,6 +1053,7 @@ func (h *Handler) buildAgentRulesResponse(ctx context.Context, agentID, server, ServerPatterns: rule.ServerPatterns, ToolPatterns: rule.ToolPatterns, Description: rule.Description, + Guidance: agentRuleGuidance(rule.Action), Order: rule.Order, }) if resp.Action == "" && server != "" && tool != "" && @@ -1066,6 +1068,20 @@ func (h *Handler) buildAgentRulesResponse(ctx context.Context, agentID, server, return resp, nil } +func newAgentRulesResponse(agentID, server, tool string) AgentRulesResponse { + return AgentRulesResponse{ + AgentID: agentID, + Server: server, + Tool: tool, + DefaultAction: invocation.RuleActionHumanApproval, + GeneratedAt: time.Now().UTC(), + EvaluationMode: "static_only", + Explanation: "This advisory response is based on the current static approval rules visible to the resolved Atryum agent. " + + "AI evaluation rules are decided only during the actual gated tool call.", + Items: []AgentRule{}, + } +} + func (h *Handler) resolveAgentRecordForRules(ctx context.Context, agentID string) string { if h.agentsRepo == nil { return "" @@ -1093,10 +1109,17 @@ func (h *Handler) resolveAgentRecordForRules(ctx context.Context, agentID string return rec.ID } -func apiRuleMatches(rule store.Rule, server, tool, agentCUID string, includeServerTool bool) bool { +func apiRuleVisibleToAgent(rule store.Rule, agentCUID string) bool { if !rule.Enabled { return false } + return apiMatchAgentCUIDs(rule.AgentCUIDs, agentCUID) +} + +func apiRuleMatches(rule store.Rule, server, tool, agentCUID string, includeServerTool bool) bool { + if !apiRuleVisibleToAgent(rule, agentCUID) { + return false + } if includeServerTool { if !apiMatchPatterns(rule.ServerPatterns, server) { return false @@ -1105,7 +1128,7 @@ func apiRuleMatches(rule store.Rule, server, tool, agentCUID string, includeServ return false } } - return apiMatchAgentCUIDs(rule.AgentCUIDs, agentCUID) + return true } func apiMatchAgentCUIDs(cuids []string, agentCUID string) bool { @@ -1221,7 +1244,11 @@ func (h *Handler) handleMCPProxy(w http.ResponseWriter, r *http.Request, server if len(resp.Error) > 0 { result := normalizeToolCallResult(resp.Error, true) if resp.Status == invocation.StatusDenied { - result = h.appendRulesContextToToolResult(r.Context(), result, server, params.Name) + agentID := auth.AgentIDFromContext(r.Context()) + if agentID == "" && h.authValidator == nil { + agentID = normalizeNoAuthRequestAgentID(requestID) + } + result = h.appendRulesContextToToolResult(r.Context(), result, agentID, server, params.Name) } h.writeRPCResult(w, req.ID, result) return @@ -1440,6 +1467,21 @@ func apiMatchPatterns(patterns []string, value string) bool { return false } +func agentRuleGuidance(action string) string { + switch action { + case invocation.RuleActionAutoApprove: + return "Likely to pass if policy and rule inputs do not change." + case invocation.RuleActionAutoDeny: + return "Will be rejected by this matched rule." + case invocation.RuleActionHumanApproval: + return "Requires reviewer approval." + case invocation.RuleActionAIEvaluation: + return "Evaluated only during the real gated call; advisory cannot guarantee the outcome." + default: + return "" + } +} + // annotatedTool is the on-the-wire shape used for tools/list when atryum is // able to compute a per-tool policy disposition. It mirrors mcp.Tool but adds // an `annotations` block plus a description prefix so models that ignore @@ -1519,11 +1561,11 @@ func effectiveActionForTool(rules []store.Rule, server, tool, agentCUID string) // appendRulesContextToToolResult adds an extra text content block to a denied // tool call result describing the applicable rules and effective action. -func (h *Handler) appendRulesContextToToolResult(ctx context.Context, result any, server, tool string) any { +func (h *Handler) appendRulesContextToToolResult(ctx context.Context, result any, agentID, server, tool string) any { if h.rulesRepo == nil { return result } - rulesResp, err := h.buildAgentRulesResponse(ctx, auth.AgentIDFromContext(ctx), server, tool) + rulesResp, err := h.buildAgentRulesResponse(ctx, agentID, server, tool) if err != nil { return result } @@ -3922,6 +3964,16 @@ func compactRequestID(id json.RawMessage) string { return string(id) } +func normalizeNoAuthRequestAgentID(value string) string { + value = strings.TrimSpace(value) + if strings.HasPrefix(value, `"`) { + if unquoted, err := strconv.Unquote(value); err == nil { + value = unquoted + } + } + return normalizeNoAuthAgentID(value) +} + func stringPtr(v string) *string { return &v } func invocationSignature(items []invocation.InvocationResponse) string { diff --git a/internal/api/handlers_test.go b/internal/api/handlers_test.go index 4718736..e7d498a 100644 --- a/internal/api/handlers_test.go +++ b/internal/api/handlers_test.go @@ -1173,12 +1173,62 @@ func TestAgentRulesListsApplicableRulesAndDisposition(t *testing.T) { if resp.MatchedRuleID == nil || *resp.MatchedRuleID != "read-auto" { t.Fatalf("expected matched rule read-auto, got %#v", resp.MatchedRuleID) } + if resp.GeneratedAt.IsZero() { + t.Fatal("expected generated_at to be populated") + } + if resp.EvaluationMode != "static_only" { + t.Fatalf("expected static_only evaluation mode, got %q", resp.EvaluationMode) + } + if !strings.Contains(resp.Explanation, "advisory") { + t.Fatalf("expected advisory explanation, got %q", resp.Explanation) + } if len(resp.Items) != 3 { t.Fatalf("expected three applicable rules, got %#v", resp.Items) } if resp.Items[0].ID != "bash-deny" || resp.Items[1].ID != "read-auto" || resp.Items[2].ID != "fallback-human" { t.Fatalf("unexpected applicable rules order: %#v", resp.Items) } + if resp.Items[1].Guidance == "" { + t.Fatalf("expected rule guidance, got %#v", resp.Items[1]) + } +} + +func TestAgentRulesFiltersOutRulesScopedToOtherAgents(t *testing.T) { + agent := store.AgentRecord{ID: "agent-cuid-007", AgentIDs: `["agent-007"]`} + other := store.AgentRecord{ID: "agent-cuid-other", AgentIDs: `["other-agent"]`} + rules := &stubRulesRepo{rules: []store.Rule{ + {ID: "other-agent", Action: invocation.RuleActionAutoDeny, ServerPatterns: []string{"amp"}, ToolPatterns: []string{"Read"}, AgentCUIDs: []string{other.ID}, Enabled: true, Order: 0}, + {ID: "this-agent", Action: invocation.RuleActionAutoApprove, ServerPatterns: []string{"amp"}, ToolPatterns: []string{"Read"}, AgentCUIDs: []string{agent.ID}, Enabled: true, Order: 1}, + {ID: "unscoped", Action: invocation.RuleActionHumanApproval, ServerPatterns: []string{"*"}, ToolPatterns: []string{"*"}, Enabled: true, Order: 2}, + }} + agents := &stubAgentsRepo{records: []store.AgentRecord{agent, other}} + h := NewHandler(&stubService{}, stubServerService{}, nil, rules, agents, nil, nil, nil, nil, nil) + req := httptest.NewRequest(http.MethodGet, "/api/v1/agent/rules?agent_id=agent-007&source=amp&tool=Read", nil) + w := httptest.NewRecorder() + + h.Routes().ServeHTTP(w, req) + + if w.Code != http.StatusOK { + t.Fatalf("expected 200, got %d body=%s", w.Code, w.Body.String()) + } + var resp AgentRulesResponse + if err := json.Unmarshal(w.Body.Bytes(), &resp); err != nil { + t.Fatal(err) + } + if resp.Action != invocation.RuleActionAutoApprove { + t.Fatalf("expected this-agent auto approve disposition, got %q", resp.Action) + } + if resp.MatchedRuleID == nil || *resp.MatchedRuleID != "this-agent" { + t.Fatalf("expected matched rule this-agent, got %#v", resp.MatchedRuleID) + } + if len(resp.Items) != 2 { + t.Fatalf("expected two visible rules, got %#v", resp.Items) + } + for _, item := range resp.Items { + if item.ID == "other-agent" { + t.Fatalf("other-agent scoped rule leaked into response: %#v", resp.Items) + } + } } func TestAgentRulesUsesDefaultAgentRecordForAgentScopedRules(t *testing.T) { @@ -1217,6 +1267,70 @@ func TestAgentRulesUsesDefaultAgentRecordForAgentScopedRules(t *testing.T) { } } +func TestAgentRulesNoAuthAgentIDFiltering(t *testing.T) { + agent := store.AgentRecord{ID: "agent-cuid-007", AgentIDs: `["agent-007"]`} + rules := &stubRulesRepo{rules: []store.Rule{ + {ID: "scoped", Action: invocation.RuleActionAutoApprove, ServerPatterns: []string{"amp"}, ToolPatterns: []string{"Read"}, AgentCUIDs: []string{agent.ID}, Enabled: true, Order: 0}, + {ID: "other", Action: invocation.RuleActionAutoDeny, ServerPatterns: []string{"amp"}, ToolPatterns: []string{"Read"}, AgentCUIDs: []string{"agent-other"}, Enabled: true, Order: 1}, + }} + agents := &stubAgentsRepo{records: []store.AgentRecord{agent}} + h := NewHandler(&stubService{}, stubServerService{}, nil, rules, agents, nil, nil, nil, nil, nil) + req := httptest.NewRequest(http.MethodGet, "/api/v1/agent/rules?agent_id=agent-007&source=amp&tool=Read", nil) + w := httptest.NewRecorder() + + h.Routes().ServeHTTP(w, req) + + if w.Code != http.StatusOK { + t.Fatalf("expected 200, got %d body=%s", w.Code, w.Body.String()) + } + var resp AgentRulesResponse + if err := json.Unmarshal(w.Body.Bytes(), &resp); err != nil { + t.Fatal(err) + } + if resp.Action != invocation.RuleActionAutoApprove { + t.Fatalf("expected scoped auto approve disposition, got %q", resp.Action) + } + if len(resp.Items) != 1 || resp.Items[0].ID != "scoped" { + t.Fatalf("expected only scoped visible rule, got %#v", resp.Items) + } +} + +func TestAgentRulesGuidanceForEachAction(t *testing.T) { + rules := &stubRulesRepo{rules: []store.Rule{ + {ID: "auto-approve", Action: invocation.RuleActionAutoApprove, ServerPatterns: []string{"amp"}, ToolPatterns: []string{"Read"}, Enabled: true, Order: 0}, + {ID: "auto-deny", Action: invocation.RuleActionAutoDeny, ServerPatterns: []string{"amp"}, ToolPatterns: []string{"Bash"}, Enabled: true, Order: 1}, + {ID: "human", Action: invocation.RuleActionHumanApproval, ServerPatterns: []string{"*"}, ToolPatterns: []string{"*"}, Enabled: true, Order: 2}, + {ID: "ai", Action: invocation.RuleActionAIEvaluation, ServerPatterns: []string{"amp"}, ToolPatterns: []string{"Search"}, Enabled: true, Order: 3}, + }} + h := NewHandler(&stubService{}, stubServerService{}, nil, rules, nil, nil, nil, nil, nil, nil) + req := httptest.NewRequest(http.MethodGet, "/api/v1/agent/rules", nil) + w := httptest.NewRecorder() + + h.Routes().ServeHTTP(w, req) + + if w.Code != http.StatusOK { + t.Fatalf("expected 200, got %d body=%s", w.Code, w.Body.String()) + } + var resp AgentRulesResponse + if err := json.Unmarshal(w.Body.Bytes(), &resp); err != nil { + t.Fatal(err) + } + byID := map[string]string{} + for _, item := range resp.Items { + byID[item.ID] = item.Guidance + } + for id, want := range map[string]string{ + "auto-approve": "Likely to pass", + "auto-deny": "Will be rejected", + "human": "Requires reviewer approval", + "ai": "real gated call", + } { + if !strings.Contains(byID[id], want) { + t.Fatalf("expected %s guidance to contain %q, got %q", id, want, byID[id]) + } + } +} + func TestMCPToolsListAnnotatesEffectiveAction(t *testing.T) { rules := &stubRulesRepo{rules: []store.Rule{ {ID: "read-auto", Action: invocation.RuleActionAutoApprove, ServerPatterns: []string{"demo"}, ToolPatterns: []string{"Read"}, Enabled: true, Order: 0}, @@ -1368,6 +1482,49 @@ func TestMCPToolsCallDenialIncludesRulesContext(t *testing.T) { } } +func TestMCPToolsCallDenialRulesContextFiltersAgentScopedRules(t *testing.T) { + now := time.Now().UTC() + agent := store.AgentRecord{ID: "agent-cuid-007", AgentIDs: `["agent-007"]`} + other := store.AgentRecord{ID: "agent-cuid-other", AgentIDs: `["other-agent"]`} + rules := &stubRulesRepo{rules: []store.Rule{ + {ID: "other-deny", Action: invocation.RuleActionAutoDeny, ServerPatterns: []string{"demo"}, ToolPatterns: []string{"Bash"}, AgentCUIDs: []string{other.ID}, Enabled: true, Order: 0}, + {ID: "agent-deny", Action: invocation.RuleActionAutoDeny, ServerPatterns: []string{"demo"}, ToolPatterns: []string{"Bash"}, AgentCUIDs: []string{agent.ID}, Enabled: true, Order: 1}, + }} + svc := &stubService{invoke: invocation.InvocationResponse{ + InvocationID: "inv_denied", ServerName: "demo", ToolName: "Bash", + Status: invocation.StatusDenied, + SubmittedAt: now, CompletedAt: &now, + Error: json.RawMessage(`{"content":[{"type":"text","text":"Tool call denied by approval rule (auto_deny)."}],"isError":true}`), + }} + agents := &stubAgentsRepo{records: []store.AgentRecord{agent, other}} + h := NewHandler(svc, stubServerService{}, nil, rules, agents, nil, nil, nil, nil, nil) + req := httptest.NewRequest(http.MethodPost, "/mcp/demo", strings.NewReader(`{"jsonrpc":"2.0","id":"agent-007","method":"tools/call","params":{"name":"Bash","arguments":{"cmd":"ls"}}}`)) + w := httptest.NewRecorder() + + h.Routes().ServeHTTP(w, req) + + var rpcResp struct { + Result struct { + Content []struct { + Text string `json:"text"` + } `json:"content"` + } `json:"result"` + } + if err := json.Unmarshal(w.Body.Bytes(), &rpcResp); err != nil { + t.Fatal(err) + } + if len(rpcResp.Result.Content) < 2 { + t.Fatalf("expected denial text plus rules context, got %#v", rpcResp.Result.Content) + } + last := rpcResp.Result.Content[len(rpcResp.Result.Content)-1].Text + if !strings.Contains(last, "agent-deny") { + t.Fatalf("expected agent scoped rule in denial context, got %q", last) + } + if strings.Contains(last, "other-deny") { + t.Fatalf("other agent rule leaked into denial context: %q", last) + } +} + func TestAdminInvocationsResponsesIncludeServerToolAndInput(t *testing.T) { now := time.Now().UTC() svc := &stubService{invoke: invocation.InvocationResponse{InvocationID: "inv_123", ServerName: "demo-server", ToolName: "demo_tool", Status: invocation.StatusSucceeded, Input: json.RawMessage(`{"issue":123,"verbose":true}`), SubmittedAt: now, CompletedAt: &now}} From 76257369e02c67fd698d56b049d77aca77e369e9 Mon Sep 17 00:00:00 2001 From: matt Date: Wed, 24 Jun 2026 10:43:50 -0700 Subject: [PATCH 2/7] update agent hooks --- examples/amp-plugin/atryum.ts | 80 +++++++++++++++++++++- examples/pi-extension/index.ts | 80 +++++++++++++++++++++- examples/shared-agent-hook/atryum-hook.mjs | 73 +++++++++++++++++++- 3 files changed, 224 insertions(+), 9 deletions(-) diff --git a/examples/amp-plugin/atryum.ts b/examples/amp-plugin/atryum.ts index 5046f43..33e965a 100644 --- a/examples/amp-plugin/atryum.ts +++ b/examples/amp-plugin/atryum.ts @@ -347,12 +347,86 @@ function describe(input: Record): string { return parts.join(" | ") || "(no string params)"; } +const RULES_CACHE_TTL_MS = 5 * 60 * 1000; +const rulesCache = new Map(); + +function formatRulesContext(rules: unknown): string { + if (!rules || typeof rules !== "object") return ""; + const record = rules as Record; + const lines = [ + "Atryum advisory rules visible to this harness before the gated call:", + `- server: ${String(record.server || SOURCE)}`, + `- tool: ${String(record.tool || "unknown")}`, + `- effective action: ${String(record.action || record.default_action || "unknown")}`, + ]; + if (record.matched_rule_id) { + lines.push(`- matched rule: ${String(record.matched_rule_id)}`); + } + if (record.generated_at) { + lines.push(`- as of: ${String(record.generated_at)}`); + } + if (Array.isArray(record.items) && record.items.length > 0) { + lines.push("- visible rules:"); + for (const item of record.items.slice(0, 20)) { + const rule = item as Record; + const guidance = rule.guidance ? ` (${String(rule.guidance)})` : ""; + lines.push( + ` - ${String(rule.id || "(unnamed)")}: ${String(rule.action)}${guidance}` + ); + } + if (record.items.length > 20) { + lines.push(` - ...${record.items.length - 20} more`); + } + } + lines.push("- advisory only; Atryum re-checks policy during the actual gated call."); + return lines.join("\n"); +} + +async function rulesContext(tool: string): Promise { + const cacheKey = [SOURCE, tool, ACCESS_TOKEN ? "auth" : "no-auth", AGENT_ID].join("\x00"); + const cached = rulesCache.get(cacheKey); + if (cached !== undefined && cached.expiresAt > Date.now()) return cached.value; + if (cached !== undefined) rulesCache.delete(cacheKey); + const url = new URL("/api/v1/agent/rules", API); + url.searchParams.set("server", SOURCE); + url.searchParams.set("tool", tool); + if (AGENT_ID && !ACCESS_TOKEN) { + url.searchParams.set("agent_id", AGENT_ID); + } + const controller = new AbortController(); + const timer = setTimeout(() => controller.abort(), 3000); + try { + const res = await fetch(url, { headers: atryumHeaders(), signal: controller.signal }); + if (!res.ok) return ""; + const result = formatRulesContext(await res.json()); + rulesCache.set(cacheKey, { + value: result, + expiresAt: Date.now() + RULES_CACHE_TTL_MS, + }); + return result; + } catch { + return ""; + } finally { + clearTimeout(timer); + } +} + +function combineContext( + rules: string, + chat: { context: string; count: number } | undefined +): { context: string; count: number | undefined } | undefined { + const context = [rules, chat?.context].filter(Boolean).join("\n\n"); + if (!context) return undefined; + return { context, count: chat?.count }; +} + async function submit( tool: string, toolUseID: string, input: Record, chat: { context: string; count: number } | undefined ): Promise { + const context = combineContext(await rulesContext(tool), chat); const res = await fetch(`${API}/api/v1/external/invocations`, { method: "POST", headers: atryumHeaders(true), @@ -363,9 +437,9 @@ async function submit( input, request_id: toolUseID, thread_id: activeThreadID() || undefined, - chat_context: chat?.context, - chat_context_messages: chat?.count, - context: chat?.context, + chat_context: context?.context, + chat_context_messages: context?.count, + context: context?.context, client_name: CLIENT_NAME, client_version: CLIENT_VERSION || undefined, agent_id: AGENT_ID || undefined, diff --git a/examples/pi-extension/index.ts b/examples/pi-extension/index.ts index 2981092..5145097 100644 --- a/examples/pi-extension/index.ts +++ b/examples/pi-extension/index.ts @@ -61,6 +61,79 @@ function describe(input: ToolInput): string { return parts.join(" | ") || "(no string params)"; } +const RULES_CACHE_TTL_MS = 5 * 60 * 1000; +const rulesCache = new Map(); + +function formatRulesContext(rules: unknown): string { + if (!rules || typeof rules !== "object") return ""; + const record = rules as Record; + const lines = [ + "Atryum advisory rules visible to this harness before the gated call:", + `- server: ${String(record.server || SOURCE)}`, + `- tool: ${String(record.tool || "unknown")}`, + `- effective action: ${String(record.action || record.default_action || "unknown")}`, + ]; + if (record.matched_rule_id) { + lines.push(`- matched rule: ${String(record.matched_rule_id)}`); + } + if (record.generated_at) { + lines.push(`- as of: ${String(record.generated_at)}`); + } + if (Array.isArray(record.items) && record.items.length > 0) { + lines.push("- visible rules:"); + for (const item of record.items.slice(0, 20)) { + const rule = item as Record; + const guidance = rule.guidance ? ` (${String(rule.guidance)})` : ""; + lines.push( + ` - ${String(rule.id || "(unnamed)")}: ${String(rule.action)}${guidance}` + ); + } + if (record.items.length > 20) { + lines.push(` - ...${record.items.length - 20} more`); + } + } + lines.push("- advisory only; Atryum re-checks policy during the actual gated call."); + return lines.join("\n"); +} + +async function rulesContext(tool: string): Promise { + const cacheKey = [SOURCE, tool, ACCESS_TOKEN ? "auth" : "no-auth", AGENT_ID].join("\x00"); + const cached = rulesCache.get(cacheKey); + if (cached !== undefined && cached.expiresAt > Date.now()) return cached.value; + if (cached !== undefined) rulesCache.delete(cacheKey); + const url = new URL("/api/v1/agent/rules", API); + url.searchParams.set("server", SOURCE); + url.searchParams.set("tool", tool); + if (AGENT_ID && !ACCESS_TOKEN) { + url.searchParams.set("agent_id", AGENT_ID); + } + const controller = new AbortController(); + const timer = setTimeout(() => controller.abort(), 3000); + try { + const res = await fetch(url, { headers: atryumHeaders(), signal: controller.signal }); + if (!res.ok) return ""; + const result = formatRulesContext(await res.json()); + rulesCache.set(cacheKey, { + value: result, + expiresAt: Date.now() + RULES_CACHE_TTL_MS, + }); + return result; + } catch { + return ""; + } finally { + clearTimeout(timer); + } +} + +function combineContext( + rules: string, + chat: { context: string; count: number } | undefined +): { context: string; count: number | undefined } | undefined { + const context = [rules, chat?.context].filter(Boolean).join("\n\n"); + if (!context) return undefined; + return { context, count: chat?.count }; +} + function sessionID(ctx: unknown): string | undefined { const manager = (ctx as { sessionManager?: unknown }).sessionManager as | { getSessionFile?: () => string; sessionId?: string; id?: string } @@ -245,6 +318,7 @@ async function submit( threadID: string | undefined, chat: { context: string; count: number } | undefined ): Promise { + const context = combineContext(await rulesContext(tool), chat); const res = await fetch(`${API}/api/v1/external/invocations`, { method: "POST", headers: atryumHeaders(true), @@ -255,9 +329,9 @@ async function submit( input, request_id: toolCallID, thread_id: threadID, - chat_context: chat?.context, - chat_context_messages: chat?.count, - context: chat?.context, + chat_context: context?.context, + chat_context_messages: context?.count, + context: context?.context, agent_id: AGENT_ID || undefined, client_name: CLIENT_NAME, client_version: CLIENT_VERSION || undefined, diff --git a/examples/shared-agent-hook/atryum-hook.mjs b/examples/shared-agent-hook/atryum-hook.mjs index a297232..26b197b 100644 --- a/examples/shared-agent-hook/atryum-hook.mjs +++ b/examples/shared-agent-hook/atryum-hook.mjs @@ -118,6 +118,72 @@ function describe(input) { return parts.join(" | ") || "(no string params)"; } +const RULES_CACHE_TTL_MS = 5 * 60 * 1000; +const rulesCache = new Map(); + +function formatRulesContext(rules) { + if (!rules || typeof rules !== "object") return ""; + const lines = [ + "Atryum advisory rules visible to this harness before the gated call:", + `- server: ${rules.server || SOURCE}`, + `- tool: ${rules.tool || "unknown"}`, + `- effective action: ${rules.action || rules.default_action || "unknown"}`, + ]; + if (rules.matched_rule_id) { + lines.push(`- matched rule: ${rules.matched_rule_id}`); + } + if (rules.generated_at) { + lines.push(`- as of: ${rules.generated_at}`); + } + if (Array.isArray(rules.items) && rules.items.length > 0) { + lines.push("- visible rules:"); + for (const rule of rules.items.slice(0, 20)) { + const guidance = rule.guidance ? ` (${rule.guidance})` : ""; + lines.push(` - ${rule.id || "(unnamed)"}: ${rule.action}${guidance}`); + } + if (rules.items.length > 20) { + lines.push(` - ...${rules.items.length - 20} more`); + } + } + lines.push("- advisory only; Atryum re-checks policy during the actual gated call."); + return lines.join("\n"); +} + +async function rulesContext(tool) { + const cacheKey = [SOURCE, tool, ACCESS_TOKEN ? "auth" : "no-auth", AGENT_ID].join("\x00"); + const cached = rulesCache.get(cacheKey); + if (cached !== undefined && cached.expiresAt > Date.now()) return cached.value; + if (cached !== undefined) rulesCache.delete(cacheKey); + const url = new URL("/api/v1/agent/rules", API); + url.searchParams.set("server", SOURCE); + url.searchParams.set("tool", tool); + if (AGENT_ID && !ACCESS_TOKEN) { + url.searchParams.set("agent_id", AGENT_ID); + } + const controller = new AbortController(); + const timer = setTimeout(() => controller.abort(), 3000); + try { + const res = await fetch(url, { headers: atryumHeaders(), signal: controller.signal }); + if (!res.ok) return ""; + const result = formatRulesContext(await res.json()); + rulesCache.set(cacheKey, { + value: result, + expiresAt: Date.now() + RULES_CACHE_TTL_MS, + }); + return result; + } catch { + return ""; + } finally { + clearTimeout(timer); + } +} + +function combineContext(rules, chat) { + const context = [rules, chat?.context].filter(Boolean).join("\n\n"); + if (!context) return undefined; + return { context, count: chat?.count }; +} + function normalizeRole(value) { const role = String(value || "").toLowerCase(); if (role === "human") return "user"; @@ -462,6 +528,7 @@ async function submit(event) { const input = toolInput(event); const id = toolUseID(event); const chat = await chatContext(event); + const context = combineContext(await rulesContext(name), chat); const res = await fetch(`${API}/api/v1/external/invocations`, { method: "POST", headers: atryumHeaders(true), @@ -472,9 +539,9 @@ async function submit(event) { input, request_id: id, thread_id: sessionId(event) || undefined, - chat_context: chat?.context, - chat_context_messages: chat?.count, - context: chat?.context, + chat_context: context?.context, + chat_context_messages: context?.count, + context: context?.context, client_name: CLIENT_NAME, client_version: CLIENT_VERSION || undefined, agent_id: AGENT_ID || undefined, From 6a9c2883b6d2250f5b1360b51641e20ebdc690c5 Mon Sep 17 00:00:00 2001 From: matt Date: Wed, 24 Jun 2026 11:02:44 -0700 Subject: [PATCH 3/7] handle no auth mode --- internal/api/handlers.go | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/internal/api/handlers.go b/internal/api/handlers.go index d01509f..a8eae5c 100644 --- a/internal/api/handlers.go +++ b/internal/api/handlers.go @@ -1212,7 +1212,11 @@ func (h *Handler) handleMCPProxy(w http.ResponseWriter, r *http.Request, server return } _ = h.emitTraceEvent(r.Context(), server, "mcp.tools.list", map[string]any{"tool_count": len(tools), "request_id": requestID}) - annotated := h.annotateToolsWithPolicy(r.Context(), server, tools) + listAgentID := auth.AgentIDFromContext(r.Context()) + if listAgentID == "" && h.authValidator == nil { + listAgentID = normalizeNoAuthRequestAgentID(requestID) + } + annotated := h.annotateToolsWithPolicy(r.Context(), listAgentID, server, tools) h.writeRPCResult(w, req.ID, map[string]any{"tools": annotated}) case "tools/call": var params struct { @@ -1505,7 +1509,7 @@ type atryumToolPolicy struct { // annotateToolsWithPolicy decorates each tool with its effective approval // disposition for the current agent so the model sees the policy at the moment // it picks a tool. Annotation requires both rulesRepo and a concrete server. -func (h *Handler) annotateToolsWithPolicy(ctx context.Context, server string, tools []mcp.Tool) []any { +func (h *Handler) annotateToolsWithPolicy(ctx context.Context, agentID, server string, tools []mcp.Tool) []any { out := make([]any, len(tools)) if h.rulesRepo == nil || strings.TrimSpace(server) == "" { for i, t := range tools { @@ -1520,7 +1524,6 @@ func (h *Handler) annotateToolsWithPolicy(ctx context.Context, server string, to } return out } - agentID := auth.AgentIDFromContext(ctx) agentCUID := h.resolveAgentRecordForRules(ctx, agentID) for i, t := range tools { action, matched := effectiveActionForTool(rules, server, t.Name, agentCUID) From 4924d0694d01005e0dc569dc993180caf34de495 Mon Sep 17 00:00:00 2001 From: matt Date: Wed, 24 Jun 2026 11:58:56 -0700 Subject: [PATCH 4/7] undo change --- internal/api/handlers.go | 22 ++-------------- internal/api/handlers_test.go | 48 ++++++++++++++++++++++++++++++++++- 2 files changed, 49 insertions(+), 21 deletions(-) diff --git a/internal/api/handlers.go b/internal/api/handlers.go index a8eae5c..6075588 100644 --- a/internal/api/handlers.go +++ b/internal/api/handlers.go @@ -1212,11 +1212,7 @@ func (h *Handler) handleMCPProxy(w http.ResponseWriter, r *http.Request, server return } _ = h.emitTraceEvent(r.Context(), server, "mcp.tools.list", map[string]any{"tool_count": len(tools), "request_id": requestID}) - listAgentID := auth.AgentIDFromContext(r.Context()) - if listAgentID == "" && h.authValidator == nil { - listAgentID = normalizeNoAuthRequestAgentID(requestID) - } - annotated := h.annotateToolsWithPolicy(r.Context(), listAgentID, server, tools) + annotated := h.annotateToolsWithPolicy(r.Context(), auth.AgentIDFromContext(r.Context()), server, tools) h.writeRPCResult(w, req.ID, map[string]any{"tools": annotated}) case "tools/call": var params struct { @@ -1248,11 +1244,7 @@ func (h *Handler) handleMCPProxy(w http.ResponseWriter, r *http.Request, server if len(resp.Error) > 0 { result := normalizeToolCallResult(resp.Error, true) if resp.Status == invocation.StatusDenied { - agentID := auth.AgentIDFromContext(r.Context()) - if agentID == "" && h.authValidator == nil { - agentID = normalizeNoAuthRequestAgentID(requestID) - } - result = h.appendRulesContextToToolResult(r.Context(), result, agentID, server, params.Name) + result = h.appendRulesContextToToolResult(r.Context(), result, auth.AgentIDFromContext(r.Context()), server, params.Name) } h.writeRPCResult(w, req.ID, result) return @@ -3967,16 +3959,6 @@ func compactRequestID(id json.RawMessage) string { return string(id) } -func normalizeNoAuthRequestAgentID(value string) string { - value = strings.TrimSpace(value) - if strings.HasPrefix(value, `"`) { - if unquoted, err := strconv.Unquote(value); err == nil { - value = unquoted - } - } - return normalizeNoAuthAgentID(value) -} - func stringPtr(v string) *string { return &v } func invocationSignature(items []invocation.InvocationResponse) string { diff --git a/internal/api/handlers_test.go b/internal/api/handlers_test.go index e7d498a..240b421 100644 --- a/internal/api/handlers_test.go +++ b/internal/api/handlers_test.go @@ -1441,6 +1441,52 @@ func TestMCPToolsListAnnotationsUseDefaultAgentScopedRules(t *testing.T) { } } +func TestMCPToolsListAnnotationsIgnoreJSONRPCIDForNoAuthAgent(t *testing.T) { + agent := store.AgentRecord{ID: "agent-cuid-007", AgentIDs: `["agent-007"]`} + rules := &stubRulesRepo{rules: []store.Rule{ + {ID: "agent-auto", Action: invocation.RuleActionAutoApprove, ServerPatterns: []string{"demo"}, ToolPatterns: []string{"Bash"}, AgentCUIDs: []string{agent.ID}, Enabled: true, Order: 0}, + }} + svc := &stubService{tools: []mcp.Tool{{Name: "Bash", Description: "run a shell command"}}} + agents := &stubAgentsRepo{records: []store.AgentRecord{agent}} + h := NewHandler(svc, stubServerService{}, nil, rules, agents, nil, nil, nil, nil, nil) + req := httptest.NewRequest(http.MethodPost, "/mcp/demo", strings.NewReader(`{"jsonrpc":"2.0","id":"agent-007","method":"tools/list","params":{}}`)) + w := httptest.NewRecorder() + + h.Routes().ServeHTTP(w, req) + + if w.Code != http.StatusOK { + t.Fatalf("expected 200, got %d body=%s", w.Code, w.Body.String()) + } + var rpcResp struct { + Result struct { + Tools []struct { + Annotations *struct { + Atryum struct { + EffectiveAction string `json:"effective_action"` + MatchedRuleID string `json:"matched_rule_id"` + } `json:"atryum"` + } `json:"annotations"` + } `json:"tools"` + } `json:"result"` + } + if err := json.Unmarshal(w.Body.Bytes(), &rpcResp); err != nil { + t.Fatal(err) + } + if len(rpcResp.Result.Tools) != 1 { + t.Fatalf("expected one tool, got %#v", rpcResp.Result.Tools) + } + got := rpcResp.Result.Tools[0].Annotations + if got == nil { + t.Fatal("expected annotation") + } + if got.Atryum.EffectiveAction != invocation.RuleActionHumanApproval { + t.Fatalf("expected anonymous/default policy, got %#v", got.Atryum) + } + if got.Atryum.MatchedRuleID != "" { + t.Fatalf("json-rpc id should not match agent-scoped rule, got %q", got.Atryum.MatchedRuleID) + } +} + func TestMCPToolsCallDenialIncludesRulesContext(t *testing.T) { now := time.Now().UTC() rules := &stubRulesRepo{rules: []store.Rule{ @@ -1498,7 +1544,7 @@ func TestMCPToolsCallDenialRulesContextFiltersAgentScopedRules(t *testing.T) { }} agents := &stubAgentsRepo{records: []store.AgentRecord{agent, other}} h := NewHandler(svc, stubServerService{}, nil, rules, agents, nil, nil, nil, nil, nil) - req := httptest.NewRequest(http.MethodPost, "/mcp/demo", strings.NewReader(`{"jsonrpc":"2.0","id":"agent-007","method":"tools/call","params":{"name":"Bash","arguments":{"cmd":"ls"}}}`)) + req := httptest.NewRequest(http.MethodPost, "/mcp/demo?agent_id=agent-007", strings.NewReader(`{"jsonrpc":"2.0","id":"request-1","method":"tools/call","params":{"name":"Bash","arguments":{"cmd":"ls"}}}`)) w := httptest.NewRecorder() h.Routes().ServeHTTP(w, req) From 17ad13126344fb5b7375a6b10663e7d32b24ed81 Mon Sep 17 00:00:00 2001 From: matt Date: Wed, 24 Jun 2026 15:32:41 -0700 Subject: [PATCH 5/7] minor enhancements --- README.md | 2 +- internal/api/handlers.go | 168 +++++++++++++++++++++++++++++----- internal/api/handlers_test.go | 78 ++++++++++++++-- 3 files changed, 219 insertions(+), 29 deletions(-) diff --git a/README.md b/README.md index 036f123..ff19799 100644 --- a/README.md +++ b/README.md @@ -129,7 +129,7 @@ The frontend fetches `/api/v1/admin-auth/config`, shows a sign-in screen, redire Public (auth-protected when `[[auth]]` is configured): -- `POST /mcp/{server}` — MCP JSON-RPC. `tools/list` annotates each tool with its policy disposition for the calling agent; a synthetic `atryum.rules.get` tool lets an agent inspect its applicable rules before deciding what to call. +- `POST /mcp/{server}` — MCP JSON-RPC. `tools/list` annotates each tool with its policy disposition for the calling agent; a synthetic `atryum_rules_get` tool lets an agent inspect its applicable rules before deciding what to call. The underscore name is intentional: dotted MCP tool names are spec-compliant, but some common harness implementations reject them. - `GET /mcp/{server}` — Streamable HTTP / legacy SSE channel for MCP clients that need a long-lived event stream. - `POST /api/v1/invocations` — direct invocation (Atryum executes). - `POST /api/v1/external/invocations`, `PATCH /api/v1/external/invocations/{id}` — hook path (harness executes, Atryum gates and records). diff --git a/internal/api/handlers.go b/internal/api/handlers.go index 6075588..6493669 100644 --- a/internal/api/handlers.go +++ b/internal/api/handlers.go @@ -156,6 +156,16 @@ type Handler struct { clientInfoMu sync.RWMutex clientInfoCache map[string]clientInfoSnapshot + // rulesListMu / rulesListCache avoids redundant DB round-trips when + // multiple operations in the same request window need the full rule list. + // Entries expire after 10 s — rules are admin-configured and change rarely. + rulesListMu sync.RWMutex + rulesListCache *rulesListCacheEntry + + // agentCUIDCache avoids repeated DB lookups for agentID → CUID resolution + // on every tools/list. Entries expire after 30 s. + agentCUIDCache sync.Map + // managedAgents is the optional Claude Managed Agents events bridge. // nil when not configured (no anthropic api key). managedAgents managedAgentsAdmin @@ -617,10 +627,14 @@ func parseAgentIDs(raw string) []string { const atryumInitializeInstructions = "This MCP server is gated by the Atryum harness. " + "Atryum may approve, deny, or request human approval for tool calls according to configured rules. " + "Those rules may change between conversation turns, so treat each tool call as subject to the current Atryum policy. " + - "To see the static approval rules that currently apply to you, issue an HTTP GET to /api/v1/agent/rules " + + "To see the static approval rules that currently apply to you, call the atryum_rules_get MCP tool or issue an HTTP GET to /api/v1/agent/rules " + "(optionally with ?server={server}&tool={tool} to preview the disposition for a specific tool); " + "the response is advisory only, as AI-evaluation and human-approval outcomes are decided during the actual gated call." +// Dotted tool names are valid MCP names, but common harnesses have rejected +// them in practice. Keep this synthetic helper underscore-only for compatibility. +const atryumRulesToolName = "atryum_rules_get" + type AgentRule struct { ID string `json:"id"` Action string `json:"action"` @@ -1005,11 +1019,6 @@ func (h *Handler) agentRules(w http.ResponseWriter, r *http.Request) { writeError(w, http.StatusMethodNotAllowed, "method not allowed") return } - if h.rulesRepo == nil { - writeJSON(w, http.StatusOK, newAgentRulesResponse("", "", "")) - return - } - agentID := auth.AgentIDFromContext(r.Context()) if agentID == "" && h.authValidator == nil { agentID = normalizeNoAuthAgentID(r.URL.Query().Get("agent_id")) @@ -1023,6 +1032,11 @@ func (h *Handler) agentRules(w http.ResponseWriter, r *http.Request) { } tool := strings.TrimSpace(r.URL.Query().Get("tool")) + if h.rulesRepo == nil { + writeJSON(w, http.StatusOK, newAgentRulesResponse(agentID, server, tool)) + return + } + resp, err := h.buildAgentRulesResponse(r.Context(), agentID, server, tool) if err != nil { writeError(w, http.StatusInternalServerError, err.Error()) @@ -1037,7 +1051,7 @@ func (h *Handler) buildAgentRulesResponse(ctx context.Context, agentID, server, return resp, nil } - rules, err := h.rulesRepo.List(ctx) + rules, err := h.cachedRulesList(ctx) if err != nil { return AgentRulesResponse{}, err } @@ -1068,6 +1082,56 @@ func (h *Handler) buildAgentRulesResponse(ctx context.Context, agentID, server, return resp, nil } +func (h *Handler) handleAtryumRulesToolCall(w http.ResponseWriter, r *http.Request, id json.RawMessage, mcpServer string, arguments map[string]any) { + if arguments == nil { + arguments = map[string]any{} + } + agentID := auth.AgentIDFromContext(r.Context()) + if agentID == "" && h.authValidator == nil { + agentID = normalizeNoAuthAgentID(stringArgument(arguments, "agent_id")) + } + if agentID == "" && h.authValidator == nil { + agentID = strings.TrimSpace(stringArgument(arguments, "request_id")) + } + server := strings.TrimSpace(stringArgument(arguments, "server")) + if server == "" { + server = strings.TrimSpace(stringArgument(arguments, "source")) + } + if server == "" { + server = mcpServer + } + tool := strings.TrimSpace(stringArgument(arguments, "tool")) + + resp, err := h.buildAgentRulesResponse(r.Context(), agentID, server, tool) + if err != nil { + h.writeRPCError(w, id, -32000, err.Error()) + return + } + body, err := json.MarshalIndent(resp, "", " ") + if err != nil { + h.writeRPCError(w, id, -32000, err.Error()) + return + } + h.writeRPCResult(w, id, map[string]any{ + "content": []map[string]any{{ + "type": "text", + "text": string(body), + }}, + }) +} + +func stringArgument(arguments map[string]any, key string) string { + v, ok := arguments[key] + if !ok { + return "" + } + s, ok := v.(string) + if !ok { + return "" + } + return s +} + func newAgentRulesResponse(agentID, server, tool string) AgentRulesResponse { return AgentRulesResponse{ AgentID: agentID, @@ -1082,10 +1146,55 @@ func newAgentRulesResponse(agentID, server, tool string) AgentRulesResponse { } } +type rulesListCacheEntry struct { + rules []store.Rule + expiresAt time.Time +} + +func (h *Handler) cachedRulesList(ctx context.Context) ([]store.Rule, error) { + h.rulesListMu.RLock() + if h.rulesListCache != nil && time.Now().Before(h.rulesListCache.expiresAt) { + rules := h.rulesListCache.rules + h.rulesListMu.RUnlock() + return rules, nil + } + h.rulesListMu.RUnlock() + + rules, err := h.rulesRepo.List(ctx) + if err != nil { + return nil, err + } + + h.rulesListMu.Lock() + defer h.rulesListMu.Unlock() + h.rulesListCache = &rulesListCacheEntry{rules: rules, expiresAt: time.Now().Add(10 * time.Second)} + return rules, nil +} + +type agentCUIDCacheEntry struct { + cuid string + expiresAt time.Time +} + func (h *Handler) resolveAgentRecordForRules(ctx context.Context, agentID string) string { if h.agentsRepo == nil { return "" } + cacheKey := agentID + if cacheKey == "" { + cacheKey = "\x00default" + } + if v, ok := h.agentCUIDCache.Load(cacheKey); ok { + if e := v.(agentCUIDCacheEntry); time.Now().Before(e.expiresAt) { + return e.cuid + } + } + cuid := h.resolveAgentCUIDUncached(ctx, agentID) + h.agentCUIDCache.Store(cacheKey, agentCUIDCacheEntry{cuid: cuid, expiresAt: time.Now().Add(30 * time.Second)}) + return cuid +} + +func (h *Handler) resolveAgentCUIDUncached(ctx context.Context, agentID string) string { if agentID != "" { if rec, err := h.agentsRepo.GetByAgentID(ctx, agentID); err == nil { return rec.ID @@ -1223,6 +1332,10 @@ func (h *Handler) handleMCPProxy(w http.ResponseWriter, r *http.Request, server h.writeRPCError(w, req.ID, -32602, "invalid params") return } + if params.Name == atryumRulesToolName { + h.handleAtryumRulesToolCall(w, r, req.ID, server, params.Arguments) + return + } toolReq := invocation.CreateInvocationRequest{Server: server, Tool: params.Name, Input: params.Arguments} if requestID != "" { toolReq.RequestID = stringPtr(requestID) @@ -1466,9 +1579,9 @@ func apiMatchPatterns(patterns []string, value string) bool { func agentRuleGuidance(action string) string { switch action { case invocation.RuleActionAutoApprove: - return "Likely to pass if policy and rule inputs do not change." + return "Auto-approved by this rule." case invocation.RuleActionAutoDeny: - return "Will be rejected by this matched rule." + return "Auto-denied by this rule." case invocation.RuleActionHumanApproval: return "Requires reviewer approval." case invocation.RuleActionAIEvaluation: @@ -1502,22 +1615,29 @@ type atryumToolPolicy struct { // disposition for the current agent so the model sees the policy at the moment // it picks a tool. Annotation requires both rulesRepo and a concrete server. func (h *Handler) annotateToolsWithPolicy(ctx context.Context, agentID, server string, tools []mcp.Tool) []any { - out := make([]any, len(tools)) + out := make([]any, 0, len(tools)+1) if h.rulesRepo == nil || strings.TrimSpace(server) == "" { - for i, t := range tools { - out[i] = t + for _, t := range tools { + if t.Name != atryumRulesToolName { + out = append(out, t) + } } - return out + return append(out, atryumRulesTool) } - rules, err := h.rulesRepo.List(ctx) + rules, err := h.cachedRulesList(ctx) if err != nil { - for i, t := range tools { - out[i] = t + for _, t := range tools { + if t.Name != atryumRulesToolName { + out = append(out, t) + } } - return out + return append(out, atryumRulesTool) } agentCUID := h.resolveAgentRecordForRules(ctx, agentID) - for i, t := range tools { + for _, t := range tools { + if t.Name == atryumRulesToolName { + continue + } action, matched := effectiveActionForTool(rules, server, t.Name, agentCUID) desc := t.Description if action != "" { @@ -1528,7 +1648,7 @@ func (h *Handler) annotateToolsWithPolicy(ctx context.Context, agentID, server s desc = prefix + desc } } - out[i] = annotatedTool{ + out = append(out, annotatedTool{ Name: t.Name, Description: desc, InputSchema: t.InputSchema, @@ -1536,9 +1656,15 @@ func (h *Handler) annotateToolsWithPolicy(ctx context.Context, agentID, server s EffectiveAction: action, MatchedRuleID: matched, }}, - } + }) } - return out + return append(out, atryumRulesTool) +} + +var atryumRulesTool = mcp.Tool{ + Name: atryumRulesToolName, + Description: "Return the current static Atryum approval rules visible to this agent. Optional arguments: server/source and tool preview a specific disposition. The response is advisory only; AI evaluation and human approval are decided during the actual gated tool call.", + InputSchema: json.RawMessage(`{"type":"object","properties":{"server":{"type":"string","description":"MCP server name to preview."},"source":{"type":"string","description":"Alias for server, useful for non-MCP harness terminology."},"tool":{"type":"string","description":"Tool name to preview."},"agent_id":{"type":"string","description":"Agent identity for no-auth local development only; ignored when inbound auth is enabled."}},"additionalProperties":false}`), } // effectiveActionForTool returns the action of the first enabled rule that diff --git a/internal/api/handlers_test.go b/internal/api/handlers_test.go index 240b421..9dafddb 100644 --- a/internal/api/handlers_test.go +++ b/internal/api/handlers_test.go @@ -1055,7 +1055,10 @@ func TestMCPToolsList(t *testing.T) { t.Fatalf("expected tools list, got %s", w.Body.String()) } if strings.Contains(w.Body.String(), `"atryum.rules.get"`) { - t.Fatalf("did not expect synthetic rules tool, got %s", w.Body.String()) + t.Fatalf("did not expect dotted synthetic rules tool, got %s", w.Body.String()) + } + if !strings.Contains(w.Body.String(), `"`+atryumRulesToolName+`"`) { + t.Fatalf("expected compatible synthetic rules tool, got %s", w.Body.String()) } } @@ -1082,6 +1085,48 @@ func TestMCPToolsCallInterceptsInvocation(t *testing.T) { } } +func TestMCPRulesToolReturnsAgentRulesWithoutInvocation(t *testing.T) { + rules := &stubRulesRepo{rules: []store.Rule{ + {ID: "read-auto", Action: invocation.RuleActionAutoApprove, ServerPatterns: []string{"demo"}, ToolPatterns: []string{"Read"}, Enabled: true, Order: 0}, + }} + svc := &stubService{} + h := NewHandler(svc, stubServerService{}, nil, rules, nil, nil, nil, nil, nil, nil) + req := httptest.NewRequest(http.MethodPost, "/mcp/demo", strings.NewReader(`{"jsonrpc":"2.0","id":7,"method":"tools/call","params":{"name":"atryum_rules_get","arguments":{"tool":"Read"}}}`)) + w := httptest.NewRecorder() + + h.Routes().ServeHTTP(w, req) + + if svc.invokedReq != nil { + t.Fatalf("rules helper should not create a gated invocation, got %#v", svc.invokedReq) + } + var rpcResp struct { + Result struct { + Content []struct { + Text string `json:"text"` + } `json:"content"` + } `json:"result"` + } + if err := json.Unmarshal(w.Body.Bytes(), &rpcResp); err != nil { + t.Fatal(err) + } + if len(rpcResp.Result.Content) != 1 { + t.Fatalf("expected one content block, got %#v", rpcResp.Result.Content) + } + var rulesResp AgentRulesResponse + if err := json.Unmarshal([]byte(rpcResp.Result.Content[0].Text), &rulesResp); err != nil { + t.Fatal(err) + } + if rulesResp.Server != "demo" || rulesResp.Tool != "Read" { + t.Fatalf("expected demo/Read preview, got %#v", rulesResp) + } + if rulesResp.Action != invocation.RuleActionAutoApprove { + t.Fatalf("expected auto approve disposition, got %q", rulesResp.Action) + } + if rulesResp.MatchedRuleID == nil || *rulesResp.MatchedRuleID != "read-auto" { + t.Fatalf("expected read-auto match, got %#v", rulesResp.MatchedRuleID) + } +} + func TestMCPNoAuthAgentIDQueryHintSetsIdentity(t *testing.T) { now := time.Now().UTC() svc := &stubService{invoke: invocation.InvocationResponse{InvocationID: "inv_123", ServerName: "demo", ToolName: "demo_tool", Status: invocation.StatusSucceeded, SubmittedAt: now, CompletedAt: &now, Result: json.RawMessage(`{"content":[{"type":"text","text":"ok"}]}`)}} @@ -1320,8 +1365,8 @@ func TestAgentRulesGuidanceForEachAction(t *testing.T) { byID[item.ID] = item.Guidance } for id, want := range map[string]string{ - "auto-approve": "Likely to pass", - "auto-deny": "Will be rejected", + "auto-approve": "Auto-approved", + "auto-deny": "Auto-denied", "human": "Requires reviewer approval", "ai": "real gated call", } { @@ -1383,7 +1428,10 @@ func TestMCPToolsListAnnotatesEffectiveAction(t *testing.T) { t.Fatalf("Other tool annotations: %#v", otherTool.Annotations) } if _, ok := byName["atryum.rules.get"]; ok { - t.Fatalf("did not expect synthetic rules tool in tools/list") + t.Fatalf("did not expect dotted synthetic rules tool in tools/list") + } + if _, ok := byName[atryumRulesToolName]; !ok { + t.Fatalf("expected compatible synthetic rules tool in tools/list") } } @@ -1460,6 +1508,7 @@ func TestMCPToolsListAnnotationsIgnoreJSONRPCIDForNoAuthAgent(t *testing.T) { var rpcResp struct { Result struct { Tools []struct { + Name string `json:"name"` Annotations *struct { Atryum struct { EffectiveAction string `json:"effective_action"` @@ -1472,10 +1521,25 @@ func TestMCPToolsListAnnotationsIgnoreJSONRPCIDForNoAuthAgent(t *testing.T) { if err := json.Unmarshal(w.Body.Bytes(), &rpcResp); err != nil { t.Fatal(err) } - if len(rpcResp.Result.Tools) != 1 { - t.Fatalf("expected one tool, got %#v", rpcResp.Result.Tools) + var bashTool *struct { + Name string `json:"name"` + Annotations *struct { + Atryum struct { + EffectiveAction string `json:"effective_action"` + MatchedRuleID string `json:"matched_rule_id"` + } `json:"atryum"` + } `json:"annotations"` + } + for i := range rpcResp.Result.Tools { + if rpcResp.Result.Tools[i].Name == "Bash" { + bashTool = &rpcResp.Result.Tools[i] + break + } + } + if bashTool == nil { + t.Fatalf("expected Bash tool, got %#v", rpcResp.Result.Tools) } - got := rpcResp.Result.Tools[0].Annotations + got := bashTool.Annotations if got == nil { t.Fatal("expected annotation") } From 881a0453f030dff9b6fe3c640eeb3ffbfc179d9b Mon Sep 17 00:00:00 2001 From: matt Date: Wed, 24 Jun 2026 15:39:07 -0700 Subject: [PATCH 6/7] remove incorrect caching --- internal/api/handlers.go | 59 ++-------------------------------------- 1 file changed, 2 insertions(+), 57 deletions(-) diff --git a/internal/api/handlers.go b/internal/api/handlers.go index 6493669..57a149b 100644 --- a/internal/api/handlers.go +++ b/internal/api/handlers.go @@ -156,16 +156,6 @@ type Handler struct { clientInfoMu sync.RWMutex clientInfoCache map[string]clientInfoSnapshot - // rulesListMu / rulesListCache avoids redundant DB round-trips when - // multiple operations in the same request window need the full rule list. - // Entries expire after 10 s — rules are admin-configured and change rarely. - rulesListMu sync.RWMutex - rulesListCache *rulesListCacheEntry - - // agentCUIDCache avoids repeated DB lookups for agentID → CUID resolution - // on every tools/list. Entries expire after 30 s. - agentCUIDCache sync.Map - // managedAgents is the optional Claude Managed Agents events bridge. // nil when not configured (no anthropic api key). managedAgents managedAgentsAdmin @@ -1051,7 +1041,7 @@ func (h *Handler) buildAgentRulesResponse(ctx context.Context, agentID, server, return resp, nil } - rules, err := h.cachedRulesList(ctx) + rules, err := h.rulesRepo.List(ctx) if err != nil { return AgentRulesResponse{}, err } @@ -1146,55 +1136,10 @@ func newAgentRulesResponse(agentID, server, tool string) AgentRulesResponse { } } -type rulesListCacheEntry struct { - rules []store.Rule - expiresAt time.Time -} - -func (h *Handler) cachedRulesList(ctx context.Context) ([]store.Rule, error) { - h.rulesListMu.RLock() - if h.rulesListCache != nil && time.Now().Before(h.rulesListCache.expiresAt) { - rules := h.rulesListCache.rules - h.rulesListMu.RUnlock() - return rules, nil - } - h.rulesListMu.RUnlock() - - rules, err := h.rulesRepo.List(ctx) - if err != nil { - return nil, err - } - - h.rulesListMu.Lock() - defer h.rulesListMu.Unlock() - h.rulesListCache = &rulesListCacheEntry{rules: rules, expiresAt: time.Now().Add(10 * time.Second)} - return rules, nil -} - -type agentCUIDCacheEntry struct { - cuid string - expiresAt time.Time -} - func (h *Handler) resolveAgentRecordForRules(ctx context.Context, agentID string) string { if h.agentsRepo == nil { return "" } - cacheKey := agentID - if cacheKey == "" { - cacheKey = "\x00default" - } - if v, ok := h.agentCUIDCache.Load(cacheKey); ok { - if e := v.(agentCUIDCacheEntry); time.Now().Before(e.expiresAt) { - return e.cuid - } - } - cuid := h.resolveAgentCUIDUncached(ctx, agentID) - h.agentCUIDCache.Store(cacheKey, agentCUIDCacheEntry{cuid: cuid, expiresAt: time.Now().Add(30 * time.Second)}) - return cuid -} - -func (h *Handler) resolveAgentCUIDUncached(ctx context.Context, agentID string) string { if agentID != "" { if rec, err := h.agentsRepo.GetByAgentID(ctx, agentID); err == nil { return rec.ID @@ -1624,7 +1569,7 @@ func (h *Handler) annotateToolsWithPolicy(ctx context.Context, agentID, server s } return append(out, atryumRulesTool) } - rules, err := h.cachedRulesList(ctx) + rules, err := h.rulesRepo.List(ctx) if err != nil { for _, t := range tools { if t.Name != atryumRulesToolName { From bb215d4d7ff0e365abc3512dbdd5d002d087a048 Mon Sep 17 00:00:00 2001 From: matt Date: Mon, 6 Jul 2026 16:18:09 -0700 Subject: [PATCH 7/7] add request_id to mcp --- internal/api/handlers.go | 2 +- internal/api/handlers_test.go | 32 ++++++++++++++++++++++++++++++++ 2 files changed, 33 insertions(+), 1 deletion(-) diff --git a/internal/api/handlers.go b/internal/api/handlers.go index 57a149b..9b884b4 100644 --- a/internal/api/handlers.go +++ b/internal/api/handlers.go @@ -1609,7 +1609,7 @@ func (h *Handler) annotateToolsWithPolicy(ctx context.Context, agentID, server s var atryumRulesTool = mcp.Tool{ Name: atryumRulesToolName, Description: "Return the current static Atryum approval rules visible to this agent. Optional arguments: server/source and tool preview a specific disposition. The response is advisory only; AI evaluation and human approval are decided during the actual gated tool call.", - InputSchema: json.RawMessage(`{"type":"object","properties":{"server":{"type":"string","description":"MCP server name to preview."},"source":{"type":"string","description":"Alias for server, useful for non-MCP harness terminology."},"tool":{"type":"string","description":"Tool name to preview."},"agent_id":{"type":"string","description":"Agent identity for no-auth local development only; ignored when inbound auth is enabled."}},"additionalProperties":false}`), + InputSchema: json.RawMessage(`{"type":"object","properties":{"server":{"type":"string","description":"MCP server name to preview."},"source":{"type":"string","description":"Alias for server, useful for non-MCP harness terminology."},"tool":{"type":"string","description":"Tool name to preview."},"agent_id":{"type":"string","description":"Agent identity for no-auth local development only; ignored when inbound auth is enabled."},"request_id":{"type":"string","description":"Deprecated compatibility alias for agent_id in no-auth local development only; ignored when inbound auth is enabled."}},"additionalProperties":false}`), } // effectiveActionForTool returns the action of the first enabled rule that diff --git a/internal/api/handlers_test.go b/internal/api/handlers_test.go index 9dafddb..1b4f011 100644 --- a/internal/api/handlers_test.go +++ b/internal/api/handlers_test.go @@ -1062,6 +1062,38 @@ func TestMCPToolsList(t *testing.T) { } } +func TestMCPRulesToolSchemaIncludesRequestIDFallback(t *testing.T) { + h := NewHandler(&stubService{tools: []mcp.Tool{{Name: "demo_tool"}}}, stubServerService{}, nil, nil, nil, nil, nil, nil, nil, nil) + req := httptest.NewRequest(http.MethodPost, "/mcp/demo", strings.NewReader(`{"jsonrpc":"2.0","id":1,"method":"tools/list","params":{}}`)) + w := httptest.NewRecorder() + + h.Routes().ServeHTTP(w, req) + + var rpcResp struct { + Result struct { + Tools []struct { + Name string `json:"name"` + InputSchema struct { + Properties map[string]json.RawMessage `json:"properties"` + } `json:"inputSchema"` + } `json:"tools"` + } `json:"result"` + } + if err := json.Unmarshal(w.Body.Bytes(), &rpcResp); err != nil { + t.Fatal(err) + } + for _, tool := range rpcResp.Result.Tools { + if tool.Name != atryumRulesToolName { + continue + } + if _, ok := tool.InputSchema.Properties["request_id"]; !ok { + t.Fatalf("expected request_id in %s schema properties, got %#v", atryumRulesToolName, tool.InputSchema.Properties) + } + return + } + t.Fatalf("expected %s in tools/list, got %#v", atryumRulesToolName, rpcResp.Result.Tools) +} + func TestMCPToolsCallInterceptsInvocation(t *testing.T) { now := time.Now().UTC() svc := &stubService{invoke: invocation.InvocationResponse{InvocationID: "inv_123", ServerName: "demo", ToolName: "demo_tool", Status: invocation.StatusSucceeded, Input: json.RawMessage(`{"a":1}`), SubmittedAt: now, CompletedAt: &now, Result: json.RawMessage(`{"content":[{"type":"text","text":"ok"}]}`)}}