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/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, 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 6365b26..dc568cd 100644 --- a/internal/api/handlers.go +++ b/internal/api/handlers.go @@ -614,7 +614,14 @@ 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, 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"` @@ -622,17 +629,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"` } // ───────────────────────────────────────────────────────────────────────────── @@ -996,11 +1007,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, AgentRulesResponse{DefaultAction: invocation.RuleActionHumanApproval, Items: []AgentRule{}}) - return - } - agentID := auth.AgentIDFromContext(r.Context()) if agentID == "" && h.authValidator == nil { agentID = normalizeNoAuthAgentID(r.URL.Query().Get("agent_id")) @@ -1014,6 +1020,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()) @@ -1023,13 +1034,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 } @@ -1041,7 +1046,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{ @@ -1050,6 +1055,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 != "" && @@ -1064,6 +1070,70 @@ 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, + 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 "" @@ -1091,10 +1161,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 @@ -1103,7 +1180,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 { @@ -1187,7 +1264,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}) - annotated := h.annotateToolsWithPolicy(r.Context(), 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 { @@ -1198,6 +1275,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) @@ -1219,7 +1300,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 { - result = h.appendRulesContextToToolResult(r.Context(), result, server, params.Name) + result = h.appendRulesContextToToolResult(r.Context(), result, auth.AgentIDFromContext(r.Context()), server, params.Name) } h.writeRPCResult(w, req.ID, result) return @@ -1438,6 +1519,21 @@ func apiMatchPatterns(patterns []string, value string) bool { return false } +func agentRuleGuidance(action string) string { + switch action { + case invocation.RuleActionAutoApprove: + return "Auto-approved by this rule." + case invocation.RuleActionAutoDeny: + return "Auto-denied by this 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 @@ -1461,24 +1557,30 @@ 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 { - out := make([]any, len(tools)) +func (h *Handler) annotateToolsWithPolicy(ctx context.Context, agentID, server string, tools []mcp.Tool) []any { + 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) 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) } - agentID := auth.AgentIDFromContext(ctx) 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 != "" { @@ -1489,7 +1591,7 @@ func (h *Handler) annotateToolsWithPolicy(ctx context.Context, server string, to desc = prefix + desc } } - out[i] = annotatedTool{ + out = append(out, annotatedTool{ Name: t.Name, Description: desc, InputSchema: t.InputSchema, @@ -1497,9 +1599,15 @@ func (h *Handler) annotateToolsWithPolicy(ctx context.Context, server string, to 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 @@ -1517,11 +1625,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 } diff --git a/internal/api/handlers_test.go b/internal/api/handlers_test.go index 1dfd7ff..041576e 100644 --- a/internal/api/handlers_test.go +++ b/internal/api/handlers_test.go @@ -908,7 +908,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()) } } @@ -935,6 +938,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"}]}`)}} @@ -1026,12 +1071,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) { @@ -1070,6 +1165,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": "Auto-approved", + "auto-deny": "Auto-denied", + "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}, @@ -1122,7 +1281,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") } } @@ -1180,6 +1342,68 @@ 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 { + Name string `json:"name"` + 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) + } + 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 := bashTool.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{ @@ -1221,6 +1445,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?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) + + 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}}